Skip to content

Commit 01198eb

Browse files
alambalamb
committed
Minor: refactor bloom filter tests to reduce duplication
1 parent d33ca4d commit 01198eb

1 file changed

Lines changed: 132 additions & 170 deletions

File tree

  • datafusion/core/src/datasource/physical_plan/parquet

‎datafusion/core/src/datasource/physical_plan/parquet/row_groups.rs‎

Lines changed: 132 additions & 170 deletions
Original file line numberDiff line numberDiff line change
@@ -1097,60 +1097,26 @@ mod tests {
10971097
// +-----------+
10981098
#[tokio::test]
10991099
async fn test_row_group_bloom_filter_pruning_predicate_simple_expr() {
1100-
// load parquet file
1101-
let testdata = datafusion_common::test_util::parquet_test_data();
1102-
let file_name = "data_index_bloom_encoding_stats.parquet";
1103-
let path = format!("{testdata}/{file_name}");
1104-
let data = bytes::Bytes::from(std::fs::read(path).unwrap());
1105-
1106-
// generate pruning predicate `(String = "Hello_Not_exists")`
1107-
let schema = Schema::new(vec![Field::new("String", DataType::Utf8, false)]);
1108-
let expr = col(r#""String""#).eq(lit("Hello_Not_Exists"));
1109-
let expr = logical2physical(&expr, &schema);
1110-
let pruning_predicate =
1111-
PruningPredicate::try_new(expr, Arc::new(schema)).unwrap();
1112-
1113-
let row_groups = vec![0];
1114-
let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate(
1115-
file_name,
1116-
data,
1117-
&pruning_predicate,
1118-
&row_groups,
1119-
)
1120-
.await
1121-
.unwrap();
1122-
assert!(pruned_row_groups.is_empty());
1100+
BloomFilterTest::new_data_index_bloom_encoding_stats()
1101+
.with_expect_all_pruned()
1102+
// generate pruning predicate `(String = "Hello_Not_exists")`
1103+
.run(col(r#""String""#).eq(lit("Hello_Not_Exists")))
1104+
.await
11231105
}
11241106

11251107
#[tokio::test]
11261108
async fn test_row_group_bloom_filter_pruning_predicate_mutiple_expr() {
1127-
// load parquet file
1128-
let testdata = datafusion_common::test_util::parquet_test_data();
1129-
let file_name = "data_index_bloom_encoding_stats.parquet";
1130-
let path = format!("{testdata}/{file_name}");
1131-
let data = bytes::Bytes::from(std::fs::read(path).unwrap());
1132-
1133-
// generate pruning predicate `(String = "Hello_Not_exists" OR String = "Hello_Not_exists2")`
1134-
let schema = Schema::new(vec![Field::new("String", DataType::Utf8, false)]);
1135-
let expr = lit("1").eq(lit("1")).and(
1136-
col(r#""String""#)
1137-
.eq(lit("Hello_Not_Exists"))
1138-
.or(col(r#""String""#).eq(lit("Hello_Not_Exists2"))),
1139-
);
1140-
let expr = logical2physical(&expr, &schema);
1141-
let pruning_predicate =
1142-
PruningPredicate::try_new(expr, Arc::new(schema)).unwrap();
1143-
1144-
let row_groups = vec![0];
1145-
let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate(
1146-
file_name,
1147-
data,
1148-
&pruning_predicate,
1149-
&row_groups,
1150-
)
1151-
.await
1152-
.unwrap();
1153-
assert!(pruned_row_groups.is_empty());
1109+
BloomFilterTest::new_data_index_bloom_encoding_stats()
1110+
.with_expect_all_pruned()
1111+
// generate pruning predicate `(String = "Hello_Not_exists" OR String = "Hello_Not_exists2")`
1112+
.run(
1113+
lit("1").eq(lit("1")).and(
1114+
col(r#""String""#)
1115+
.eq(lit("Hello_Not_Exists"))
1116+
.or(col(r#""String""#).eq(lit("Hello_Not_Exists2"))),
1117+
),
1118+
)
1119+
.await
11541120
}
11551121

11561122
#[tokio::test]
@@ -1186,144 +1152,140 @@ mod tests {
11861152

11871153
#[tokio::test]
11881154
async fn test_row_group_bloom_filter_pruning_predicate_with_exists_value() {
1189-
// load parquet file
1190-
let testdata = datafusion_common::test_util::parquet_test_data();
1191-
let file_name = "data_index_bloom_encoding_stats.parquet";
1192-
let path = format!("{testdata}/{file_name}");
1193-
let data = bytes::Bytes::from(std::fs::read(path).unwrap());
1194-
1195-
// generate pruning predicate `(String = "Hello")`
1196-
let schema = Schema::new(vec![Field::new("String", DataType::Utf8, false)]);
1197-
let expr = col(r#""String""#).eq(lit("Hello"));
1198-
let expr = logical2physical(&expr, &schema);
1199-
let pruning_predicate =
1200-
PruningPredicate::try_new(expr, Arc::new(schema)).unwrap();
1201-
1202-
let row_groups = vec![0];
1203-
let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate(
1204-
file_name,
1205-
data,
1206-
&pruning_predicate,
1207-
&row_groups,
1208-
)
1209-
.await
1210-
.unwrap();
1211-
assert_eq!(pruned_row_groups, row_groups);
1155+
BloomFilterTest::new_data_index_bloom_encoding_stats()
1156+
.with_expect_none_pruned()
1157+
// generate pruning predicate `(String = "Hello")`
1158+
.run(col(r#""String""#).eq(lit("Hello")))
1159+
.await
12121160
}
12131161

12141162
#[tokio::test]
12151163
async fn test_row_group_bloom_filter_pruning_predicate_with_exists_2_values() {
1216-
// load parquet file
1217-
let testdata = datafusion_common::test_util::parquet_test_data();
1218-
let file_name = "data_index_bloom_encoding_stats.parquet";
1219-
let path = format!("{testdata}/{file_name}");
1220-
let data = bytes::Bytes::from(std::fs::read(path).unwrap());
1221-
1222-
// generate pruning predicate `(String = "Hello") OR (String = "the quick")`
1223-
let schema = Schema::new(vec![Field::new("String", DataType::Utf8, false)]);
1224-
let expr = col(r#""String""#)
1225-
.eq(lit("Hello"))
1226-
.or(col(r#""String""#).eq(lit("the quick")));
1227-
let expr = logical2physical(&expr, &schema);
1228-
let pruning_predicate =
1229-
PruningPredicate::try_new(expr, Arc::new(schema)).unwrap();
1230-
1231-
let row_groups = vec![0];
1232-
let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate(
1233-
file_name,
1234-
data,
1235-
&pruning_predicate,
1236-
&row_groups,
1237-
)
1238-
.await
1239-
.unwrap();
1240-
assert_eq!(pruned_row_groups, row_groups);
1164+
BloomFilterTest::new_data_index_bloom_encoding_stats()
1165+
.with_expect_none_pruned()
1166+
// generate pruning predicate `(String = "Hello") OR (String = "the quick")`
1167+
.run(
1168+
col(r#""String""#)
1169+
.eq(lit("Hello"))
1170+
.or(col(r#""String""#).eq(lit("the quick"))),
1171+
)
1172+
.await
12411173
}
12421174

12431175
#[tokio::test]
12441176
async fn test_row_group_bloom_filter_pruning_predicate_with_exists_3_values() {
1245-
// load parquet file
1246-
let testdata = datafusion_common::test_util::parquet_test_data();
1247-
let file_name = "data_index_bloom_encoding_stats.parquet";
1248-
let path = format!("{testdata}/{file_name}");
1249-
let data = bytes::Bytes::from(std::fs::read(path).unwrap());
1250-
1251-
// generate pruning predicate `(String = "Hello") OR (String = "the quick") OR (String = "are you")`
1252-
let schema = Schema::new(vec![Field::new("String", DataType::Utf8, false)]);
1253-
let expr = col(r#""String""#)
1254-
.eq(lit("Hello"))
1255-
.or(col(r#""String""#).eq(lit("the quick")))
1256-
.or(col(r#""String""#).eq(lit("are you")));
1257-
let expr = logical2physical(&expr, &schema);
1258-
let pruning_predicate =
1259-
PruningPredicate::try_new(expr, Arc::new(schema)).unwrap();
1260-
1261-
let row_groups = vec![0];
1262-
let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate(
1263-
file_name,
1264-
data,
1265-
&pruning_predicate,
1266-
&row_groups,
1267-
)
1268-
.await
1269-
.unwrap();
1270-
assert_eq!(pruned_row_groups, row_groups);
1177+
BloomFilterTest::new_data_index_bloom_encoding_stats()
1178+
.with_expect_none_pruned()
1179+
// generate pruning predicate `(String = "Hello") OR (String = "the quick") OR (String = "are you")`
1180+
.run(
1181+
col(r#""String""#)
1182+
.eq(lit("Hello"))
1183+
.or(col(r#""String""#).eq(lit("the quick")))
1184+
.or(col(r#""String""#).eq(lit("are you"))),
1185+
)
1186+
.await
12711187
}
12721188

12731189
#[tokio::test]
12741190
async fn test_row_group_bloom_filter_pruning_predicate_with_or_not_eq() {
1275-
// load parquet file
1276-
let testdata = datafusion_common::test_util::parquet_test_data();
1277-
let file_name = "data_index_bloom_encoding_stats.parquet";
1278-
let path = format!("{testdata}/{file_name}");
1279-
let data = bytes::Bytes::from(std::fs::read(path).unwrap());
1280-
1281-
// generate pruning predicate `(String = "foo") OR (String != "bar")`
1282-
let schema = Schema::new(vec![Field::new("String", DataType::Utf8, false)]);
1283-
let expr = col(r#""String""#)
1284-
.not_eq(lit("foo"))
1285-
.or(col(r#""String""#).not_eq(lit("bar")));
1286-
let expr = logical2physical(&expr, &schema);
1287-
let pruning_predicate =
1288-
PruningPredicate::try_new(expr, Arc::new(schema)).unwrap();
1289-
1290-
let row_groups = vec![0];
1291-
let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate(
1292-
file_name,
1293-
data,
1294-
&pruning_predicate,
1295-
&row_groups,
1296-
)
1297-
.await
1298-
.unwrap();
1299-
assert_eq!(pruned_row_groups, row_groups);
1191+
BloomFilterTest::new_data_index_bloom_encoding_stats()
1192+
.with_expect_none_pruned()
1193+
// generate pruning predicate `(String = "foo") OR (String != "bar")`
1194+
.run(
1195+
col(r#""String""#)
1196+
.not_eq(lit("foo"))
1197+
.or(col(r#""String""#).not_eq(lit("bar"))),
1198+
)
1199+
.await
13001200
}
13011201

13021202
#[tokio::test]
13031203
async fn test_row_group_bloom_filter_pruning_predicate_without_bloom_filter() {
1304-
// load parquet file
1305-
let testdata = datafusion_common::test_util::parquet_test_data();
1306-
let file_name = "alltypes_plain.parquet";
1307-
let path = format!("{testdata}/{file_name}");
1308-
let data = bytes::Bytes::from(std::fs::read(path).unwrap());
1309-
13101204
// generate pruning predicate on a column without a bloom filter
1311-
let schema = Schema::new(vec![Field::new("string_col", DataType::Utf8, false)]);
1312-
let expr = col(r#""string_col""#).eq(lit("0"));
1313-
let expr = logical2physical(&expr, &schema);
1314-
let pruning_predicate =
1315-
PruningPredicate::try_new(expr, Arc::new(schema)).unwrap();
1205+
BloomFilterTest::new_all_types()
1206+
.with_expect_none_pruned()
1207+
.run(col(r#""string_col""#).eq(lit("0")))
1208+
.await
1209+
}
13161210

1317-
let row_groups = vec![0];
1318-
let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate(
1319-
file_name,
1320-
data,
1321-
&pruning_predicate,
1322-
&row_groups,
1323-
)
1324-
.await
1325-
.unwrap();
1326-
assert_eq!(pruned_row_groups, row_groups);
1211+
struct BloomFilterTest {
1212+
file_name: String,
1213+
schema: Schema,
1214+
// which row groups should be attempted to prune
1215+
row_groups: Vec<usize>,
1216+
// which row groups are expected to be left after pruning. Must be set
1217+
// otherwise will panic on run()
1218+
post_pruning_row_groups: Option<Vec<usize>>,
1219+
}
1220+
1221+
impl BloomFilterTest {
1222+
/// Return a test for data_index_bloom_encoding_stats.parquet
1223+
fn new_data_index_bloom_encoding_stats() -> Self {
1224+
Self {
1225+
file_name: String::from("data_index_bloom_encoding_stats.parquet"),
1226+
schema: Schema::new(vec![Field::new("String", DataType::Utf8, false)]),
1227+
row_groups: vec![0],
1228+
post_pruning_row_groups: None,
1229+
}
1230+
}
1231+
1232+
// Return a test for alltypes_plain.parquet
1233+
fn new_all_types() -> Self {
1234+
Self {
1235+
file_name: String::from("alltypes_plain.parquet"),
1236+
schema: Schema::new(vec![Field::new(
1237+
"string_col",
1238+
DataType::Utf8,
1239+
false,
1240+
)]),
1241+
row_groups: vec![0],
1242+
post_pruning_row_groups: None,
1243+
}
1244+
}
1245+
1246+
/// Expect all row groups to be pruned
1247+
pub fn with_expect_all_pruned(mut self) -> Self {
1248+
self.post_pruning_row_groups = Some(vec![]);
1249+
self
1250+
}
1251+
1252+
/// Expect all row groups not to be pruned
1253+
pub fn with_expect_none_pruned(mut self) -> Self {
1254+
self.post_pruning_row_groups = Some(self.row_groups.clone());
1255+
self
1256+
}
1257+
1258+
/// Prune this file using the specified expression and check that the expected row groups are left
1259+
async fn run(self, expr: Expr) {
1260+
let Self {
1261+
file_name,
1262+
schema,
1263+
row_groups,
1264+
post_pruning_row_groups,
1265+
} = self;
1266+
1267+
let post_pruning_row_groups =
1268+
post_pruning_row_groups.expect("post_pruning_row_groups must be set");
1269+
1270+
let testdata = datafusion_common::test_util::parquet_test_data();
1271+
let path = format!("{testdata}/{file_name}");
1272+
let data = bytes::Bytes::from(std::fs::read(path).unwrap());
1273+
1274+
// generate pruning predicate on a column without a bloom filter
1275+
let expr = logical2physical(&expr, &schema);
1276+
let pruning_predicate =
1277+
PruningPredicate::try_new(expr, Arc::new(schema)).unwrap();
1278+
1279+
let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate(
1280+
&file_name,
1281+
data,
1282+
&pruning_predicate,
1283+
&row_groups,
1284+
)
1285+
.await
1286+
.unwrap();
1287+
assert_eq!(pruned_row_groups, post_pruning_row_groups);
1288+
}
13271289
}
13281290

13291291
async fn test_row_group_bloom_filter_pruning_predicate(

0 commit comments

Comments
 (0)