This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-22348-c8b784a01f5d0bcbe0dac806730fb61afc0be8ef in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 7802abc0d9561fcd29313565888127630489ebae Author: Adrian Garcia Badaracco <[email protected]> AuthorDate: Tue May 19 05:03:34 2026 -0700 refactor(parquet-datasource): split bloom_filter out of row_group_filter.rs (#22348) ## Which issue does this PR close? Relates to the discussion in #22024 about the Parquet datasource crate becoming hard to navigate. Split out of #22156, which bundled several code-motion moves into one PR — this is one of three smaller, independently-reviewable PRs that replace it. ## Rationale for this change `row_group_filter.rs` had grown to ~1,900 LOC. It mixes "data we loaded from the file" with "the access-plan filter that consumes it." This PR is **code motion only**: no behavior change and no public API change. ## What changes are included in this PR? Extracts `BloomFilterStatistics` — the loaded Split Block Bloom Filter (SBBF) data plus its `PruningStatistics` adapter — from `row_group_filter.rs` into a new `bloom_filter.rs`, and moves the bloom-filter tests alongside it. This separates `BloomFilterStatistics` (data loaded from the file) from `RowGroupAccessPlanFilter` (the access-plan filter that consumes it), leaving `row_group_filter.rs` focused on the latter. Each commit builds green on its own: 1. **Split `bloom_filter` out of `row_group_filter.rs`** — move `BloomFilterStatistics` and its `PruningStatistics` adapter into a new `bloom_filter.rs`. 2. **Build `BloomFilterStatistics` via its constructors in tests** — the test built it with a struct literal; use the existing `with_capacity`/`insert` constructors (the pattern `opener.rs` already uses) so the `column_sbbf` field stays private. 3. **Extract `ExpectedPruning` into a shared `test_util` module** — the one test helper shared between the row-group and bloom-filter tests, so the bloom-filter tests can move out. Adds a `#[cfg(test)] pub(crate) fn access_plan()` accessor on `RowGroupAccessPlanFilter` so the helper can assert from a sibling module without widening field visibility. 4. **Move the bloom-filter tests into `bloom_filter.rs`** — relocate `test_row_group_bloom_*`, the `BloomFilterTest` builder, and its helper next to the code they exercise. `BloomFilterStatistics` is crate-internal; `row_group_filter` re-exports it (`pub(crate) use`) so the existing `crate::row_group_filter::BloomFilterStatistics` path keeps resolving for in-crate callers. Aside from `row_group_filter.rs` and `bloom_filter.rs`, this adds a new `src/test_util.rs` and a one-line module declaration in `mod.rs`. ## Are these changes tested? Yes, covered by existing tests. `cargo test -p datafusion-datasource-parquet --all-features` (122 passing) and `cargo clippy -p datafusion-datasource-parquet --all-targets --all-features -- -D warnings` both pass. ## Are there any user-facing changes? No. `BloomFilterStatistics` is crate-internal; this only reorganizes files inside the crate. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]> --- datafusion/datasource-parquet/src/bloom_filter.rs | 560 ++++++++++++++++++++ datafusion/datasource-parquet/src/mod.rs | 3 + .../datasource-parquet/src/row_group_filter.rs | 577 +-------------------- datafusion/datasource-parquet/src/test_util.rs | 71 +++ 4 files changed, 650 insertions(+), 561 deletions(-) diff --git a/datafusion/datasource-parquet/src/bloom_filter.rs b/datafusion/datasource-parquet/src/bloom_filter.rs new file mode 100644 index 0000000000..9388aba438 --- /dev/null +++ b/datafusion/datasource-parquet/src/bloom_filter.rs @@ -0,0 +1,560 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Loaded Parquet Split Block Bloom Filter (SBBF) data, with a +//! [`PruningStatistics`] adapter so the predicate-pruning machinery in +//! [`datafusion_pruning`] can consume it. + +use std::collections::{HashMap, HashSet}; + +use arrow::array::{ArrayRef, BooleanArray}; +use datafusion_common::pruning::PruningStatistics; +use datafusion_common::{Column, ScalarValue}; +use parquet::basic::Type; +use parquet::bloom_filter::Sbbf; +use parquet::data_type::Decimal; + +/// In memory Parquet Split Block Bloom Filters (SBBF). +/// +/// This structure implements [`PruningStatistics`] and is used to prune +/// Parquet row groups and data pages based on the query predicate. +#[derive(Debug, Clone, Default)] +pub(crate) struct BloomFilterStatistics { + /// Per-column Bloom filters + /// Key: predicate column name + /// Value: + /// * [`Sbbf`] (Bloom filter), + /// * Parquet physical [`Type`] needed to evaluate literals against the filter + column_sbbf: HashMap<String, (Sbbf, Type)>, +} + +impl BloomFilterStatistics { + /// Create an empty [`BloomFilterStatistics`] + pub(crate) fn new() -> Self { + Default::default() + } + + /// Create an empty [`BloomFilterStatistics`] with the specified capacity + pub(crate) fn with_capacity(capacity: usize) -> Self { + Self { + column_sbbf: HashMap::with_capacity(capacity), + } + } + + /// Add a Bloom filter and type for the specified column + pub(crate) fn insert(&mut self, column: impl Into<String>, sbbf: Sbbf, ty: Type) { + self.column_sbbf.insert(column.into(), (sbbf, ty)); + } + + /// Helper function for checking if [`Sbbf`] filter contains [`ScalarValue`]. + /// + /// In case the type of scalar is not supported, returns `true`, assuming that the + /// value may be present. + fn check_scalar(sbbf: &Sbbf, value: &ScalarValue, parquet_type: &Type) -> bool { + match value { + ScalarValue::Utf8(Some(v)) + | ScalarValue::Utf8View(Some(v)) + | ScalarValue::LargeUtf8(Some(v)) => sbbf.check(&v.as_str()), + ScalarValue::Binary(Some(v)) + | ScalarValue::BinaryView(Some(v)) + | ScalarValue::LargeBinary(Some(v)) => sbbf.check(v), + ScalarValue::FixedSizeBinary(_size, Some(v)) => sbbf.check(v), + ScalarValue::Boolean(Some(v)) => sbbf.check(v), + ScalarValue::Float64(Some(v)) => sbbf.check(v), + ScalarValue::Float32(Some(v)) => sbbf.check(v), + ScalarValue::Int64(Some(v)) => sbbf.check(v), + ScalarValue::Int32(Some(v)) => sbbf.check(v), + ScalarValue::UInt64(Some(v)) => sbbf.check(v), + ScalarValue::UInt32(Some(v)) => sbbf.check(v), + ScalarValue::Decimal128(Some(v), p, s) => match parquet_type { + Type::INT32 => { + //https://github.com/apache/parquet-format/blob/eb4b31c1d64a01088d02a2f9aefc6c17c54cc6fc/Encodings.md?plain=1#L35-L42 + // All physical type are little-endian + if *p > 9 { + //DECIMAL can be used to annotate the following types: + // + // int32: for 1 <= precision <= 9 + // int64: for 1 <= precision <= 18 + return true; + } + let b = (*v as i32).to_le_bytes(); + // Use Decimal constructor after https://github.com/apache/arrow-rs/issues/5325 + let decimal = Decimal::Int32 { + value: b, + precision: *p as i32, + scale: *s as i32, + }; + sbbf.check(&decimal) + } + Type::INT64 => { + if *p > 18 { + return true; + } + let b = (*v as i64).to_le_bytes(); + let decimal = Decimal::Int64 { + value: b, + precision: *p as i32, + scale: *s as i32, + }; + sbbf.check(&decimal) + } + Type::FIXED_LEN_BYTE_ARRAY => { + // keep with from_bytes_to_i128 + let b = v.to_be_bytes().to_vec(); + // Use Decimal constructor after https://github.com/apache/arrow-rs/issues/5325 + let decimal = Decimal::Bytes { + value: b.into(), + precision: *p as i32, + scale: *s as i32, + }; + sbbf.check(&decimal) + } + _ => true, + }, + ScalarValue::Dictionary(_, inner) => { + BloomFilterStatistics::check_scalar(sbbf, inner, parquet_type) + } + _ => true, + } + } +} + +impl PruningStatistics for BloomFilterStatistics { + fn min_values(&self, _column: &Column) -> Option<ArrayRef> { + None + } + + fn max_values(&self, _column: &Column) -> Option<ArrayRef> { + None + } + + fn num_containers(&self) -> usize { + 1 + } + + fn null_counts(&self, _column: &Column) -> Option<ArrayRef> { + None + } + + fn row_counts(&self) -> Option<ArrayRef> { + None + } + + /// Use bloom filters to determine if we are sure this column can not + /// possibly contain `values` + /// + /// The `contained` API returns false if the bloom filters knows that *ALL* + /// of the values in a column are not present. + fn contained( + &self, + column: &Column, + values: &HashSet<ScalarValue>, + ) -> Option<BooleanArray> { + let (sbbf, parquet_type) = self.column_sbbf.get(column.name.as_str())?; + + // Bloom filters are probabilistic data structures that can return false + // positives (i.e. it might return true even if the value is not + // present) however, the bloom filter will return `false` if the value is + // definitely not present. + + let known_not_present = values + .iter() + .map(|value| BloomFilterStatistics::check_scalar(sbbf, value, parquet_type)) + // The row group doesn't contain any of the values if + // all the checks are false + .all(|v| !v); + + let contains = if known_not_present { + Some(false) + } else { + // Given the bloom filter is probabilistic, we can't be sure that + // the row group actually contains the values. Return `None` to + // indicate this uncertainty + None + }; + + Some(BooleanArray::from(vec![contains])) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::sync::Arc; + + use crate::reader::ParquetFileReader; + use crate::test_util::ExpectedPruning; + use crate::{ParquetAccessPlan, ParquetFileMetrics, RowGroupAccessPlanFilter}; + + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::Result; + use datafusion_expr::{Expr, col, lit}; + use datafusion_physical_expr::planner::logical2physical; + use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; + use datafusion_pruning::PruningPredicate; + use object_store::ObjectStoreExt; + use parquet::arrow::ParquetRecordBatchStreamBuilder; + use parquet::arrow::async_reader::ParquetObjectReader; + + #[tokio::test] + async fn test_row_group_bloom_filter_pruning_predicate_simple_expr() { + 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_multiple_expr() { + BloomFilterTest::new_data_index_bloom_encoding_stats() + .with_expect_all_pruned() + // generate pruning predicate `(String = "Hello_Not_exists" OR String = "Hello_Not_exists2")` + .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] + async fn test_row_group_bloom_filter_pruning_predicate_multiple_expr_view() { + BloomFilterTest::new_data_index_bloom_encoding_stats() + .with_expect_all_pruned() + // generate pruning predicate `(String = "Hello_Not_exists" OR String = "Hello_Not_exists2")` + .run( + lit("1").eq(lit("1")).and( + col(r#""String""#) + .eq(Expr::Literal( + ScalarValue::Utf8View(Some(String::from("Hello_Not_Exists"))), + None, + )) + .or(col(r#""String""#).eq(Expr::Literal( + ScalarValue::Utf8View(Some(String::from( + "Hello_Not_Exists2", + ))), + None, + ))), + ), + ) + .await + } + + #[tokio::test] + async fn test_row_group_bloom_filter_pruning_predicate_sql_in() { + // 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 + let schema = Schema::new(vec![Field::new("String", DataType::Utf8, false)]); + + let expr = col(r#""String""#).in_list( + (1..25) + .map(|i| lit(format!("Hello_Not_Exists{i}"))) + .collect::<Vec<_>>(), + false, + ); + 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, + ) + .await + .unwrap(); + assert!( + pruned_row_groups + .access_plan() + .row_group_indexes() + .is_empty() + ); + } + + #[tokio::test] + async fn test_row_group_bloom_filter_pruning_predicate_with_exists_value() { + 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() { + 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() { + 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_exists_3_values_view() { + 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(Expr::Literal( + ScalarValue::Utf8View(Some(String::from("Hello"))), + None, + )) + .or(col(r#""String""#).eq(Expr::Literal( + ScalarValue::Utf8View(Some(String::from("the quick"))), + None, + ))) + .or(col(r#""String""#).eq(Expr::Literal( + ScalarValue::Utf8View(Some(String::from("are you"))), + None, + ))), + ) + .await + } + + #[tokio::test] + async fn test_row_group_bloom_filter_pruning_predicate_with_or_not_eq() { + 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() { + // generate pruning predicate on a column without a bloom filter + BloomFilterTest::new_all_types() + .with_expect_none_pruned() + .run(col(r#""string_col""#).eq(lit("0"))) + .await + } + + struct BloomFilterTest { + file_name: String, + schema: Schema, + // which row groups are expected to be left after pruning + post_pruning_row_groups: ExpectedPruning, + } + + 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)]), + post_pruning_row_groups: ExpectedPruning::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, + )]), + post_pruning_row_groups: ExpectedPruning::None, + } + } + + /// Expect all row groups to be pruned + pub fn with_expect_all_pruned(mut self) -> Self { + self.post_pruning_row_groups = ExpectedPruning::All; + self + } + + /// Expect all row groups not to be pruned + pub fn with_expect_none_pruned(mut self) -> Self { + self.post_pruning_row_groups = ExpectedPruning::None; + 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, + post_pruning_row_groups, + } = self; + + 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, + ) + .await + .unwrap(); + + post_pruning_row_groups.assert(&pruned_row_groups); + } + } + + /// Evaluates the pruning predicate on the specified row groups and returns the row groups that are left + async fn test_row_group_bloom_filter_pruning_predicate( + file_name: &str, + data: bytes::Bytes, + pruning_predicate: &PruningPredicate, + ) -> Result<RowGroupAccessPlanFilter> { + use datafusion_datasource::PartitionedFile; + use object_store::ObjectMeta; + + let object_meta = ObjectMeta { + location: object_store::path::Path::parse(file_name).expect("creating path"), + last_modified: chrono::DateTime::from(std::time::SystemTime::now()), + size: data.len() as u64, + e_tag: None, + version: None, + }; + let in_memory = object_store::memory::InMemory::new(); + in_memory + .put(&object_meta.location, data.into()) + .await + .expect("put parquet file into in memory object store"); + + let metrics = ExecutionPlanMetricsSet::new(); + let file_metrics = + ParquetFileMetrics::new(0, object_meta.location.as_ref(), &metrics); + let inner = + ParquetObjectReader::new(Arc::new(in_memory), object_meta.location.clone()) + .with_file_size(object_meta.size); + + let partitioned_file = PartitionedFile::new_from_meta(object_meta); + + let reader = ParquetFileReader { + inner, + file_metrics: file_metrics.clone(), + partitioned_file, + }; + let mut builder = ParquetRecordBatchStreamBuilder::new(reader).await.unwrap(); + + let access_plan = ParquetAccessPlan::new_all(builder.metadata().num_row_groups()); + let mut pruned_row_groups = RowGroupAccessPlanFilter::new(access_plan); + let literal_columns = pruning_predicate.literal_columns(); + let parquet_columns: Vec<_> = literal_columns + .into_iter() + .filter_map(|column_name| { + let (column_idx, _) = parquet::arrow::parquet_column( + builder.parquet_schema(), + pruning_predicate.schema(), + &column_name, + )?; + Some(( + column_name.to_string(), + column_idx, + builder.parquet_schema().column(column_idx).physical_type(), + )) + }) + .collect::<Vec<_>>(); + let mut row_group_bloom_filters = + Vec::with_capacity(builder.metadata().num_row_groups()); + row_group_bloom_filters.resize_with( + builder.metadata().num_row_groups(), + BloomFilterStatistics::new, + ); + for idx in pruned_row_groups.row_group_indexes() { + let mut bloom_filters = + BloomFilterStatistics::with_capacity(parquet_columns.len()); + for (column_name, column_idx, physical_type) in &parquet_columns { + let bf = match builder + .get_row_group_column_bloom_filter(idx, *column_idx) + .await + { + Ok(Some(bf)) => bf, + Ok(None) => continue, + Err(e) => { + log::debug!("Ignoring error reading bloom filter: {e}"); + file_metrics.predicate_evaluation_errors.add(1); + continue; + } + }; + bloom_filters.insert(column_name.clone(), bf, *physical_type); + } + row_group_bloom_filters[idx] = bloom_filters; + } + pruned_row_groups.prune_by_bloom_filters( + pruning_predicate, + &file_metrics, + &row_group_bloom_filters, + ); + + Ok(pruned_row_groups) + } +} diff --git a/datafusion/datasource-parquet/src/mod.rs b/datafusion/datasource-parquet/src/mod.rs index af2e103e6a..b3a1328607 100644 --- a/datafusion/datasource-parquet/src/mod.rs +++ b/datafusion/datasource-parquet/src/mod.rs @@ -25,6 +25,7 @@ #![cfg_attr(test, allow(clippy::needless_pass_by_value))] pub mod access_plan; +mod bloom_filter; pub mod file_format; pub mod metadata; mod metrics; @@ -39,6 +40,8 @@ mod sink; mod sort; pub mod source; mod supported_predicates; +#[cfg(test)] +mod test_util; mod writer; pub use access_plan::{ParquetAccessPlan, RowGroupAccess}; diff --git a/datafusion/datasource-parquet/src/row_group_filter.rs b/datafusion/datasource-parquet/src/row_group_filter.rs index c45e69600f..07f4fe92cf 100644 --- a/datafusion/datasource-parquet/src/row_group_filter.rs +++ b/datafusion/datasource-parquet/src/row_group_filter.rs @@ -15,10 +15,13 @@ // specific language governing permissions and limitations // under the License. -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use std::sync::Arc; use super::{ParquetAccessPlan, ParquetFileMetrics}; +// Re-exported so the existing `crate::row_group_filter::BloomFilterStatistics` +// path keeps resolving for in-crate callers (e.g. `opener`). +pub(crate) use crate::bloom_filter::BloomFilterStatistics; use arrow::array::{ArrayRef, BooleanArray, UInt64Array}; use arrow::datatypes::Schema; use datafusion_common::pruning::PruningStatistics; @@ -30,10 +33,8 @@ use datafusion_physical_expr::utils::collect_columns; use datafusion_physical_expr::{PhysicalExpr, PhysicalExprSimplifier}; use datafusion_pruning::PruningPredicate; use parquet::arrow::arrow_reader::statistics::StatisticsConverter; -use parquet::basic::Type; -use parquet::data_type::Decimal; +use parquet::file::metadata::RowGroupMetaData; use parquet::schema::types::SchemaDescriptor; -use parquet::{bloom_filter::Sbbf, file::metadata::RowGroupMetaData}; /// Reduces the [`ParquetAccessPlan`] based on row group level metadata. /// @@ -74,6 +75,15 @@ impl RowGroupAccessPlanFilter { self.access_plan } + /// Returns a reference to the inner access plan. + /// + /// Test-only accessor used by the shared assertion helpers in + /// [`crate::test_util`]. + #[cfg(test)] + pub(crate) fn access_plan(&self) -> &ParquetAccessPlan { + &self.access_plan + } + /// Returns the is_fully_matched vector. pub fn is_fully_matched(&self) -> &Vec<bool> { self.access_plan.fully_matched() @@ -443,169 +453,6 @@ impl RowGroupAccessPlanFilter { } } -/// In memory Parquet Split Block Bloom Filters (SBBF). -/// -/// This structure implements [`PruningStatistics`] and is used to prune -/// Parquet row groups and data pages based on the query predicate. -#[derive(Debug, Clone, Default)] -pub(crate) struct BloomFilterStatistics { - /// Per-column Bloom filters - /// Key: predicate column name - /// Value: - /// * [`Sbbf`] (Bloom filter), - /// * Parquet physical [`Type`] needed to evaluate literals against the filter - column_sbbf: HashMap<String, (Sbbf, Type)>, -} - -impl BloomFilterStatistics { - /// Create an empty [`BloomFilterStatistics`] - pub(crate) fn new() -> Self { - Default::default() - } - - /// Create an empty [`BloomFilterStatistics`] with the specified capacity - pub(crate) fn with_capacity(capacity: usize) -> Self { - Self { - column_sbbf: HashMap::with_capacity(capacity), - } - } - - /// Add a Bloom filter and type for the specified column - pub(crate) fn insert(&mut self, column: impl Into<String>, sbbf: Sbbf, ty: Type) { - self.column_sbbf.insert(column.into(), (sbbf, ty)); - } - - /// Helper function for checking if [`Sbbf`] filter contains [`ScalarValue`]. - /// - /// In case the type of scalar is not supported, returns `true`, assuming that the - /// value may be present. - fn check_scalar(sbbf: &Sbbf, value: &ScalarValue, parquet_type: &Type) -> bool { - match value { - ScalarValue::Utf8(Some(v)) - | ScalarValue::Utf8View(Some(v)) - | ScalarValue::LargeUtf8(Some(v)) => sbbf.check(&v.as_str()), - ScalarValue::Binary(Some(v)) - | ScalarValue::BinaryView(Some(v)) - | ScalarValue::LargeBinary(Some(v)) => sbbf.check(v), - ScalarValue::FixedSizeBinary(_size, Some(v)) => sbbf.check(v), - ScalarValue::Boolean(Some(v)) => sbbf.check(v), - ScalarValue::Float64(Some(v)) => sbbf.check(v), - ScalarValue::Float32(Some(v)) => sbbf.check(v), - ScalarValue::Int64(Some(v)) => sbbf.check(v), - ScalarValue::Int32(Some(v)) => sbbf.check(v), - ScalarValue::UInt64(Some(v)) => sbbf.check(v), - ScalarValue::UInt32(Some(v)) => sbbf.check(v), - ScalarValue::Decimal128(Some(v), p, s) => match parquet_type { - Type::INT32 => { - //https://github.com/apache/parquet-format/blob/eb4b31c1d64a01088d02a2f9aefc6c17c54cc6fc/Encodings.md?plain=1#L35-L42 - // All physical type are little-endian - if *p > 9 { - //DECIMAL can be used to annotate the following types: - // - // int32: for 1 <= precision <= 9 - // int64: for 1 <= precision <= 18 - return true; - } - let b = (*v as i32).to_le_bytes(); - // Use Decimal constructor after https://github.com/apache/arrow-rs/issues/5325 - let decimal = Decimal::Int32 { - value: b, - precision: *p as i32, - scale: *s as i32, - }; - sbbf.check(&decimal) - } - Type::INT64 => { - if *p > 18 { - return true; - } - let b = (*v as i64).to_le_bytes(); - let decimal = Decimal::Int64 { - value: b, - precision: *p as i32, - scale: *s as i32, - }; - sbbf.check(&decimal) - } - Type::FIXED_LEN_BYTE_ARRAY => { - // keep with from_bytes_to_i128 - let b = v.to_be_bytes().to_vec(); - // Use Decimal constructor after https://github.com/apache/arrow-rs/issues/5325 - let decimal = Decimal::Bytes { - value: b.into(), - precision: *p as i32, - scale: *s as i32, - }; - sbbf.check(&decimal) - } - _ => true, - }, - ScalarValue::Dictionary(_, inner) => { - BloomFilterStatistics::check_scalar(sbbf, inner, parquet_type) - } - _ => true, - } - } -} - -impl PruningStatistics for BloomFilterStatistics { - fn min_values(&self, _column: &Column) -> Option<ArrayRef> { - None - } - - fn max_values(&self, _column: &Column) -> Option<ArrayRef> { - None - } - - fn num_containers(&self) -> usize { - 1 - } - - fn null_counts(&self, _column: &Column) -> Option<ArrayRef> { - None - } - - fn row_counts(&self) -> Option<ArrayRef> { - None - } - - /// Use bloom filters to determine if we are sure this column can not - /// possibly contain `values` - /// - /// The `contained` API returns false if the bloom filters knows that *ALL* - /// of the values in a column are not present. - fn contained( - &self, - column: &Column, - values: &HashSet<ScalarValue>, - ) -> Option<BooleanArray> { - let (sbbf, parquet_type) = self.column_sbbf.get(column.name.as_str())?; - - // Bloom filters are probabilistic data structures that can return false - // positives (i.e. it might return true even if the value is not - // present) however, the bloom filter will return `false` if the value is - // definitely not present. - - let known_not_present = values - .iter() - .map(|value| BloomFilterStatistics::check_scalar(sbbf, value, parquet_type)) - // The row group doesn't contain any of the values if - // all the checks are false - .all(|v| !v); - - let contains = if known_not_present { - Some(false) - } else { - // Given the bloom filter is probabilistic, we can't be sure that - // the row group actually contains the values. Return `None` to - // indicate this uncertainty - None - }; - - Some(BooleanArray::from(vec![contains])) - } -} - /// Wraps a slice of [`RowGroupMetaData`] in a way that implements [`PruningStatistics`] struct RowGroupPruningStatistics<'a> { parquet_schema: &'a SchemaDescriptor, @@ -680,17 +527,14 @@ mod tests { use std::ops::Rem; use super::*; - use crate::reader::ParquetFileReader; + use crate::test_util::ExpectedPruning; use arrow::datatypes::DataType::Decimal128; use arrow::datatypes::{DataType, Field}; - use datafusion_expr::{Expr, cast, col, lit}; + use datafusion_expr::{cast, col, lit}; use datafusion_physical_expr::planner::logical2physical; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; - use object_store::ObjectStoreExt; use parquet::arrow::ArrowSchemaConverter; - use parquet::arrow::ParquetRecordBatchStreamBuilder; - use parquet::arrow::async_reader::ParquetObjectReader; use parquet::basic::LogicalType; use parquet::data_type::{ByteArray, FixedLenByteArray}; use parquet::file::metadata::ColumnChunkMetaData; @@ -1528,396 +1372,7 @@ mod tests { ParquetFileMetrics::new(0, "file.parquet", &metrics) } - #[tokio::test] - async fn test_row_group_bloom_filter_pruning_predicate_simple_expr() { - 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_multiple_expr() { - BloomFilterTest::new_data_index_bloom_encoding_stats() - .with_expect_all_pruned() - // generate pruning predicate `(String = "Hello_Not_exists" OR String = "Hello_Not_exists2")` - .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] - async fn test_row_group_bloom_filter_pruning_predicate_multiple_expr_view() { - BloomFilterTest::new_data_index_bloom_encoding_stats() - .with_expect_all_pruned() - // generate pruning predicate `(String = "Hello_Not_exists" OR String = "Hello_Not_exists2")` - .run( - lit("1").eq(lit("1")).and( - col(r#""String""#) - .eq(Expr::Literal( - ScalarValue::Utf8View(Some(String::from("Hello_Not_Exists"))), - None, - )) - .or(col(r#""String""#).eq(Expr::Literal( - ScalarValue::Utf8View(Some(String::from( - "Hello_Not_Exists2", - ))), - None, - ))), - ), - ) - .await - } - - #[tokio::test] - async fn test_row_group_bloom_filter_pruning_predicate_sql_in() { - // 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 - let schema = Schema::new(vec![Field::new("String", DataType::Utf8, false)]); - - let expr = col(r#""String""#).in_list( - (1..25) - .map(|i| lit(format!("Hello_Not_Exists{i}"))) - .collect::<Vec<_>>(), - false, - ); - 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, - ) - .await - .unwrap(); - assert!(pruned_row_groups.access_plan.row_group_indexes().is_empty()); - } - - #[tokio::test] - async fn test_row_group_bloom_filter_pruning_predicate_with_exists_value() { - 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() { - 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() { - 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_exists_3_values_view() { - 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(Expr::Literal( - ScalarValue::Utf8View(Some(String::from("Hello"))), - None, - )) - .or(col(r#""String""#).eq(Expr::Literal( - ScalarValue::Utf8View(Some(String::from("the quick"))), - None, - ))) - .or(col(r#""String""#).eq(Expr::Literal( - ScalarValue::Utf8View(Some(String::from("are you"))), - None, - ))), - ) - .await - } - - #[tokio::test] - async fn test_row_group_bloom_filter_pruning_predicate_with_or_not_eq() { - 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() { - // generate pruning predicate on a column without a bloom filter - BloomFilterTest::new_all_types() - .with_expect_none_pruned() - .run(col(r#""string_col""#).eq(lit("0"))) - .await - } - - // What row groups are expected to be left after pruning - #[derive(Debug)] - enum ExpectedPruning { - All, - /// Only the specified row groups are expected to REMAIN (not what is pruned) - Some(Vec<usize>), - None, - } - - impl ExpectedPruning { - /// asserts that the pruned row group match this expectation - fn assert(&self, row_groups: &RowGroupAccessPlanFilter) { - let num_row_groups = row_groups.access_plan.len(); - assert!(num_row_groups > 0); - let num_pruned = (0..num_row_groups) - .filter_map(|i| { - if row_groups.access_plan.should_scan(i) { - None - } else { - Some(1) - } - }) - .sum::<usize>(); - - match self { - Self::All => { - assert_eq!( - num_row_groups, num_pruned, - "Expected all row groups to be pruned, but got {row_groups:?}" - ); - } - ExpectedPruning::None => { - assert_eq!( - num_pruned, 0, - "Expected no row groups to be pruned, but got {row_groups:?}" - ); - } - ExpectedPruning::Some(expected) => { - let actual = row_groups.access_plan.row_group_indexes(); - assert_eq!( - expected, &actual, - "Unexpected row groups pruned. Expected {expected:?}, got {actual:?}" - ); - } - } - } - } - fn assert_pruned(row_groups: RowGroupAccessPlanFilter, expected: ExpectedPruning) { expected.assert(&row_groups); } - - struct BloomFilterTest { - file_name: String, - schema: Schema, - // which row groups are expected to be left after pruning - post_pruning_row_groups: ExpectedPruning, - } - - 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)]), - post_pruning_row_groups: ExpectedPruning::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, - )]), - post_pruning_row_groups: ExpectedPruning::None, - } - } - - /// Expect all row groups to be pruned - pub fn with_expect_all_pruned(mut self) -> Self { - self.post_pruning_row_groups = ExpectedPruning::All; - self - } - - /// Expect all row groups not to be pruned - pub fn with_expect_none_pruned(mut self) -> Self { - self.post_pruning_row_groups = ExpectedPruning::None; - 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, - post_pruning_row_groups, - } = self; - - 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, - ) - .await - .unwrap(); - - post_pruning_row_groups.assert(&pruned_row_groups); - } - } - - /// Evaluates the pruning predicate on the specified row groups and returns the row groups that are left - async fn test_row_group_bloom_filter_pruning_predicate( - file_name: &str, - data: bytes::Bytes, - pruning_predicate: &PruningPredicate, - ) -> Result<RowGroupAccessPlanFilter> { - use datafusion_datasource::PartitionedFile; - use object_store::ObjectMeta; - - let object_meta = ObjectMeta { - location: object_store::path::Path::parse(file_name).expect("creating path"), - last_modified: chrono::DateTime::from(std::time::SystemTime::now()), - size: data.len() as u64, - e_tag: None, - version: None, - }; - let in_memory = object_store::memory::InMemory::new(); - in_memory - .put(&object_meta.location, data.into()) - .await - .expect("put parquet file into in memory object store"); - - let metrics = ExecutionPlanMetricsSet::new(); - let file_metrics = - ParquetFileMetrics::new(0, object_meta.location.as_ref(), &metrics); - let inner = - ParquetObjectReader::new(Arc::new(in_memory), object_meta.location.clone()) - .with_file_size(object_meta.size); - - let partitioned_file = PartitionedFile::new_from_meta(object_meta); - - let reader = ParquetFileReader { - inner, - file_metrics: file_metrics.clone(), - partitioned_file, - }; - let mut builder = ParquetRecordBatchStreamBuilder::new(reader).await.unwrap(); - - let access_plan = ParquetAccessPlan::new_all(builder.metadata().num_row_groups()); - let mut pruned_row_groups = RowGroupAccessPlanFilter::new(access_plan); - let literal_columns = pruning_predicate.literal_columns(); - let parquet_columns: Vec<_> = literal_columns - .into_iter() - .filter_map(|column_name| { - let (column_idx, _) = parquet::arrow::parquet_column( - builder.parquet_schema(), - pruning_predicate.schema(), - &column_name, - )?; - Some(( - column_name.to_string(), - column_idx, - builder.parquet_schema().column(column_idx).physical_type(), - )) - }) - .collect::<Vec<_>>(); - let mut row_group_bloom_filters = - Vec::with_capacity(builder.metadata().num_row_groups()); - row_group_bloom_filters.resize_with( - builder.metadata().num_row_groups(), - BloomFilterStatistics::new, - ); - for idx in pruned_row_groups.row_group_indexes() { - let mut column_sbbf = HashMap::with_capacity(parquet_columns.len()); - for (column_name, column_idx, physical_type) in &parquet_columns { - let bf = match builder - .get_row_group_column_bloom_filter(idx, *column_idx) - .await - { - Ok(Some(bf)) => bf, - Ok(None) => continue, - Err(e) => { - log::debug!("Ignoring error reading bloom filter: {e}"); - file_metrics.predicate_evaluation_errors.add(1); - continue; - } - }; - column_sbbf.insert(column_name.clone(), (bf, *physical_type)); - } - row_group_bloom_filters[idx] = BloomFilterStatistics { column_sbbf }; - } - pruned_row_groups.prune_by_bloom_filters( - pruning_predicate, - &file_metrics, - &row_group_bloom_filters, - ); - - Ok(pruned_row_groups) - } } diff --git a/datafusion/datasource-parquet/src/test_util.rs b/datafusion/datasource-parquet/src/test_util.rs new file mode 100644 index 0000000000..db494801f6 --- /dev/null +++ b/datafusion/datasource-parquet/src/test_util.rs @@ -0,0 +1,71 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Test helpers shared across the Parquet datasource crate. + +use crate::row_group_filter::RowGroupAccessPlanFilter; + +/// What row groups are expected to be left after pruning. +#[derive(Debug)] +pub(crate) enum ExpectedPruning { + /// All row groups are expected to be pruned. + All, + /// Only the specified row groups are expected to REMAIN (not what is pruned). + Some(Vec<usize>), + /// No row groups are expected to be pruned. + None, +} + +impl ExpectedPruning { + /// Asserts that the pruned row groups match this expectation. + pub(crate) fn assert(&self, row_groups: &RowGroupAccessPlanFilter) { + let access_plan = row_groups.access_plan(); + let num_row_groups = access_plan.len(); + assert!(num_row_groups > 0); + let num_pruned = (0..num_row_groups) + .filter_map(|i| { + if access_plan.should_scan(i) { + None + } else { + Some(1) + } + }) + .sum::<usize>(); + + match self { + Self::All => { + assert_eq!( + num_row_groups, num_pruned, + "Expected all row groups to be pruned, but got {row_groups:?}" + ); + } + ExpectedPruning::None => { + assert_eq!( + num_pruned, 0, + "Expected no row groups to be pruned, but got {row_groups:?}" + ); + } + ExpectedPruning::Some(expected) => { + let actual = access_plan.row_group_indexes(); + assert_eq!( + expected, &actual, + "Unexpected row groups pruned. Expected {expected:?}, got {actual:?}" + ); + } + } + } +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
