You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/05/19 12:03:58 UTC
(datafusion) branch gh-readonly-queue/main/pr-22348-c8b784a01f5d0bcbe0dac806730fb61afc0be8ef created (now 7802abc0d9)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-22348-c8b784a01f5d0bcbe0dac806730fb61afc0be8ef
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 7802abc0d9 refactor(parquet-datasource): split bloom_filter out of row_group_filter.rs (#22348)
This branch includes the following new commits:
new 7802abc0d9 refactor(parquet-datasource): split bloom_filter out of row_group_filter.rs (#22348)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: refactor(parquet-datasource): split bloom_filter out of row_group_filter.rs (#22348)
Posted by gi...@apache.org.
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 <17...@users.noreply.github.com>
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) <no...@anthropic.com>
---
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: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org