perf: evaluate Parquet Bloom filters incrementally during row group pruning

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"

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::<Vec<_>>(),
+    ))
+}
+
+fn create_parquet_file() -> Vec<u8> {
+    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<dyn ObjectStore> {
+    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<dyn ObjectStore>); 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::<Vec<_>>()
+        .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);

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<ScalarValue>,
+        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,
+        ));
+    }
 }