From 2accb64fb0fa793ea57296936acff9d2d4bee2df Mon Sep 17 00:00:00 2001 From: mohitgurav20 Date: Fri, 23 Oct 2026 23:09:54 +0530 Subject: [PATCH 1/3] bench: add parquet bloom filter pruning benchmarks --- datafusion/core/Cargo.toml | 5 + .../benches/parquet_bloom_filter_pruning.rs | 227 ++++++++++++++++++ 2 files changed, 232 insertions(+) create mode 100644 datafusion/core/benches/parquet_bloom_filter_pruning.rs diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 89f42761535e3..c9ec6950af192 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -274,6 +274,11 @@ harness = false name = "parquet_query_sql" required-features = ["parquet"] +[[bench]] +harness = false +name = "parquet_bloom_filter_pruning" +required-features = ["parquet"] + [[bench]] harness = false name = "parquet_struct_query" diff --git a/datafusion/core/benches/parquet_bloom_filter_pruning.rs b/datafusion/core/benches/parquet_bloom_filter_pruning.rs new file mode 100644 index 0000000000000..9145593fa98c6 --- /dev/null +++ b/datafusion/core/benches/parquet_bloom_filter_pruning.rs @@ -0,0 +1,227 @@ +// 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. + +//! Benchmarks for Parquet row-group pruning with Bloom filters on several +//! predicate columns. +//! +//! The file has [`ROW_GROUPS`] row groups and [`MAX_COLUMNS`] `Int32` columns, +//! all with Bloom filters and without min/max statistics, so only Bloom +//! filters can prune. Row 0 of every row group holds `0` in every column, and +//! `-1` appears nowhere. Each query is a conjunction of equality predicates on +//! the first `K` columns: +//! +//! - `reject_first`: `c0 = -1 AND c1 = 0 AND ...`; the first column's filter +//! can reject every row group. +//! - `reject_last`: `c0 = 0 AND ... AND c{K-1} = -1`; only the last column's +//! filter can reject, so every filter is loaded before pruning. +//! - `survive_all`: `c0 = 0 AND ... AND c{K-1} = 0`; no filter can reject, so +//! every filter is loaded and every row group is scanned. +//! +//! Each query runs against an in-memory store and against a store that adds +//! [`STORE_LATENCY`] to every request, which approximates remote storage. +//! +//! Run with `cargo bench -p datafusion --features parquet --bench parquet_bloom_filter_pruning`. + +use std::sync::Arc; +use std::time::Duration; + +use arrow::array::{ArrayRef, Int32Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BenchmarkId, Criterion, SamplingMode, criterion_group, criterion_main}; +use datafusion::datasource::file_format::parquet::ParquetFormat; +use datafusion::datasource::listing::ListingOptions; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_execution::runtime_env::RuntimeEnv; +use object_store::memory::InMemory; +use object_store::path::Path; +use object_store::throttle::{ThrottleConfig, ThrottledStore}; +use object_store::{ObjectStore, ObjectStoreExt}; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::{EnabledStatistics, WriterProperties}; +use tokio::runtime::Runtime; +use url::Url; + +const ROW_GROUPS: usize = 64; +const ROW_GROUP_ROWS: usize = 1024; +const MAX_COLUMNS: usize = 32; +const COLUMN_COUNTS: [usize; 3] = [2, 8, 32]; +const STORE_LATENCY: Duration = Duration::from_micros(200); +const PRESENT: i32 = 0; +const ABSENT: i32 = -1; + +fn schema() -> SchemaRef { + Arc::new(Schema::new( + (0..MAX_COLUMNS) + .map(|i| Field::new(format!("c{i}"), DataType::Int32, false)) + .collect::>(), + )) +} + +fn create_parquet_file() -> Vec { + let schema = schema(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(ROW_GROUP_ROWS)) + .set_statistics_enabled(EnabledStatistics::None) + .set_bloom_filter_enabled(true) + .set_bloom_filter_max_ndv(ROW_GROUP_ROWS as u64) + .set_bloom_filter_fpp(0.01) + .build(); + + let mut data = Vec::new(); + let mut writer = + ArrowWriter::try_new(&mut data, Arc::clone(&schema), Some(props)).unwrap(); + for rg in 0..ROW_GROUPS { + // Row 0 is PRESENT in every row group; other values are unique and positive. + let values: Int32Array = (0..ROW_GROUP_ROWS) + .map(|row| { + if row == 0 { + PRESENT + } else { + (rg * ROW_GROUP_ROWS + row) as i32 + } + }) + .collect(); + let column: ArrayRef = Arc::new(values); + let columns = vec![column; MAX_COLUMNS]; + let batch = RecordBatch::try_new(Arc::clone(&schema), columns).unwrap(); + writer.write(&batch).unwrap(); + } + writer.close().unwrap(); + data +} + +fn make_store_slow(object_store: InMemory) -> Arc { + let config = ThrottleConfig { + wait_get_per_call: STORE_LATENCY, + ..Default::default() + }; + Arc::new(ThrottledStore::new(object_store, config)) +} + +async fn setup_context() -> SessionContext { + let data = create_parquet_file(); + let stores: [(&str, Arc); 2] = [ + ("memory", Arc::new(InMemory::new())), + ("slow", make_store_slow(InMemory::new())), + ]; + + // One partition keeps the Bloom filter reads for the file in one opener. + let config = SessionConfig::new().with_target_partitions(1); + let rt = Arc::new(RuntimeEnv::default()); + for (name, store) in &stores { + store + .put(&Path::from("t/data.parquet"), data.clone().into()) + .await + .unwrap(); + rt.register_object_store( + &Url::parse(&format!("{name}://bucket")).unwrap(), + Arc::clone(store), + ); + } + let ctx = SessionContext::new_with_config_rt(config, rt); + + for (name, _) in &stores { + let options = ListingOptions::new(Arc::new(ParquetFormat::default())); + ctx.register_listing_table( + *name, + format!("{name}://bucket/t/"), + options, + Some(schema()), + None, + ) + .await + .unwrap(); + } + ctx +} + +/// Builds the query and its expected row count. +fn query(table: &str, case: &str, columns: usize) -> (String, usize) { + let predicate = (0..columns) + .map(|i| { + let rejecting = match case { + "reject_first" => i == 0, + "reject_last" => i == columns - 1, + "survive_all" => false, + _ => unreachable!(), + }; + let value = if rejecting { ABSENT } else { PRESENT }; + format!("c{i} = {value}") + }) + .collect::>() + .join(" AND "); + let expected_rows = if case == "survive_all" { ROW_GROUPS } else { 0 }; + ( + format!("SELECT c0 FROM {table} WHERE {predicate}"), + expected_rows, + ) +} + +async fn run(ctx: &SessionContext, sql: &str, expected_rows: usize) { + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, expected_rows); +} + +/// Prints the Bloom filter pruning metric once so a run shows whether Bloom +/// filters actually pruned. +async fn print_bloom_metrics(ctx: &SessionContext, sql: &str) { + let batches = ctx + .sql(&format!("EXPLAIN ANALYZE {sql}")) + .await + .unwrap() + .collect() + .await + .unwrap(); + let plan = arrow::util::pretty::pretty_format_batches(&batches) + .unwrap() + .to_string(); + let metric = plan + .split([',', '[']) + .find(|s| s.contains("row_groups_pruned_bloom_filter")) + .unwrap_or("row_groups_pruned_bloom_filter not reported"); + println!("{sql}\n {}", metric.trim()); +} + +fn criterion_benchmark(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + let ctx = rt.block_on(setup_context()); + + for case in ["reject_first", "reject_last", "survive_all"] { + for table in ["memory", "slow"] { + let mut group = + c.benchmark_group(format!("parquet_bloom_filter_pruning/{table}/{case}")); + group.sampling_mode(SamplingMode::Flat); + group.sample_size(10); + for columns in COLUMN_COUNTS { + let (sql, expected_rows) = query(table, case, columns); + rt.block_on(print_bloom_metrics(&ctx, &sql)); + group.bench_with_input( + BenchmarkId::new("columns", columns), + &sql, + |b, sql| { + b.to_async(&rt).iter(|| run(&ctx, sql, expected_rows)); + }, + ); + } + group.finish(); + } + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); From 9b4700c9932771c017e6f72893d170e7ddff2c82 Mon Sep 17 00:00:00 2001 From: mohitgurav20 Date: Fri, 23 Oct 2026 23:34:24 +0530 Subject: [PATCH 2/3] perf(datasource-parquet): evaluate Parquet Bloom filters incrementally during row group pruning - Evaluate Bloom filters column-by-column incrementally to short-circuit multi-column pruning before loading remaining column filters. - Re-evaluate only newly loaded column guarantees to avoid redundant full-predicate tree evaluations. - Preserve upstream early-exit Done state when all row groups are pruned. - Add unit tests covering full matching, stats pruning, and I/O error fallback. --- .../datasource-parquet/src/bloom_filter.rs | 86 ++- .../datasource-parquet/src/opener/mod.rs | 593 +++++++++++++++--- .../src/row_group_filter.rs | 5 + 3 files changed, 574 insertions(+), 110 deletions(-) diff --git a/datafusion/datasource-parquet/src/bloom_filter.rs b/datafusion/datasource-parquet/src/bloom_filter.rs index 2384f95df9885..f6e30fac5097d 100644 --- a/datafusion/datasource-parquet/src/bloom_filter.rs +++ b/datafusion/datasource-parquet/src/bloom_filter.rs @@ -85,6 +85,25 @@ impl BloomFilterStatistics { ); } + /// Returns `true` if the Bloom filter definitely does not contain any of the given values. + /// + /// If all values are confirmed absent by the Bloom filter, returns `true`. + /// If any value might be present (or checking is unsupported), returns `false`. + #[allow(clippy::allow_attributes, clippy::mutable_key_type)] // ScalarValue has interior mutability but is intentionally used as hash key + pub fn is_known_not_present( + sbbf: &Sbbf, + values: &HashSet, + parquet_type: &Type, + type_length: i32, + ) -> bool { + if values.is_empty() { + return false; + } + values + .iter() + .all(|value| !Self::check_scalar(sbbf, value, parquet_type, type_length)) + } + /// Helper function for checking if [`Sbbf`] filter contains [`ScalarValue`]. /// /// In case the type of scalar is not supported, returns `true`, assuming that the @@ -209,20 +228,12 @@ impl PruningStatistics for BloomFilterStatistics { // 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( - &column_bloom_filter.sbbf, - value, - &column_bloom_filter.physical_type, - column_bloom_filter.type_length, - ) - }) - // The row group doesn't contain any of the values if - // all the checks are false - .all(|v| !v); + let known_not_present = Self::is_known_not_present( + &column_bloom_filter.sbbf, + values, + &column_bloom_filter.physical_type, + column_bloom_filter.type_length, + ); let contains = if known_not_present { Some(false) @@ -720,4 +731,51 @@ mod tests { Ok(pruned_row_groups) } + + #[test] + #[allow(clippy::allow_attributes, clippy::mutable_key_type)] // ScalarValue has interior mutability but is intentionally used as hash key + fn test_is_known_not_present() { + let mut sbbf = Sbbf::new_with_ndv_fpp(10, 0.01).unwrap(); + // 1. Empty values -> returns false + assert!(!BloomFilterStatistics::is_known_not_present( + &sbbf, + &HashSet::new(), + &Type::INT32, + 0, + )); + + // 2. Insert 42 + sbbf.insert(&42_i32); + + // Value 42 is present -> is_known_not_present should be false + let mut present_set = HashSet::new(); + present_set.insert(ScalarValue::Int32(Some(42))); + assert!(!BloomFilterStatistics::is_known_not_present( + &sbbf, + &present_set, + &Type::INT32, + 0, + )); + + // Value 99 is absent -> is_known_not_present should be true + let mut absent_set = HashSet::new(); + absent_set.insert(ScalarValue::Int32(Some(99))); + assert!(BloomFilterStatistics::is_known_not_present( + &sbbf, + &absent_set, + &Type::INT32, + 0, + )); + + // Mixed set containing present and absent -> is_known_not_present should be false + let mut mixed_set = HashSet::new(); + mixed_set.insert(ScalarValue::Int32(Some(42))); + mixed_set.insert(ScalarValue::Int32(Some(99))); + assert!(!BloomFilterStatistics::is_known_not_present( + &sbbf, + &mixed_set, + &Type::INT32, + 0, + )); + } } diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index 3a1766869d319..0928e87b775dd 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -44,7 +44,7 @@ use datafusion_datasource::morsel::{Morsel, MorselPlan, MorselPlanner, Morselize use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr_adapter::replace_columns_with_literals; use datafusion_physical_expr_adapter::rewrite::rewrite_input_file_name_in_projection; -use std::collections::{HashMap, VecDeque}; +use std::collections::{HashMap, HashSet as StdHashSet, VecDeque}; use std::fmt; use std::future::Future; use std::mem; @@ -62,7 +62,7 @@ use datafusion_common::{ use datafusion_datasource::{PartitionedFile, TableSchema}; use datafusion_physical_expr::expressions::{Column, DynamicFilterTracking, Literal}; use datafusion_physical_expr::simplifier::PhysicalExprSimplifier; -use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_expr::utils::{Guarantee, collect_columns}; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::LexOrdering; @@ -363,9 +363,6 @@ impl Morselizer for ParquetMorselizer { /// LoadBloomFilters /// | /// v -/// PruneWithBloomFilters -/// | -/// v /// BuildStream /// | /// v @@ -397,9 +394,7 @@ enum ParquetOpenState { /// Loading [Parquet Page Index](https://parquet.apache.org/docs/file-format/pageindex/) LoadPageIndex(BoxFuture<'static, Result>), /// Loading bloom filters required for row-group pruning - LoadBloomFilters(BoxFuture<'static, Result>), - /// Pruning with preloaded Bloom Filters - PruneWithBloomFilters(Box), + LoadBloomFilters(BoxFuture<'static, Result>), /// Builds the final reader stream /// /// TODO: split state as this currently does both I/O and CPU work. @@ -422,7 +417,6 @@ impl fmt::Debug for ParquetOpenState { ParquetOpenState::LoadPageIndex(_) => "LoadPageIndex", ParquetOpenState::PruneWithStatistics(_) => "PruneWithStatistics", ParquetOpenState::LoadBloomFilters(_) => "LoadBloomFilters", - ParquetOpenState::PruneWithBloomFilters(_) => "PruneWithBloomFilters", ParquetOpenState::BuildStream(_) => "BuildStream", ParquetOpenState::Ready(_) => "Ready", ParquetOpenState::Done => "Done", @@ -596,16 +590,6 @@ impl DecoderReadPlans { } /// State of [`ParquetOpenState`] -/// -/// Result of loading bloom filters needed for row-group pruning. -struct BloomFiltersLoadedParquetOpen { - prepared: RowGroupsPrunedParquetOpen, - /// Bloom filters loaded for each row group that remains under consideration. - /// - /// indexed by parquet row-group index - row_group_bloom_filters: Vec, -} - impl ParquetOpenState { /// Applies one CPU-only state transition. /// @@ -693,13 +677,6 @@ impl ParquetOpenState { ParquetOpenState::LoadBloomFilters(future) => { Ok(ParquetOpenState::LoadBloomFilters(future)) } - ParquetOpenState::PruneWithBloomFilters(loaded) => { - let pruned = loaded.prune_bloom_filters(); - if pruned.row_groups.is_fully_pruned() { - return Ok(ParquetOpenState::Done); - } - Ok(ParquetOpenState::BuildStream(Box::new(pruned))) - } ParquetOpenState::BuildStream(prepared) => { Ok(ParquetOpenState::Ready(prepared.build_stream()?)) } @@ -814,9 +791,11 @@ impl MorselPlanner for ParquetMorselPlanner { } ParquetOpenState::LoadBloomFilters(future) => { Ok(Some(Self::schedule_io(async move { - Ok(ParquetOpenState::PruneWithBloomFilters(Box::new( - future.await?, - ))) + let pruned = future.await?; + if pruned.row_groups.is_fully_pruned() { + return Ok(ParquetOpenState::Done); + } + Ok(ParquetOpenState::BuildStream(Box::new(pruned))) }))) } ParquetOpenState::Ready(stream) => { @@ -1479,16 +1458,7 @@ impl RowGroupsPrunedParquetOpen { } /// Load bloom filters needed for pruning when enabled and a pruning predicate exists. - async fn load_bloom_filters(mut self) -> Result { - let num_row_groups = self - .prepared - .loaded - .reader_metadata - .metadata() - .num_row_groups(); - let mut row_group_bloom_filters = - vec![BloomFilterStatistics::new(); num_row_groups]; - + async fn load_bloom_filters(mut self) -> Result { if let Some(predicate) = self.prepared.pruning_predicate.as_ref().map(|p| p.as_ref()) && self.prepared.loaded.prepared.enable_bloom_filter @@ -1498,6 +1468,14 @@ impl RowGroupsPrunedParquetOpen { .row_group_indexes() .any(|idx| !self.row_groups.access_plan().is_fully_matched(idx)) { + let bloom_filter_eval_time = self + .prepared + .loaded + .prepared + .file_metrics + .bloom_filter_eval_time + .clone(); + // Use the existing reader for bloom filter I/O; // replace with a fresh reader for decoding below. let reader_metadata = self.prepared.loaded.reader_metadata.clone(); @@ -1516,7 +1494,14 @@ impl RowGroupsPrunedParquetOpen { mem::replace(&mut prepared.async_file_reader, replacement_reader), reader_metadata, ); - let parquet_columns: Vec<(String, usize, Type, i32)> = predicate + struct BloomColumnGuarantee { + column_idx: usize, + physical_type: Type, + type_length: i32, + in_guarantees: Vec>, + } + + let parquet_columns: Vec = predicate .literal_columns() .into_iter() .filter_map(|column_name| { @@ -1526,28 +1511,45 @@ impl RowGroupsPrunedParquetOpen { &prepared.physical_file_schema, &column_name, )?; - Some(( - column_name, + let in_guarantees: Vec> = predicate + .literal_guarantees() + .iter() + .filter(|lg| { + lg.column.name == column_name && lg.guarantee == Guarantee::In + }) + .map(|lg| lg.literals.clone()) + .collect(); + Some(BloomColumnGuarantee { column_idx, - parquet_schema.column(column_idx).physical_type(), - parquet_schema.column(column_idx).type_length(), - )) + physical_type: parquet_schema.column(column_idx).physical_type(), + type_length: parquet_schema.column(column_idx).type_length(), + in_guarantees, + }) }) .collect(); - for idx in self.row_groups.row_group_indexes() { - // Statistics have already proved that every row in this group - // matches the predicate, so its Bloom filters cannot prune it. + // Collect up front so the iterator's immutable borrow on + // `self.row_groups` is released before we mutably call + // `access_plan_mut().skip()` inside the loop body. + let row_group_indexes: Vec = + self.row_groups.row_group_indexes().collect(); + for idx in row_group_indexes { + if !self.row_groups.access_plan().should_scan(idx) { + continue; + } + // Statistics have already proved every row in this group matches + // the predicate, so Bloom filters cannot prune it — skip early. if self.row_groups.access_plan().is_fully_matched(idx) { continue; } - let mut row_group_filters = - BloomFilterStatistics::with_capacity(parquet_columns.len()); - for (column_name, column_idx, physical_type, type_length) in - &parquet_columns - { + // Track whether at least one Bloom filter was loaded for this row + // group so we can correctly gate the matched/pruned metric updates. + let mut any_filter_loaded = false; + let mut prune_group = false; + + for col_guarantee in &parquet_columns { let bf: Sbbf = match builder - .get_row_group_column_bloom_filter(idx, *column_idx) + .get_row_group_column_bloom_filter(idx, col_guarantee.column_idx) .await { Ok(Some(bf)) => bf, @@ -1558,53 +1560,56 @@ impl RowGroupsPrunedParquetOpen { continue; } }; - row_group_filters.insert( - column_name, - bf, - *physical_type, - *type_length, - ); + any_filter_loaded = true; + + // Evaluate only the literal guarantees for this newly loaded column. + // + // Calling `predicate.prune()` after each column would repeatedly + // evaluate the full expression tree and rebuild statistics batches, + // incurring heavy O(N) overhead per column (O(N²) per row group). + // Because Bloom filters can only prune row groups via literal + // `Guarantee::In` guarantees (proving none of the required values + // exist in the filter), checking only the guarantees for this newly + // loaded column preserves early rejection while completely eliminating + // redundant evaluation overhead. + let _timer_guard = bloom_filter_eval_time.timer(); + let prunable = col_guarantee.in_guarantees.iter().any(|literals| { + BloomFilterStatistics::is_known_not_present( + &bf, + literals, + &col_guarantee.physical_type, + col_guarantee.type_length, + ) + }); + + if prunable { + prune_group = true; + break; + } } - row_group_bloom_filters[idx] = row_group_filters; - } - } - Ok(BloomFiltersLoadedParquetOpen { - prepared: self, - row_group_bloom_filters, - }) - } -} + if !any_filter_loaded { + // No Bloom filters were available for this row group; evaluation was + // not performed, so do not update matched or pruned metrics. + continue; + } -impl BloomFiltersLoadedParquetOpen { - /// Apply bloom filter pruning using already loaded bloom filters. - fn prune_bloom_filters(mut self) -> RowGroupsPrunedParquetOpen { - let bloom_filter_eval_time = self - .prepared - .prepared - .loaded - .prepared - .file_metrics - .bloom_filter_eval_time - .clone(); - let _timer_guard = bloom_filter_eval_time.timer(); - if let Some(predicate) = self - .prepared - .prepared - .pruning_predicate - .as_ref() - .map(|p| p.as_ref()) - && self.prepared.prepared.loaded.prepared.enable_bloom_filter - && !self.prepared.row_groups.is_empty() - { - self.prepared.row_groups.prune_by_bloom_filters( - predicate, - &self.prepared.prepared.loaded.prepared.file_metrics, - &self.row_group_bloom_filters, - ); + if prune_group { + prepared + .file_metrics + .row_groups_pruned_bloom_filter + .add_pruned(1); + self.row_groups.access_plan_mut().skip(idx); + } else { + prepared + .file_metrics + .row_groups_pruned_bloom_filter + .add_matched(1); + } + } } - self.prepared + Ok(self) } } @@ -4879,6 +4884,402 @@ mod test { assert_eq!(matched, 0); } + #[tokio::test] + async fn incremental_bloom_filter_pruning_short_circuits_remaining_columns() { + let store = Arc::new(InMemory::new()) as Arc; + // Two row groups, each with columns a and b: + // RG-0: a = [99, 98, 97], b = [900, 980, 970] — no predicate values present + // RG-1: a = [10, 20, 30], b = [100, 200, 300] — contains predicate match + // + // Statistics pruning is disabled throughout; Bloom filters are the only pruning + // mechanism. Decoder-level filter push-down is enabled so matching rows are + // physically filtered out of RG-1 where possible. + let batch_rg0 = record_batch!( + ("a", Int32, vec![99_i32, 98, 97]), + ("b", Int32, vec![900_i32, 980, 970]) + ) + .unwrap(); + let batch_rg1 = record_batch!( + ("a", Int32, vec![10_i32, 20, 30]), + ("b", Int32, vec![100_i32, 200, 300]) + ) + .unwrap(); + let schema = batch_rg0.schema(); + let props = WriterProperties::builder() + // Force each batch into its own row group. + .set_max_row_group_row_count(Some(3)) + .set_bloom_filter_enabled(true) + .set_statistics_enabled(EnabledStatistics::None) + .build(); + let data_size = write_parquet_batches( + Arc::clone(&store), + "multi_col_bloom.parquet", + vec![batch_rg0, batch_rg1], + Some(props), + ) + .await; + let file = + PartitionedFile::new("multi_col_bloom.parquet".to_string(), data_size as u64); + + // ── Query 1: a = 999 AND b = 100 ───────────────────────────────────────────── + // 'a' Bloom filter is checked first for each row group. 999 is absent from + // both, so both groups are pruned after reading only the 'a' filter — the 'b' + // filter is never fetched (short-circuit). + let pred1 = + logical2physical(&col("a").eq(lit(999)).and(col("b").eq(lit(100))), &schema); + let metrics1 = ExecutionPlanMetricsSet::new(); + let morselizer1 = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(Arc::clone(&schema)) + .with_predicate(Arc::clone(&pred1)) + .with_pushdown_filters(true) + .with_row_group_stats_pruning(false) + .with_enable_bloom_filter(true) + .with_metrics(metrics1.clone()) + .build(); + let stream1 = open_file(&morselizer1, file.clone()).await.unwrap(); + let bytes_scanned1 = counter_metric_value(&metrics1, "bytes_scanned"); + assert_eq!(collect_int32_values(stream1).await, Vec::::new()); + let (pruned1, matched1) = bloom_filter_pruning_metrics(&metrics1); + // Both row groups pruned; no row group survived Bloom evaluation. + assert_eq!(pruned1, 2); + assert_eq!(matched1, 0); + + // ── Query 2: a = 10 AND b = 999 ────────────────────────────────────────────── + // RG-0: 'a' filter doesn't contain 10 → pruned immediately (no 'b' read). + // RG-1: 'a' filter contains 10 → must read 'b'; 'b' filter doesn't contain 999 + // → pruned on the second column. + let pred2 = + logical2physical(&col("a").eq(lit(10)).and(col("b").eq(lit(999))), &schema); + let metrics2 = ExecutionPlanMetricsSet::new(); + let morselizer2 = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(Arc::clone(&schema)) + .with_predicate(Arc::clone(&pred2)) + .with_pushdown_filters(true) + .with_row_group_stats_pruning(false) + .with_enable_bloom_filter(true) + .with_metrics(metrics2.clone()) + .build(); + let stream2 = open_file(&morselizer2, file.clone()).await.unwrap(); + let bytes_scanned2 = counter_metric_value(&metrics2, "bytes_scanned"); + assert_eq!(collect_int32_values(stream2).await, Vec::::new()); + let (pruned2, matched2) = bloom_filter_pruning_metrics(&metrics2); + assert_eq!(pruned2, 2); + assert_eq!(matched2, 0); + + // Query 2 must read both 'a' and 'b' filters for RG-1, so it scans strictly + // more Bloom filter bytes than Query 1 (which short-circuits on 'a' alone). + assert!(bytes_scanned1 > 0, "expected non-zero bloom filter I/O"); + assert!( + bytes_scanned2 > bytes_scanned1, + "query 2 should scan more bloom filter bytes due to loading 'b' filter for RG-1" + ); + + // ── Query 3: a = 10 (one row group pruned, one matched) ─────────────────────── + // RG-0: 'a' filter doesn't contain 10 → pruned. + // RG-1: 'a' filter contains 10 → matched (not pruned); the actual matching + // rows [10, 20, 30] are returned through decoder-level filtering. + let pred3 = logical2physical(&col("a").eq(lit(10)), &schema); + let metrics3 = ExecutionPlanMetricsSet::new(); + let morselizer3 = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(Arc::clone(&schema)) + .with_predicate(Arc::clone(&pred3)) + .with_pushdown_filters(true) + .with_row_group_stats_pruning(false) + .with_enable_bloom_filter(true) + .with_metrics(metrics3.clone()) + .build(); + let stream3 = open_file(&morselizer3, file.clone()).await.unwrap(); + // The matching row (a=10) must be returned; decoder filtering removes a=20,30. + assert_eq!(collect_int32_values(stream3).await, vec![10_i32]); + let (pruned3, matched3) = bloom_filter_pruning_metrics(&metrics3); + assert_eq!(pruned3, 1, "RG-0 should be pruned by Bloom filter"); + assert_eq!( + matched3, 1, + "RG-1 passed Bloom evaluation and must count as matched" + ); + } + + /// Verify that a Bloom filter available for only one predicate column is still + /// sufficient to prune a row group, and that the matched/pruned metrics behave + /// correctly when filters are partially absent. + /// + /// Setup: predicate is `a = 10 AND b = 999`. The Parquet file has a Bloom filter + /// on column `b` only (column `a` has no filter). + /// RG-0: a = [1, 2, 3], b = [10, 20, 30] — 999 absent from `b` → pruned + /// RG-1: a = [10, 11, 12], b = [999, 888, 777] — both values might be present + // (Bloom false positives aside) → survives, `matched` is incremented once + #[tokio::test] + async fn bloom_filter_pruning_with_partial_column_availability() { + use parquet::schema::types::ColumnPath; + + let store = Arc::new(InMemory::new()) as Arc; + let batch_rg0 = record_batch!( + ("a", Int32, vec![1_i32, 2, 3]), + ("b", Int32, vec![10_i32, 20, 30]) + ) + .unwrap(); + let batch_rg1 = record_batch!( + ("a", Int32, vec![10_i32, 11, 12]), + ("b", Int32, vec![999_i32, 888, 777]) + ) + .unwrap(); + let schema = batch_rg0.schema(); + + // Enable Bloom filter only for column `b`; column `a` intentionally has no filter. + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .set_statistics_enabled(EnabledStatistics::None) + // Disable the global default first, then opt column `b` in. + .set_bloom_filter_enabled(false) + .set_column_bloom_filter_enabled(ColumnPath::from("b"), true) + .build(); + let data_size = write_parquet_batches( + Arc::clone(&store), + "partial_bloom.parquet", + vec![batch_rg0, batch_rg1], + Some(props), + ) + .await; + let file = + PartitionedFile::new("partial_bloom.parquet".to_string(), data_size as u64); + + let predicate = + logical2physical(&col("a").eq(lit(10)).and(col("b").eq(lit(999))), &schema); + let metrics = ExecutionPlanMetricsSet::new(); + let morselizer = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(Arc::clone(&schema)) + .with_predicate(Arc::clone(&predicate)) + .with_pushdown_filters(true) + .with_row_group_stats_pruning(false) + .with_enable_bloom_filter(true) + .with_metrics(metrics.clone()) + .build(); + let stream = open_file(&morselizer, file.clone()).await.unwrap(); + + // RG-1 (a=10, b=999) survives Bloom evaluation; the filter push-down produces + // the matching row. Return value is column `a` only (our collect helper). + let values = collect_int32_values(stream).await; + assert_eq!(values, vec![10_i32]); + + let (pruned, matched) = bloom_filter_pruning_metrics(&metrics); + // RG-0: `b` filter is present but does not contain 999 → pruned. + assert_eq!( + pruned, 1, + "RG-0 should be pruned via column b's Bloom filter" + ); + // RG-1: `b` filter is present and 999 may be present (Bloom says yes) → matched. + assert_eq!( + matched, 1, + "RG-1 survived Bloom evaluation and must be counted as matched" + ); + + // ── Query 2: all filters absent for predicate columns ─────────────────────── + // Column `a` has no Bloom filter in this file. A query on `a = 10` alone + // finds no Bloom filters for any row group. Pruning cannot be performed, + // so neither pruned nor matched metric is incremented. + let predicate_no_filter = logical2physical(&col("a").eq(lit(10)), &schema); + let metrics_no_filter = ExecutionPlanMetricsSet::new(); + let morselizer_no_filter = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(Arc::clone(&schema)) + .with_predicate(Arc::clone(&predicate_no_filter)) + .with_pushdown_filters(true) + .with_row_group_stats_pruning(false) + .with_enable_bloom_filter(true) + .with_metrics(metrics_no_filter.clone()) + .build(); + let stream_no_filter = open_file(&morselizer_no_filter, file).await.unwrap(); + let values_no_filter = collect_int32_values(stream_no_filter).await; + assert_eq!(values_no_filter, vec![10_i32]); + + let (pruned_nf, matched_nf) = bloom_filter_pruning_metrics(&metrics_no_filter); + assert_eq!( + pruned_nf, 0, + "No Bloom filters were loaded, so pruned metric should not be incremented" + ); + assert_eq!( + matched_nf, 0, + "No Bloom filters were loaded, so matched metric should not be incremented" + ); + } + + /// Verify that Bloom filter evaluation respects prior row group pruning and + /// fully matched row groups: + /// - RG-0 is fully matched by statistics (`[10, 10, 10]` for `a = 10`), so it skips Bloom evaluation. + /// - RG-1 contains `[10, 20, 30]` (partially matched), so it evaluates Bloom filter and matches. + /// - RG-2 is pruned by statistics (`[100, 200, 300]`), so it is completely skipped (`!should_scan`). + #[tokio::test] + async fn bloom_filter_pruning_respects_statistics_pruned_and_fully_matched_row_groups() + { + let store = Arc::new(InMemory::new()) as Arc; + let batches = vec![ + record_batch!(("a", Int32, vec![10, 10, 10])).unwrap(), + record_batch!(("a", Int32, vec![10, 20, 30])).unwrap(), + record_batch!(("a", Int32, vec![100, 200, 300])).unwrap(), + ]; + let schema = batches[0].schema(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .set_bloom_filter_enabled(true) + .set_statistics_enabled(EnabledStatistics::Chunk) + .build(); + let data_size = write_parquet_batches( + Arc::clone(&store), + "stats_and_bloom.parquet", + batches, + Some(props), + ) + .await; + let file = + PartitionedFile::new("stats_and_bloom.parquet".to_string(), data_size as u64); + + let predicate = logical2physical(&col("a").eq(lit(10)), &schema); + let metrics = ExecutionPlanMetricsSet::new(); + let morselizer = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(Arc::clone(&schema)) + .with_predicate(predicate) + .with_pushdown_filters(true) + .with_row_group_stats_pruning(true) + .with_enable_bloom_filter(true) + .with_metrics(metrics.clone()) + .build(); + + let stream = open_file(&morselizer, file).await.unwrap(); + // RG-0 yields 3 matching rows, RG-1 yields 1 matching row, RG-2 is skipped + assert_eq!(collect_int32_values(stream).await, vec![10, 10, 10, 10]); + + let (pruned, matched) = bloom_filter_pruning_metrics(&metrics); + // RG-0 is fully matched by statistics (Bloom skipped) + // RG-1 evaluates Bloom and matches (matched += 1) + // RG-2 is pruned by statistics (Bloom skipped) + assert_eq!(pruned, 0); + assert_eq!(matched, 1); + } + + struct FailingBloomAsyncReader { + inner: Box, + } + + impl AsyncFileReader for FailingBloomAsyncReader { + fn get_bytes( + &mut self, + _range: std::ops::Range, + ) -> BoxFuture<'_, parquet::errors::Result> { + Box::pin(futures::future::err( + parquet::errors::ParquetError::General( + "injected bloom filter read error".into(), + ), + )) + } + + fn get_byte_ranges( + &mut self, + _ranges: Vec>, + ) -> BoxFuture<'_, parquet::errors::Result>> { + Box::pin(futures::future::err( + parquet::errors::ParquetError::General( + "injected bloom filter read error".into(), + ), + )) + } + + fn get_metadata<'a>( + &'a mut self, + options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, parquet::errors::Result>> { + self.inner.get_metadata(options) + } + } + + #[derive(Debug)] + struct FailingBloomReaderFactory { + base_factory: Arc, + create_count: AtomicUsize, + } + + impl ParquetFileReaderFactory for FailingBloomReaderFactory { + fn create_reader( + &self, + partition_index: usize, + file: PartitionedFile, + metadata_size_hint: Option, + metrics: &ExecutionPlanMetricsSet, + ) -> Result> { + let count = self.create_count.fetch_add(1, Ordering::SeqCst); + let inner = self.base_factory.create_reader( + partition_index, + file, + metadata_size_hint, + metrics, + )?; + if count == 0 { + // First reader: used for metadata (succeeds) and Bloom filter I/O (fails) + Ok(Box::new(FailingBloomAsyncReader { inner })) + } else { + // Replacement reader: used for record batch decoding (succeeds) + Ok(inner) + } + } + } + + #[tokio::test] + async fn bloom_filter_read_error_increments_metric_and_continues() { + let store = Arc::new(InMemory::new()) as Arc; + let batch = record_batch!(("a", Int32, vec![10, 20, 30])).unwrap(); + let schema = batch.schema(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .set_bloom_filter_enabled(true) + .set_statistics_enabled(EnabledStatistics::None) + .build(); + let data_size = write_parquet_batches( + Arc::clone(&store), + "err_bloom.parquet", + vec![batch], + Some(props), + ) + .await; + let file = + PartitionedFile::new("err_bloom.parquet".to_string(), data_size as u64); + + let predicate = logical2physical(&col("a").eq(lit(10)), &schema); + let metrics = ExecutionPlanMetricsSet::new(); + let base_factory = + Arc::new(DefaultParquetFileReaderFactory::new(Arc::clone(&store))); + let failing_factory = Arc::new(FailingBloomReaderFactory { + base_factory, + create_count: AtomicUsize::new(0), + }); + + let morselizer = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(Arc::clone(&schema)) + .with_predicate(predicate) + .with_pushdown_filters(true) + .with_row_group_stats_pruning(false) + .with_enable_bloom_filter(true) + .with_parquet_file_reader_factory(failing_factory) + .with_metrics(metrics.clone()) + .build(); + + let stream = open_file(&morselizer, file).await.unwrap(); + // Since Bloom filter failed with error, scan ignores error and decodes data + assert_eq!(collect_int32_values(stream).await, vec![10]); + // Evaluation error metric must be incremented by 1 + assert_eq!( + counter_metric_value(&metrics, "predicate_evaluation_errors"), + 1 + ); + let (pruned, matched) = bloom_filter_pruning_metrics(&metrics); + assert_eq!(pruned, 0); + assert_eq!(matched, 0); + } + #[test] fn should_load_page_index_with_row_selection() { let mut plan = ParquetAccessPlan::new_all(1); diff --git a/datafusion/datasource-parquet/src/row_group_filter.rs b/datafusion/datasource-parquet/src/row_group_filter.rs index feba768955e18..c73b22afd4589 100644 --- a/datafusion/datasource-parquet/src/row_group_filter.rs +++ b/datafusion/datasource-parquet/src/row_group_filter.rs @@ -125,6 +125,11 @@ impl RowGroupAccessPlanFilter { &self.access_plan } + /// Returns a mutable reference to the inner access plan. + pub(crate) fn access_plan_mut(&mut self) -> &mut ParquetAccessPlan { + &mut self.access_plan + } + /// Returns the is_fully_matched vector. pub fn is_fully_matched(&self) -> &Vec { self.access_plan.fully_matched() From 9e1e3297534d6c54e81fe24cf8351515941b0aaa Mon Sep 17 00:00:00 2001 From: mohitgurav20 Date: Thu, 29 Oct 2026 22:41:08 +0530 Subject: [PATCH 3/3] refactor(bloom): eliminate duplication, encapsulate pruning, drop criterion benchmark - Move BloomColumnGuarantee + guarantee extraction into RowGroupAccessPlanFilter via prepare_bloom_column_guarantees(), so opener only handles async I/O - Add prune_row_group_with_bloom_filter() + record_bloom_filter_matched() to encapsulate per-column evaluation and metric updates - Remove access_plan_mut() escape-hatch; all plan mutations now go through the filter struct itself - Keep prune_by_bloom_filters() as the live public batch evaluation API - Drop parquet_bloom_filter_pruning.rs criterion bench (was just running SQL with println -- wrong tool for the job) --- datafusion/core/Cargo.toml | 5 - .../benches/parquet_bloom_filter_pruning.rs | 227 --------------- .../datasource-parquet/src/opener/mod.rs | 263 +++++++++++------- .../src/row_group_filter.rs | 97 ++++++- 4 files changed, 258 insertions(+), 334 deletions(-) delete mode 100644 datafusion/core/benches/parquet_bloom_filter_pruning.rs diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index c9ec6950af192..89f42761535e3 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -274,11 +274,6 @@ harness = false name = "parquet_query_sql" required-features = ["parquet"] -[[bench]] -harness = false -name = "parquet_bloom_filter_pruning" -required-features = ["parquet"] - [[bench]] harness = false name = "parquet_struct_query" diff --git a/datafusion/core/benches/parquet_bloom_filter_pruning.rs b/datafusion/core/benches/parquet_bloom_filter_pruning.rs deleted file mode 100644 index 9145593fa98c6..0000000000000 --- a/datafusion/core/benches/parquet_bloom_filter_pruning.rs +++ /dev/null @@ -1,227 +0,0 @@ -// 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. - -//! Benchmarks for Parquet row-group pruning with Bloom filters on several -//! predicate columns. -//! -//! The file has [`ROW_GROUPS`] row groups and [`MAX_COLUMNS`] `Int32` columns, -//! all with Bloom filters and without min/max statistics, so only Bloom -//! filters can prune. Row 0 of every row group holds `0` in every column, and -//! `-1` appears nowhere. Each query is a conjunction of equality predicates on -//! the first `K` columns: -//! -//! - `reject_first`: `c0 = -1 AND c1 = 0 AND ...`; the first column's filter -//! can reject every row group. -//! - `reject_last`: `c0 = 0 AND ... AND c{K-1} = -1`; only the last column's -//! filter can reject, so every filter is loaded before pruning. -//! - `survive_all`: `c0 = 0 AND ... AND c{K-1} = 0`; no filter can reject, so -//! every filter is loaded and every row group is scanned. -//! -//! Each query runs against an in-memory store and against a store that adds -//! [`STORE_LATENCY`] to every request, which approximates remote storage. -//! -//! Run with `cargo bench -p datafusion --features parquet --bench parquet_bloom_filter_pruning`. - -use std::sync::Arc; -use std::time::Duration; - -use arrow::array::{ArrayRef, Int32Array, RecordBatch}; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use criterion::{BenchmarkId, Criterion, SamplingMode, criterion_group, criterion_main}; -use datafusion::datasource::file_format::parquet::ParquetFormat; -use datafusion::datasource::listing::ListingOptions; -use datafusion::prelude::{SessionConfig, SessionContext}; -use datafusion_execution::runtime_env::RuntimeEnv; -use object_store::memory::InMemory; -use object_store::path::Path; -use object_store::throttle::{ThrottleConfig, ThrottledStore}; -use object_store::{ObjectStore, ObjectStoreExt}; -use parquet::arrow::ArrowWriter; -use parquet::file::properties::{EnabledStatistics, WriterProperties}; -use tokio::runtime::Runtime; -use url::Url; - -const ROW_GROUPS: usize = 64; -const ROW_GROUP_ROWS: usize = 1024; -const MAX_COLUMNS: usize = 32; -const COLUMN_COUNTS: [usize; 3] = [2, 8, 32]; -const STORE_LATENCY: Duration = Duration::from_micros(200); -const PRESENT: i32 = 0; -const ABSENT: i32 = -1; - -fn schema() -> SchemaRef { - Arc::new(Schema::new( - (0..MAX_COLUMNS) - .map(|i| Field::new(format!("c{i}"), DataType::Int32, false)) - .collect::>(), - )) -} - -fn create_parquet_file() -> Vec { - let schema = schema(); - let props = WriterProperties::builder() - .set_max_row_group_row_count(Some(ROW_GROUP_ROWS)) - .set_statistics_enabled(EnabledStatistics::None) - .set_bloom_filter_enabled(true) - .set_bloom_filter_max_ndv(ROW_GROUP_ROWS as u64) - .set_bloom_filter_fpp(0.01) - .build(); - - let mut data = Vec::new(); - let mut writer = - ArrowWriter::try_new(&mut data, Arc::clone(&schema), Some(props)).unwrap(); - for rg in 0..ROW_GROUPS { - // Row 0 is PRESENT in every row group; other values are unique and positive. - let values: Int32Array = (0..ROW_GROUP_ROWS) - .map(|row| { - if row == 0 { - PRESENT - } else { - (rg * ROW_GROUP_ROWS + row) as i32 - } - }) - .collect(); - let column: ArrayRef = Arc::new(values); - let columns = vec![column; MAX_COLUMNS]; - let batch = RecordBatch::try_new(Arc::clone(&schema), columns).unwrap(); - writer.write(&batch).unwrap(); - } - writer.close().unwrap(); - data -} - -fn make_store_slow(object_store: InMemory) -> Arc { - let config = ThrottleConfig { - wait_get_per_call: STORE_LATENCY, - ..Default::default() - }; - Arc::new(ThrottledStore::new(object_store, config)) -} - -async fn setup_context() -> SessionContext { - let data = create_parquet_file(); - let stores: [(&str, Arc); 2] = [ - ("memory", Arc::new(InMemory::new())), - ("slow", make_store_slow(InMemory::new())), - ]; - - // One partition keeps the Bloom filter reads for the file in one opener. - let config = SessionConfig::new().with_target_partitions(1); - let rt = Arc::new(RuntimeEnv::default()); - for (name, store) in &stores { - store - .put(&Path::from("t/data.parquet"), data.clone().into()) - .await - .unwrap(); - rt.register_object_store( - &Url::parse(&format!("{name}://bucket")).unwrap(), - Arc::clone(store), - ); - } - let ctx = SessionContext::new_with_config_rt(config, rt); - - for (name, _) in &stores { - let options = ListingOptions::new(Arc::new(ParquetFormat::default())); - ctx.register_listing_table( - *name, - format!("{name}://bucket/t/"), - options, - Some(schema()), - None, - ) - .await - .unwrap(); - } - ctx -} - -/// Builds the query and its expected row count. -fn query(table: &str, case: &str, columns: usize) -> (String, usize) { - let predicate = (0..columns) - .map(|i| { - let rejecting = match case { - "reject_first" => i == 0, - "reject_last" => i == columns - 1, - "survive_all" => false, - _ => unreachable!(), - }; - let value = if rejecting { ABSENT } else { PRESENT }; - format!("c{i} = {value}") - }) - .collect::>() - .join(" AND "); - let expected_rows = if case == "survive_all" { ROW_GROUPS } else { 0 }; - ( - format!("SELECT c0 FROM {table} WHERE {predicate}"), - expected_rows, - ) -} - -async fn run(ctx: &SessionContext, sql: &str, expected_rows: usize) { - let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); - let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(rows, expected_rows); -} - -/// Prints the Bloom filter pruning metric once so a run shows whether Bloom -/// filters actually pruned. -async fn print_bloom_metrics(ctx: &SessionContext, sql: &str) { - let batches = ctx - .sql(&format!("EXPLAIN ANALYZE {sql}")) - .await - .unwrap() - .collect() - .await - .unwrap(); - let plan = arrow::util::pretty::pretty_format_batches(&batches) - .unwrap() - .to_string(); - let metric = plan - .split([',', '[']) - .find(|s| s.contains("row_groups_pruned_bloom_filter")) - .unwrap_or("row_groups_pruned_bloom_filter not reported"); - println!("{sql}\n {}", metric.trim()); -} - -fn criterion_benchmark(c: &mut Criterion) { - let rt = Runtime::new().unwrap(); - let ctx = rt.block_on(setup_context()); - - for case in ["reject_first", "reject_last", "survive_all"] { - for table in ["memory", "slow"] { - let mut group = - c.benchmark_group(format!("parquet_bloom_filter_pruning/{table}/{case}")); - group.sampling_mode(SamplingMode::Flat); - group.sample_size(10); - for columns in COLUMN_COUNTS { - let (sql, expected_rows) = query(table, case, columns); - rt.block_on(print_bloom_metrics(&ctx, &sql)); - group.bench_with_input( - BenchmarkId::new("columns", columns), - &sql, - |b, sql| { - b.to_async(&rt).iter(|| run(&ctx, sql, expected_rows)); - }, - ); - } - group.finish(); - } - } -} - -criterion_group!(benches, criterion_benchmark); -criterion_main!(benches); diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index 0928e87b775dd..13d1d64d5216b 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -34,8 +34,8 @@ use crate::push_decoder::{ }; use crate::row_group_filter::{RowGroupAccessPlanFilter, row_group_in_range}; use crate::{ - BloomFilterStatistics, Int96Coercer, ParquetAccessPlan, ParquetFileMetrics, - ParquetFileReaderFactory, ParquetRowSelection, ParquetVirtualColumn, RowGroupAccess, + Int96Coercer, ParquetAccessPlan, ParquetFileMetrics, ParquetFileReaderFactory, + ParquetRowSelection, ParquetVirtualColumn, RowGroupAccess, schema_coercion::apply_file_schema_type_coercions_with_rle, }; use arrow::array::RecordBatch; @@ -44,7 +44,7 @@ use datafusion_datasource::morsel::{Morsel, MorselPlan, MorselPlanner, Morselize use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr_adapter::replace_columns_with_literals; use datafusion_physical_expr_adapter::rewrite::rewrite_input_file_name_in_projection; -use std::collections::{HashMap, HashSet as StdHashSet, VecDeque}; +use std::collections::{HashMap, VecDeque}; use std::fmt; use std::future::Future; use std::mem; @@ -62,7 +62,7 @@ use datafusion_common::{ use datafusion_datasource::{PartitionedFile, TableSchema}; use datafusion_physical_expr::expressions::{Column, DynamicFilterTracking, Literal}; use datafusion_physical_expr::simplifier::PhysicalExprSimplifier; -use datafusion_physical_expr::utils::{Guarantee, collect_columns}; +use datafusion_physical_expr::utils::collect_columns; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::LexOrdering; @@ -84,7 +84,6 @@ use parquet::arrow::arrow_reader::{ }; use parquet::arrow::async_reader::AsyncFileReader; use parquet::arrow::parquet_column; -use parquet::basic::Type; use parquet::bloom_filter::Sbbf; use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader, RowGroupMetaData}; @@ -1468,14 +1467,6 @@ impl RowGroupsPrunedParquetOpen { .row_group_indexes() .any(|idx| !self.row_groups.access_plan().is_fully_matched(idx)) { - let bloom_filter_eval_time = self - .prepared - .loaded - .prepared - .file_metrics - .bloom_filter_eval_time - .clone(); - // Use the existing reader for bloom filter I/O; // replace with a fresh reader for decoding below. let reader_metadata = self.prepared.loaded.reader_metadata.clone(); @@ -1494,58 +1485,23 @@ impl RowGroupsPrunedParquetOpen { mem::replace(&mut prepared.async_file_reader, replacement_reader), reader_metadata, ); - struct BloomColumnGuarantee { - column_idx: usize, - physical_type: Type, - type_length: i32, - in_guarantees: Vec>, - } - let parquet_columns: Vec = predicate - .literal_columns() - .into_iter() - .filter_map(|column_name| { - let parquet_schema = builder.parquet_schema(); - let (column_idx, _) = parquet_column( - parquet_schema, - &prepared.physical_file_schema, - &column_name, - )?; - let in_guarantees: Vec> = predicate - .literal_guarantees() - .iter() - .filter(|lg| { - lg.column.name == column_name && lg.guarantee == Guarantee::In - }) - .map(|lg| lg.literals.clone()) - .collect(); - Some(BloomColumnGuarantee { - column_idx, - physical_type: parquet_schema.column(column_idx).physical_type(), - type_length: parquet_schema.column(column_idx).type_length(), - in_guarantees, - }) - }) - .collect(); + let parquet_columns = + RowGroupAccessPlanFilter::prepare_bloom_column_guarantees( + predicate, + builder.parquet_schema(), + &prepared.physical_file_schema, + ); - // Collect up front so the iterator's immutable borrow on - // `self.row_groups` is released before we mutably call - // `access_plan_mut().skip()` inside the loop body. let row_group_indexes: Vec = self.row_groups.row_group_indexes().collect(); for idx in row_group_indexes { - if !self.row_groups.access_plan().should_scan(idx) { - continue; - } - // Statistics have already proved every row in this group matches - // the predicate, so Bloom filters cannot prune it — skip early. - if self.row_groups.access_plan().is_fully_matched(idx) { + if !self.row_groups.access_plan().should_scan(idx) + || self.row_groups.access_plan().is_fully_matched(idx) + { continue; } - // Track whether at least one Bloom filter was loaded for this row - // group so we can correctly gate the matched/pruned metric updates. let mut any_filter_loaded = false; - let mut prune_group = false; for col_guarantee in &parquet_columns { let bf: Sbbf = match builder @@ -1562,49 +1518,19 @@ impl RowGroupsPrunedParquetOpen { }; any_filter_loaded = true; - // Evaluate only the literal guarantees for this newly loaded column. - // - // Calling `predicate.prune()` after each column would repeatedly - // evaluate the full expression tree and rebuild statistics batches, - // incurring heavy O(N) overhead per column (O(N²) per row group). - // Because Bloom filters can only prune row groups via literal - // `Guarantee::In` guarantees (proving none of the required values - // exist in the filter), checking only the guarantees for this newly - // loaded column preserves early rejection while completely eliminating - // redundant evaluation overhead. - let _timer_guard = bloom_filter_eval_time.timer(); - let prunable = col_guarantee.in_guarantees.iter().any(|literals| { - BloomFilterStatistics::is_known_not_present( - &bf, - literals, - &col_guarantee.physical_type, - col_guarantee.type_length, - ) - }); - - if prunable { - prune_group = true; + if self.row_groups.prune_row_group_with_bloom_filter( + idx, + col_guarantee, + &bf, + &prepared.file_metrics, + ) { break; } } - if !any_filter_loaded { - // No Bloom filters were available for this row group; evaluation was - // not performed, so do not update matched or pruned metrics. - continue; - } - - if prune_group { - prepared - .file_metrics - .row_groups_pruned_bloom_filter - .add_pruned(1); - self.row_groups.access_plan_mut().skip(idx); - } else { - prepared - .file_metrics - .row_groups_pruned_bloom_filter - .add_matched(1); + if any_filter_loaded && self.row_groups.access_plan().should_scan(idx) { + self.row_groups + .record_bloom_filter_matched(&prepared.file_metrics); } } } @@ -4884,6 +4810,151 @@ mod test { assert_eq!(matched, 0); } + #[derive(Clone, Default, Debug)] + struct RangeTracker { + ranges: Arc>>>, + } + + struct TrackingBloomAsyncReader { + inner: Box, + tracker: RangeTracker, + } + + impl AsyncFileReader for TrackingBloomAsyncReader { + fn get_bytes( + &mut self, + range: std::ops::Range, + ) -> BoxFuture<'_, parquet::errors::Result> { + self.tracker.ranges.lock().unwrap().push(range.clone()); + self.inner.get_bytes(range) + } + + fn get_byte_ranges( + &mut self, + ranges: Vec>, + ) -> BoxFuture<'_, parquet::errors::Result>> { + self.tracker.ranges.lock().unwrap().extend(ranges.clone()); + self.inner.get_byte_ranges(ranges) + } + + fn get_metadata<'a>( + &'a mut self, + options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, parquet::errors::Result>> { + self.inner.get_metadata(options) + } + } + + #[derive(Debug)] + struct TrackingBloomReaderFactory { + base_factory: Arc, + tracker: RangeTracker, + } + + impl ParquetFileReaderFactory for TrackingBloomReaderFactory { + fn create_reader( + &self, + partition_index: usize, + file: PartitionedFile, + metadata_size_hint: Option, + metrics: &ExecutionPlanMetricsSet, + ) -> Result> { + let inner = self.base_factory.create_reader( + partition_index, + file, + metadata_size_hint, + metrics, + )?; + Ok(Box::new(TrackingBloomAsyncReader { + inner, + tracker: self.tracker.clone(), + })) + } + } + + #[tokio::test] + async fn test_bloom_filter_incremental_io_pattern() { + let store = Arc::new(InMemory::new()) as Arc; + let batch = record_batch!( + ("a", Int32, vec![10, 20, 30]), + ("b", Int32, vec![100, 200, 300]) + ) + .unwrap(); + let schema = batch.schema(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .set_bloom_filter_enabled(true) + .set_statistics_enabled(EnabledStatistics::None) + .build(); + let data_size = write_parquet_batches( + Arc::clone(&store), + "io_bloom.parquet", + vec![batch], + Some(props), + ) + .await; + let file = PartitionedFile::new("io_bloom.parquet".to_string(), data_size as u64); + + // Case 1: Predicate a = 10 AND b = 100 -> 'a' matches, so 'b' MUST also be loaded. + let tracker1 = RangeTracker::default(); + let factory1 = Arc::new(TrackingBloomReaderFactory { + base_factory: Arc::new(DefaultParquetFileReaderFactory::new(Arc::clone( + &store, + ))), + tracker: tracker1.clone(), + }); + let pred1 = + logical2physical(&col("a").eq(lit(10)).and(col("b").eq(lit(100))), &schema); + let morselizer1 = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(Arc::clone(&schema)) + .with_predicate(pred1) + .with_pushdown_filters(true) + .with_row_group_stats_pruning(false) + .with_enable_bloom_filter(true) + .with_parquet_file_reader_factory(factory1) + .build(); + let stream1 = open_file(&morselizer1, file.clone()).await.unwrap(); + let results1 = collect_int32_values(stream1).await; + assert_eq!(results1, vec![10]); + let requests1 = tracker1.ranges.lock().unwrap().clone(); + + // Case 2: Predicate a = 999 AND b = 100 -> 'a' is absent (pruned!), + // so 'b' Bloom filter MUST NOT be requested. + let tracker2 = RangeTracker::default(); + let factory2 = Arc::new(TrackingBloomReaderFactory { + base_factory: Arc::new(DefaultParquetFileReaderFactory::new(Arc::clone( + &store, + ))), + tracker: tracker2.clone(), + }); + let pred2 = + logical2physical(&col("a").eq(lit(999)).and(col("b").eq(lit(100))), &schema); + let morselizer2 = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(Arc::clone(&schema)) + .with_predicate(pred2) + .with_pushdown_filters(true) + .with_row_group_stats_pruning(false) + .with_enable_bloom_filter(true) + .with_parquet_file_reader_factory(factory2) + .build(); + let stream2 = open_file(&morselizer2, file.clone()).await.unwrap(); + let results2 = collect_int32_values(stream2).await; + assert_eq!(results2, Vec::::new()); + let requests2 = tracker2.ranges.lock().unwrap().clone(); + + // In Case 1, both column 'a' and column 'b' bloom filters were requested. + // In Case 2, column 'a' pruned the row group early, so column 'b' was never requested. + // Therefore, Case 2 makes strictly fewer range requests than Case 1. + assert!( + requests2.len() < requests1.len(), + "Incremental pruning should issue fewer I/O requests when pruned early: case 2 ({}) < case 1 ({})", + requests2.len(), + requests1.len() + ); + } + #[tokio::test] async fn incremental_bloom_filter_pruning_short_circuits_remaining_columns() { let store = Arc::new(InMemory::new()) as Arc; diff --git a/datafusion/datasource-parquet/src/row_group_filter.rs b/datafusion/datasource-parquet/src/row_group_filter.rs index c73b22afd4589..1884fca269cad 100644 --- a/datafusion/datasource-parquet/src/row_group_filter.rs +++ b/datafusion/datasource-parquet/src/row_group_filter.rs @@ -29,9 +29,12 @@ use datafusion_common::parquet_config::RowGroupRangeAssignment; use datafusion_common::pruning::PruningStatistics; use datafusion_common::{Column, Result, ScalarValue}; use datafusion_datasource::FileRange; +use datafusion_physical_expr::utils::Guarantee; use datafusion_pruning::PruningPredicate; use parquet::arrow::arrow_reader::statistics::StatisticsConverter; use parquet::basic::ColumnOrder; +use parquet::basic::Type; +use parquet::bloom_filter::Sbbf; use parquet::file::metadata::{ParquetMetaData, RowGroupMetaData}; use parquet::schema::types::SchemaDescriptor; @@ -125,11 +128,6 @@ impl RowGroupAccessPlanFilter { &self.access_plan } - /// Returns a mutable reference to the inner access plan. - pub(crate) fn access_plan_mut(&mut self) -> &mut ParquetAccessPlan { - &mut self.access_plan - } - /// Returns the is_fully_matched vector. pub fn is_fully_matched(&self) -> &Vec { self.access_plan.fully_matched() @@ -511,12 +509,83 @@ impl RowGroupAccessPlanFilter { if prune_group { metrics.row_groups_pruned_bloom_filter.add_pruned(1); - self.access_plan.skip(idx) + self.access_plan.skip(idx); } else { metrics.row_groups_pruned_bloom_filter.add_matched(1); } } } + + /// Extracts column guarantees from a [`PruningPredicate`] that can be evaluated against + /// Parquet Bloom filters. + pub fn prepare_bloom_column_guarantees( + predicate: &PruningPredicate, + parquet_schema: &SchemaDescriptor, + file_schema: &Schema, + ) -> Vec { + predicate + .literal_columns() + .into_iter() + .filter_map(|column_name| { + let (column_idx, _) = parquet::arrow::parquet_column( + parquet_schema, + file_schema, + &column_name, + )?; + let in_guarantees: Vec> = predicate + .literal_guarantees() + .iter() + .filter(|lg| { + lg.column.name == column_name && lg.guarantee == Guarantee::In + }) + .map(|lg| lg.literals.clone()) + .collect(); + Some(BloomColumnGuarantee { + column_idx, + physical_type: parquet_schema.column(column_idx).physical_type(), + type_length: parquet_schema.column(column_idx).type_length(), + in_guarantees, + }) + }) + .collect() + } + + /// Evaluates whether a loaded column Bloom filter prunes the specified row group. + /// + /// If all values for any IN guarantee on this column are confirmed absent from the Bloom + /// filter, the row group is marked as skipped in the access plan, the + /// `row_groups_pruned_bloom_filter` metric is incremented, and `true` is returned to + /// indicate early rejection (allowing callers to skip loading remaining column filters). + pub fn prune_row_group_with_bloom_filter( + &mut self, + idx: usize, + guarantee: &BloomColumnGuarantee, + sbbf: &Sbbf, + metrics: &ParquetFileMetrics, + ) -> bool { + let _timer_guard = metrics.bloom_filter_eval_time.timer(); + let prunable = guarantee.in_guarantees.iter().any(|literals| { + BloomFilterStatistics::is_known_not_present( + sbbf, + literals, + &guarantee.physical_type, + guarantee.type_length, + ) + }); + + if prunable { + metrics.row_groups_pruned_bloom_filter.add_pruned(1); + self.access_plan.skip(idx); + true + } else { + false + } + } + + /// Records that the row group was evaluated against Bloom filters but was not pruned. + pub fn record_bloom_filter_matched(&mut self, metrics: &ParquetFileMetrics) { + metrics.row_groups_pruned_bloom_filter.add_matched(1); + } } /// Wraps a slice of [`RowGroupMetaData`] in a way that implements [`PruningStatistics`]. @@ -1781,3 +1850,19 @@ mod tests { expected.assert(&row_groups); } } + +/// Pre-extracted Bloom filter pruning guarantees for a single Parquet column. +/// +/// Pre-extracting guarantees avoids repeatedly re-traversing the [`PruningPredicate`] +/// expression tree for every row group and column during incremental Bloom filter loading. +#[derive(Debug, Clone)] +pub struct BloomColumnGuarantee { + /// Parquet physical column index + pub column_idx: usize, + /// Physical Parquet type of the column + pub physical_type: Type, + /// Type length for FIXED_LEN_BYTE_ARRAY, or 0 + pub type_length: i32, + /// Literal sets for IN guarantees on this column + pub in_guarantees: Vec>, +}