diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 89f42761535e3..60bed6087614c 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -245,10 +245,6 @@ name = "math_query_sql" harness = false name = "filter_query_sql" -[[bench]] -harness = false -name = "in_list_rewrite" - [[bench]] harness = false name = "struct_query_sql" @@ -274,6 +270,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); diff --git a/datafusion/datasource-parquet/src/bloom_filter.rs b/datafusion/datasource-parquet/src/bloom_filter.rs index 2384f95df9885..a10932b102210 100644 --- a/datafusion/datasource-parquet/src/bloom_filter.rs +++ b/datafusion/datasource-parquet/src/bloom_filter.rs @@ -85,6 +85,24 @@ 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`. + 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 +227,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 +730,50 @@ mod tests { Ok(pruned_row_groups) } + + #[test] + 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()