Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Minor: refactor bloom filter tests to reduce duplication #8435

Merged
merged 3 commits into from
Jan 2, 2024
Merged

Minor: refactor bloom filter tests to reduce duplication #8435

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
08a274a
Minor: refactor bloom filter tests to reduce duplication
alamb Dec 6, 2023
72054ca
Merge remote-tracking branch 'apache/main' into alamb/less_test_copy_…
alamb Jan 1, 2024
287cdcf
Merge remote-tracking branch 'apache/main' into alamb/less_test_copy_…
alamb Jan 1, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
343 changes: 153 additions & 190 deletions datafusion/core/src/datasource/physical_plan/parquet/row_groups.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1013,82 +1013,28 @@ mod tests {
create_physical_expr(expr, &df_schema, schema, &execution_props).unwrap()
}

// Note the values in the `String` column are:
// ❯ select * from './parquet-testing/data/data_index_bloom_encoding_stats.parquet';
// +-----------+
// | String |
// +-----------+
// | Hello |
// | This is |
// | a |
// | test |
// | How |
// | are you |
// | doing |
// | today |
// | the quick |
// | brown fox |
// | jumps |
// | over |
// | the lazy |
// | dog |
// +-----------+
#[tokio::test]
async fn test_row_group_bloom_filter_pruning_predicate_simple_expr() {
// load parquet file
let testdata = datafusion_common::test_util::parquet_test_data();
let file_name = "data_index_bloom_encoding_stats.parquet";
let path = format!("{testdata}/{file_name}");
let data = bytes::Bytes::from(std::fs::read(path).unwrap());

// generate pruning predicate `(String = "Hello_Not_exists")`
let schema = Schema::new(vec![Field::new("String", DataType::Utf8, false)]);
let expr = col(r#""String""#).eq(lit("Hello_Not_Exists"));
let expr = logical2physical(&expr, &schema);
let pruning_predicate =
PruningPredicate::try_new(expr, Arc::new(schema)).unwrap();

let row_groups = vec![0];
let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate(
file_name,
data,
&pruning_predicate,
&row_groups,
)
.await
.unwrap();
assert!(pruned_row_groups.is_empty());
BloomFilterTest::new_data_index_bloom_encoding_stats()
.with_expect_all_pruned()
// generate pruning predicate `(String = "Hello_Not_exists")`
.run(col(r#""String""#).eq(lit("Hello_Not_Exists")))
.await
}

#[tokio::test]
async fn test_row_group_bloom_filter_pruning_predicate_mutiple_expr() {
// load parquet file
let testdata = datafusion_common::test_util::parquet_test_data();
let file_name = "data_index_bloom_encoding_stats.parquet";
let path = format!("{testdata}/{file_name}");
let data = bytes::Bytes::from(std::fs::read(path).unwrap());

// generate pruning predicate `(String = "Hello_Not_exists" OR String = "Hello_Not_exists2")`
let schema = Schema::new(vec![Field::new("String", DataType::Utf8, false)]);
let expr = lit("1").eq(lit("1")).and(
col(r#""String""#)
.eq(lit("Hello_Not_Exists"))
.or(col(r#""String""#).eq(lit("Hello_Not_Exists2"))),
);
let expr = logical2physical(&expr, &schema);
let pruning_predicate =
PruningPredicate::try_new(expr, Arc::new(schema)).unwrap();

let row_groups = vec![0];
let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate(
file_name,
data,
&pruning_predicate,
&row_groups,
)
.await
.unwrap();
assert!(pruned_row_groups.is_empty());
BloomFilterTest::new_data_index_bloom_encoding_stats()
.with_expect_all_pruned()
// generate pruning predicate `(String = "Hello_Not_exists" OR String = "Hello_Not_exists2")`
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A bit controversial due to "1" = "1" part in actual test case expression.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree -- I am not sure what the "1" = "1" is all about. @haohuaijin or @waynexia do you remember?

.run(
lit("1").eq(lit("1")).and(
col(r#""String""#)
.eq(lit("Hello_Not_Exists"))
.or(col(r#""String""#).eq(lit("Hello_Not_Exists2"))),
),
)
.await
}

#[tokio::test]
Expand Down Expand Up @@ -1129,144 +1075,161 @@ mod tests {

#[tokio::test]
async fn test_row_group_bloom_filter_pruning_predicate_with_exists_value() {
// load parquet file
let testdata = datafusion_common::test_util::parquet_test_data();
let file_name = "data_index_bloom_encoding_stats.parquet";
let path = format!("{testdata}/{file_name}");
let data = bytes::Bytes::from(std::fs::read(path).unwrap());

// generate pruning predicate `(String = "Hello")`
let schema = Schema::new(vec![Field::new("String", DataType::Utf8, false)]);
let expr = col(r#""String""#).eq(lit("Hello"));
let expr = logical2physical(&expr, &schema);
let pruning_predicate =
PruningPredicate::try_new(expr, Arc::new(schema)).unwrap();

let row_groups = vec![0];
let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate(
file_name,
data,
&pruning_predicate,
&row_groups,
)
.await
.unwrap();
assert_eq!(pruned_row_groups, row_groups);
BloomFilterTest::new_data_index_bloom_encoding_stats()
.with_expect_none_pruned()
// generate pruning predicate `(String = "Hello")`
.run(col(r#""String""#).eq(lit("Hello")))
.await
}

#[tokio::test]
async fn test_row_group_bloom_filter_pruning_predicate_with_exists_2_values() {
// load parquet file
let testdata = datafusion_common::test_util::parquet_test_data();
let file_name = "data_index_bloom_encoding_stats.parquet";
let path = format!("{testdata}/{file_name}");
let data = bytes::Bytes::from(std::fs::read(path).unwrap());

// generate pruning predicate `(String = "Hello") OR (String = "the quick")`
let schema = Schema::new(vec![Field::new("String", DataType::Utf8, false)]);
let expr = col(r#""String""#)
.eq(lit("Hello"))
.or(col(r#""String""#).eq(lit("the quick")));
let expr = logical2physical(&expr, &schema);
let pruning_predicate =
PruningPredicate::try_new(expr, Arc::new(schema)).unwrap();

let row_groups = vec![0];
let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate(
file_name,
data,
&pruning_predicate,
&row_groups,
)
.await
.unwrap();
assert_eq!(pruned_row_groups, row_groups);
BloomFilterTest::new_data_index_bloom_encoding_stats()
.with_expect_none_pruned()
// generate pruning predicate `(String = "Hello") OR (String = "the quick")`
.run(
col(r#""String""#)
.eq(lit("Hello"))
.or(col(r#""String""#).eq(lit("the quick"))),
)
.await
}

#[tokio::test]
async fn test_row_group_bloom_filter_pruning_predicate_with_exists_3_values() {
// load parquet file
let testdata = datafusion_common::test_util::parquet_test_data();
let file_name = "data_index_bloom_encoding_stats.parquet";
let path = format!("{testdata}/{file_name}");
let data = bytes::Bytes::from(std::fs::read(path).unwrap());

// generate pruning predicate `(String = "Hello") OR (String = "the quick") OR (String = "are you")`
let schema = Schema::new(vec![Field::new("String", DataType::Utf8, false)]);
let expr = col(r#""String""#)
.eq(lit("Hello"))
.or(col(r#""String""#).eq(lit("the quick")))
.or(col(r#""String""#).eq(lit("are you")));
let expr = logical2physical(&expr, &schema);
let pruning_predicate =
PruningPredicate::try_new(expr, Arc::new(schema)).unwrap();

let row_groups = vec![0];
let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate(
file_name,
data,
&pruning_predicate,
&row_groups,
)
.await
.unwrap();
assert_eq!(pruned_row_groups, row_groups);
BloomFilterTest::new_data_index_bloom_encoding_stats()
.with_expect_none_pruned()
// generate pruning predicate `(String = "Hello") OR (String = "the quick") OR (String = "are you")`
.run(
col(r#""String""#)
.eq(lit("Hello"))
.or(col(r#""String""#).eq(lit("the quick")))
.or(col(r#""String""#).eq(lit("are you"))),
)
.await
}

#[tokio::test]
async fn test_row_group_bloom_filter_pruning_predicate_with_or_not_eq() {
// load parquet file
let testdata = datafusion_common::test_util::parquet_test_data();
let file_name = "data_index_bloom_encoding_stats.parquet";
let path = format!("{testdata}/{file_name}");
let data = bytes::Bytes::from(std::fs::read(path).unwrap());

// generate pruning predicate `(String = "foo") OR (String != "bar")`
let schema = Schema::new(vec![Field::new("String", DataType::Utf8, false)]);
let expr = col(r#""String""#)
.not_eq(lit("foo"))
.or(col(r#""String""#).not_eq(lit("bar")));
let expr = logical2physical(&expr, &schema);
let pruning_predicate =
PruningPredicate::try_new(expr, Arc::new(schema)).unwrap();

let row_groups = vec![0];
let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate(
file_name,
data,
&pruning_predicate,
&row_groups,
)
.await
.unwrap();
assert_eq!(pruned_row_groups, row_groups);
BloomFilterTest::new_data_index_bloom_encoding_stats()
.with_expect_none_pruned()
// generate pruning predicate `(String = "foo") OR (String != "bar")`
.run(
col(r#""String""#)
.not_eq(lit("foo"))
.or(col(r#""String""#).not_eq(lit("bar"))),
)
.await
}

#[tokio::test]
async fn test_row_group_bloom_filter_pruning_predicate_without_bloom_filter() {
// load parquet file
let testdata = datafusion_common::test_util::parquet_test_data();
let file_name = "alltypes_plain.parquet";
let path = format!("{testdata}/{file_name}");
let data = bytes::Bytes::from(std::fs::read(path).unwrap());

// generate pruning predicate on a column without a bloom filter
let schema = Schema::new(vec![Field::new("string_col", DataType::Utf8, false)]);
let expr = col(r#""string_col""#).eq(lit("0"));
let expr = logical2physical(&expr, &schema);
let pruning_predicate =
PruningPredicate::try_new(expr, Arc::new(schema)).unwrap();
BloomFilterTest::new_all_types()
.with_expect_none_pruned()
.run(col(r#""string_col""#).eq(lit("0")))
.await
}

let row_groups = vec![0];
let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate(
file_name,
data,
&pruning_predicate,
&row_groups,
)
.await
.unwrap();
assert_eq!(pruned_row_groups, row_groups);
struct BloomFilterTest {
file_name: String,
schema: Schema,
// which row groups should be attempted to prune
row_groups: Vec<usize>,
// which row groups are expected to be left after pruning. Must be set
// otherwise will panic on run()
post_pruning_row_groups: Option<Vec<usize>>,
}

impl BloomFilterTest {
/// Return a test for data_index_bloom_encoding_stats.parquet
/// Note the values in the `String` column are:
/// ```sql
/// ❯ select * from './parquet-testing/data/data_index_bloom_encoding_stats.parquet';
/// +-----------+
/// | String |
/// +-----------+
/// | Hello |
/// | This is |
/// | a |
/// | test |
/// | How |
/// | are you |
/// | doing |
/// | today |
/// | the quick |
/// | brown fox |
/// | jumps |
/// | over |
/// | the lazy |
/// | dog |
/// +-----------+
/// ```
fn new_data_index_bloom_encoding_stats() -> Self {
Self {
file_name: String::from("data_index_bloom_encoding_stats.parquet"),
schema: Schema::new(vec![Field::new("String", DataType::Utf8, false)]),
row_groups: vec![0],
post_pruning_row_groups: None,
}
}

// Return a test for alltypes_plain.parquet
fn new_all_types() -> Self {
Self {
file_name: String::from("alltypes_plain.parquet"),
schema: Schema::new(vec![Field::new(
"string_col",
DataType::Utf8,
false,
)]),
row_groups: vec![0],
post_pruning_row_groups: None,
}
}

/// Expect all row groups to be pruned
pub fn with_expect_all_pruned(mut self) -> Self {
self.post_pruning_row_groups = Some(vec![]);
self
}

/// Expect all row groups not to be pruned
pub fn with_expect_none_pruned(mut self) -> Self {
self.post_pruning_row_groups = Some(self.row_groups.clone());
self
}

/// Prune this file using the specified expression and check that the expected row groups are left
async fn run(self, expr: Expr) {
let Self {
file_name,
schema,
row_groups,
post_pruning_row_groups,
} = self;

let post_pruning_row_groups =
post_pruning_row_groups.expect("post_pruning_row_groups must be set");

let testdata = datafusion_common::test_util::parquet_test_data();
let path = format!("{testdata}/{file_name}");
let data = bytes::Bytes::from(std::fs::read(path).unwrap());

let expr = logical2physical(&expr, &schema);
let pruning_predicate =
PruningPredicate::try_new(expr, Arc::new(schema)).unwrap();

let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate(
&file_name,
data,
&pruning_predicate,
&row_groups,
)
.await
.unwrap();
assert_eq!(pruned_row_groups, post_pruning_row_groups);
}
}

async fn test_row_group_bloom_filter_pruning_predicate(
Expand Down