From c0a70fb733f81d0165b5efb05a72f98e451eb212 Mon Sep 17 00:00:00 2001 From: mohitgurav20 Date: Fri, 30 Oct 2026 00:20:25 +0530 Subject: [PATCH 1/7] test(parquet): add SQL-level I/O snapshot test for Bloom filter loading baseline - Add InstrumentedObjectStore that correlates requested byte ranges to Parquet file structures (footers, page indexes, bloom filters, data pages) - Add test_parquet_bloom_filter_io_baseline using public SessionContext SQL APIs - Use insta snapshot to capture and compare the storage I/O trace - Documents baseline on main where all candidate column Bloom filters are loaded upfront, establishing the target diff for incremental pruning (#25811) --- .../core/tests/parquet/bloom_filter_io.rs | 357 ++++++++++++++++++ datafusion/core/tests/parquet/mod.rs | 1 + ..._io__parquet_bloom_filter_io_baseline.snap | 17 + 3 files changed, 375 insertions(+) create mode 100644 datafusion/core/tests/parquet/bloom_filter_io.rs create mode 100644 datafusion/core/tests/parquet/snapshots/parquet_integration__parquet__bloom_filter_io__parquet_bloom_filter_io_baseline.snap diff --git a/datafusion/core/tests/parquet/bloom_filter_io.rs b/datafusion/core/tests/parquet/bloom_filter_io.rs new file mode 100644 index 0000000000000..d50a7543e9b83 --- /dev/null +++ b/datafusion/core/tests/parquet/bloom_filter_io.rs @@ -0,0 +1,357 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! Integration tests tracking storage I/O events during Parquet Bloom filter evaluation. +//! +//! Uses an instrumented [`ObjectStore`] to record storage requests (footers, page indexes, +//! bloom filters, data pages) and compares the resulting trace against an `insta` snapshot. + +use std::fmt::Write as _; +use std::ops::Range; +use std::sync::{Arc, Mutex}; + +use arrow::array::{ArrayRef, Int32Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema}; +use async_trait::async_trait; +use bytes::Bytes; +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use futures::StreamExt; +use futures::stream::BoxStream; +use insta::assert_snapshot; +use object_store::memory::InMemory; +use object_store::path::Path; +use object_store::{ + GetOptions, GetRange, GetResult, ListResult, MultipartUpload, ObjectMeta, + ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, +}; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::{EnabledStatistics, WriterProperties}; +use parquet::file::reader::{FileReader, SerializedFileReader}; +use url::Url; + +#[derive(Debug, Clone)] +struct ByteRangeDescriptor { + range: Range, + description: String, +} + +#[derive(Debug, Clone, Default)] +struct ParquetFileLayout { + descriptors: Vec, +} + +impl ParquetFileLayout { + fn new(bytes: &Bytes) -> Self { + let cursor = bytes.clone(); + let reader = SerializedFileReader::new(cursor).expect("valid parquet file"); + let metadata = reader.metadata(); + let mut descriptors = Vec::new(); + + for (rg_idx, rg) in metadata.row_groups().iter().enumerate() { + for col in rg.columns() { + let col_name = col.column_descr().name(); + + // Bloom filter byte range + if let (Some(bf_offset), Some(bf_len)) = + (col.bloom_filter_offset(), col.bloom_filter_length()) + { + let start = bf_offset as u64; + let end = start + bf_len as u64; + descriptors.push(ByteRangeDescriptor { + range: start..end, + description: format!( + "Bloom filter for column '{col_name}' (rg {rg_idx})" + ), + }); + } + + // Page index byte ranges + if let Some(r) = col.column_index_range() { + descriptors.push(ByteRangeDescriptor { + range: r, + description: format!( + "ColumnIndex for column '{col_name}' (rg {rg_idx})" + ), + }); + } + if let Some(r) = col.offset_index_range() { + descriptors.push(ByteRangeDescriptor { + range: r, + description: format!( + "OffsetIndex for column '{col_name}' (rg {rg_idx})" + ), + }); + } + + // Data pages / column chunk byte range + let data_start = col + .dictionary_page_offset() + .and_then(|o| u64::try_from(o).ok()) + .unwrap_or_else(|| col.data_page_offset() as u64); + let data_end = data_start + col.compressed_size() as u64; + descriptors.push(ByteRangeDescriptor { + range: data_start..data_end, + description: format!( + "data page for column '{col_name}' (rg {rg_idx})" + ), + }); + } + } + + Self { descriptors } + } + + fn describe_range(&self, range: &Range, file_size: u64) -> String { + if range.end >= file_size && (file_size.saturating_sub(range.start)) <= 8192 { + return format!("Parquet footer (offsets {}..{})", range.start, range.end); + } + for desc in &self.descriptors { + if range.start < desc.range.end && range.end > desc.range.start { + return format!( + "{} (offsets {}..{})", + desc.description, range.start, range.end + ); + } + } + format!("byte range (offsets {}..{})", range.start, range.end) + } +} + +#[derive(Debug)] +struct InstrumentedObjectStore { + inner: Arc, + layout: ParquetFileLayout, + file_size: u64, + events: Arc>>, +} + +impl InstrumentedObjectStore { + fn new( + inner: Arc, + layout: ParquetFileLayout, + file_size: u64, + events: Arc>>, + ) -> Self { + Self { + inner, + layout, + file_size, + events, + } + } + + fn record_range(&self, range: &Range) { + let desc = self.layout.describe_range(range, self.file_size); + self.events.lock().unwrap().push(format!("Read {desc}")); + } +} + +impl std::fmt::Display for InstrumentedObjectStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "InstrumentedObjectStore") + } +} + +#[async_trait] +impl ObjectStore for InstrumentedObjectStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> object_store::Result { + self.inner.put_opts(location, payload, opts).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> object_store::Result> { + self.inner.put_multipart_opts(location, opts).await + } + + async fn get_opts( + &self, + location: &Path, + options: GetOptions, + ) -> object_store::Result { + if let Some(range) = &options.range { + match range { + GetRange::Bounded(r) => { + self.record_range(&(r.start..r.end)); + } + GetRange::Offset(offset) => { + self.record_range(&(*offset..self.file_size)); + } + GetRange::Suffix(suffix) => { + let start = self.file_size.saturating_sub(*suffix); + self.record_range(&(start..self.file_size)); + } + } + } else { + self.record_range(&(0..self.file_size)); + } + self.inner.get_opts(location, options).await + } + + async fn get_ranges( + &self, + location: &Path, + ranges: &[Range], + ) -> object_store::Result> { + for range in ranges { + self.record_range(range); + } + self.inner.get_ranges(location, ranges).await + } + + fn delete_stream( + &self, + locations: BoxStream<'static, object_store::Result>, + ) -> BoxStream<'static, object_store::Result> { + self.inner.delete_stream(locations) + } + + fn list( + &self, + prefix: Option<&Path>, + ) -> BoxStream<'static, object_store::Result> { + self.inner.list(prefix) + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, object_store::Result> { + self.inner.list_with_offset(prefix, offset) + } + + async fn list_with_delimiter( + &self, + prefix: Option<&Path>, + ) -> object_store::Result { + self.inner.list_with_delimiter(prefix).await + } + + async fn copy_opts( + &self, + from: &Path, + to: &Path, + options: object_store::CopyOptions, + ) -> object_store::Result<()> { + self.inner.copy_opts(from, to, options).await + } +} + +async fn run_query_and_collect_events( + ctx: &SessionContext, + sql: &str, + events: &Arc>>, +) -> String { + events.lock().unwrap().clear(); + + let df = ctx.sql(sql).await.unwrap(); + let mut stream = df.execute_stream().await.unwrap(); + + let mut row_count = 0; + while let Some(batch) = stream.next().await { + let batch = batch.unwrap(); + row_count += batch.num_rows(); + events + .lock() + .unwrap() + .push(format!("Produce batch ({} rows)", batch.num_rows())); + } + + let recorded = events.lock().unwrap().clone(); + let mut output = format!("--- Query: {sql} ({row_count} rows) ---\n"); + for event in recorded { + let _ = writeln!(output, "{event}"); + } + output +} + +#[tokio::test] +async fn test_parquet_bloom_filter_io_baseline() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef, + Arc::new(Int32Array::from(vec![100, 200, 300])) as ArrayRef, + ], + ) + .unwrap(); + + // Disable row group and page statistics so pruning relies entirely on Bloom filters + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .set_bloom_filter_enabled(true) + .set_statistics_enabled(EnabledStatistics::None) + .build(); + + let mut buffer = Vec::new(); + let mut writer = + ArrowWriter::try_new(&mut buffer, batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let file_bytes = Bytes::from(buffer); + let file_size = file_bytes.len() as u64; + let layout = ParquetFileLayout::new(&file_bytes); + + let memory_store = Arc::new(InMemory::new()); + let path = Path::from("table.parquet"); + memory_store + .put_opts(&path, PutPayload::from(file_bytes), PutOptions::default()) + .await + .unwrap(); + + let events = Arc::new(Mutex::new(Vec::new())); + let instrumented_store = Arc::new(InstrumentedObjectStore::new( + memory_store, + layout, + file_size, + events.clone(), + )); + + let config = SessionConfig::new() + .with_target_partitions(1) + .with_parquet_bloom_filter_pruning(true); + + let ctx = SessionContext::new_with_config(config); + let table_url = Url::parse("mem://").unwrap(); + ctx.register_object_store(&table_url, instrumented_store); + + ctx.register_parquet("t", "mem:///table.parquet", ParquetReadOptions::default()) + .await + .unwrap(); + + // Case 1: Matching row group (a = 10 AND b = 100) + // Loads Bloom filters, verifies match, then reads data pages and produces batch + let trace_matching = run_query_and_collect_events( + &ctx, + "SELECT * FROM t WHERE a = 10 AND b = 100", + &events, + ) + .await; + + // Case 2: Column 'a' does not match (a = 999 AND b = 100) + // On main, all candidate column Bloom filters ('a' and 'b') are loaded upfront. + // In #25811, incremental pruning will evaluate 'a' first, short-circuit, and skip loading 'b'. + let trace_pruned = run_query_and_collect_events( + &ctx, + "SELECT * FROM t WHERE a = 999 AND b = 100", + &events, + ) + .await; + + let combined_trace = format!("{trace_matching}\n{trace_pruned}"); + assert_snapshot!(combined_trace); +} diff --git a/datafusion/core/tests/parquet/mod.rs b/datafusion/core/tests/parquet/mod.rs index be7cce00ddab4..8fe7184ce6d54 100644 --- a/datafusion/core/tests/parquet/mod.rs +++ b/datafusion/core/tests/parquet/mod.rs @@ -31,6 +31,7 @@ use parquet::file::properties::{EnabledStatistics, WriterProperties}; use std::sync::Arc; use tempfile::NamedTempFile; +mod bloom_filter_io; mod content_defined_chunking; mod custom_reader; mod dynamic_row_group_pruning; diff --git a/datafusion/core/tests/parquet/snapshots/parquet_integration__parquet__bloom_filter_io__parquet_bloom_filter_io_baseline.snap b/datafusion/core/tests/parquet/snapshots/parquet_integration__parquet__bloom_filter_io__parquet_bloom_filter_io_baseline.snap new file mode 100644 index 0000000000000..ae78d000b305a --- /dev/null +++ b/datafusion/core/tests/parquet/snapshots/parquet_integration__parquet__bloom_filter_io__parquet_bloom_filter_io_baseline.snap @@ -0,0 +1,17 @@ +--- +source: datafusion/core/tests/parquet/bloom_filter_io.rs +expression: combined_trace +--- +--- Query: SELECT * FROM t WHERE a = 10 AND b = 100 (1 rows) --- +Read Parquet footer (offsets 0..678) +Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) +Read Bloom filter for column 'a' (rg 0) (offsets 98..145) +Read Bloom filter for column 'b' (rg 0) (offsets 145..192) +Read data page for column 'a' (rg 0) (offsets 4..51) +Read data page for column 'b' (rg 0) (offsets 51..98) +Produce batch (1 rows) + +--- Query: SELECT * FROM t WHERE a = 999 AND b = 100 (0 rows) --- +Read Parquet footer (offsets 0..678) +Read Bloom filter for column 'a' (rg 0) (offsets 98..145) +Read Bloom filter for column 'b' (rg 0) (offsets 145..192) From 73ba64e20e15c8bffd78d399fb00a6f5f5e80734 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 2 Nov 2026 15:37:23 -0500 Subject: [PATCH 2/7] test(parquet): use inline snapshot for Bloom filter I/O test Co-Authored-By: Claude Opus 5.5 --- .../core/tests/parquet/bloom_filter_io.rs | 16 +++++++++++++++- ...er_io__parquet_bloom_filter_io_baseline.snap | 17 ----------------- 2 files changed, 15 insertions(+), 18 deletions(-) delete mode 100644 datafusion/core/tests/parquet/snapshots/parquet_integration__parquet__bloom_filter_io__parquet_bloom_filter_io_baseline.snap diff --git a/datafusion/core/tests/parquet/bloom_filter_io.rs b/datafusion/core/tests/parquet/bloom_filter_io.rs index d50a7543e9b83..ce90ff6344896 100644 --- a/datafusion/core/tests/parquet/bloom_filter_io.rs +++ b/datafusion/core/tests/parquet/bloom_filter_io.rs @@ -353,5 +353,19 @@ async fn test_parquet_bloom_filter_io_baseline() { .await; let combined_trace = format!("{trace_matching}\n{trace_pruned}"); - assert_snapshot!(combined_trace); + assert_snapshot!(combined_trace, @r" + --- Query: SELECT * FROM t WHERE a = 10 AND b = 100 (1 rows) --- + Read Parquet footer (offsets 0..678) + Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) + Read Bloom filter for column 'a' (rg 0) (offsets 98..145) + Read Bloom filter for column 'b' (rg 0) (offsets 145..192) + Read data page for column 'a' (rg 0) (offsets 4..51) + Read data page for column 'b' (rg 0) (offsets 51..98) + Produce batch (1 rows) + + --- Query: SELECT * FROM t WHERE a = 999 AND b = 100 (0 rows) --- + Read Parquet footer (offsets 0..678) + Read Bloom filter for column 'a' (rg 0) (offsets 98..145) + Read Bloom filter for column 'b' (rg 0) (offsets 145..192) + "); } diff --git a/datafusion/core/tests/parquet/snapshots/parquet_integration__parquet__bloom_filter_io__parquet_bloom_filter_io_baseline.snap b/datafusion/core/tests/parquet/snapshots/parquet_integration__parquet__bloom_filter_io__parquet_bloom_filter_io_baseline.snap deleted file mode 100644 index ae78d000b305a..0000000000000 --- a/datafusion/core/tests/parquet/snapshots/parquet_integration__parquet__bloom_filter_io__parquet_bloom_filter_io_baseline.snap +++ /dev/null @@ -1,17 +0,0 @@ ---- -source: datafusion/core/tests/parquet/bloom_filter_io.rs -expression: combined_trace ---- ---- Query: SELECT * FROM t WHERE a = 10 AND b = 100 (1 rows) --- -Read Parquet footer (offsets 0..678) -Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) -Read Bloom filter for column 'a' (rg 0) (offsets 98..145) -Read Bloom filter for column 'b' (rg 0) (offsets 145..192) -Read data page for column 'a' (rg 0) (offsets 4..51) -Read data page for column 'b' (rg 0) (offsets 51..98) -Produce batch (1 rows) - ---- Query: SELECT * FROM t WHERE a = 999 AND b = 100 (0 rows) --- -Read Parquet footer (offsets 0..678) -Read Bloom filter for column 'a' (rg 0) (offsets 98..145) -Read Bloom filter for column 'b' (rg 0) (offsets 145..192) From 37110ee091411ca5a6ad37eff7aabda903b93677 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 2 Nov 2026 15:38:45 -0500 Subject: [PATCH 3/7] test(parquet): move I/O tracing setup into parquet/io module Split the Bloom filter I/O test into reusable setup code in tests/parquet/io/mod.rs and the test itself in tests/parquet/io/bloom_filter_io.rs. Co-Authored-By: Claude Opus 5.5 --- .../core/tests/parquet/io/bloom_filter_io.rs | 117 ++++++++++++++++ .../parquet/{bloom_filter_io.rs => io/mod.rs} | 127 +++--------------- datafusion/core/tests/parquet/mod.rs | 2 +- 3 files changed, 133 insertions(+), 113 deletions(-) create mode 100644 datafusion/core/tests/parquet/io/bloom_filter_io.rs rename datafusion/core/tests/parquet/{bloom_filter_io.rs => io/mod.rs} (65%) diff --git a/datafusion/core/tests/parquet/io/bloom_filter_io.rs b/datafusion/core/tests/parquet/io/bloom_filter_io.rs new file mode 100644 index 0000000000000..703b45932c64f --- /dev/null +++ b/datafusion/core/tests/parquet/io/bloom_filter_io.rs @@ -0,0 +1,117 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! Tests tracking storage I/O events during Parquet Bloom filter evaluation. + +use std::sync::{Arc, Mutex}; + +use arrow::array::{ArrayRef, Int32Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema}; +use bytes::Bytes; +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use insta::assert_snapshot; +use object_store::memory::InMemory; +use object_store::path::Path; +use object_store::{ObjectStore, PutOptions, PutPayload}; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::{EnabledStatistics, WriterProperties}; +use url::Url; + +use super::{InstrumentedObjectStore, ParquetFileLayout, run_query_and_collect_events}; + +#[tokio::test] +async fn test_parquet_bloom_filter_io_baseline() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef, + Arc::new(Int32Array::from(vec![100, 200, 300])) as ArrayRef, + ], + ) + .unwrap(); + + // Disable row group and page statistics so pruning relies entirely on Bloom filters + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .set_bloom_filter_enabled(true) + .set_statistics_enabled(EnabledStatistics::None) + .build(); + + let mut buffer = Vec::new(); + let mut writer = + ArrowWriter::try_new(&mut buffer, batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let file_bytes = Bytes::from(buffer); + let file_size = file_bytes.len() as u64; + let layout = ParquetFileLayout::new(&file_bytes); + + let memory_store = Arc::new(InMemory::new()); + let path = Path::from("table.parquet"); + memory_store + .put_opts(&path, PutPayload::from(file_bytes), PutOptions::default()) + .await + .unwrap(); + + let events = Arc::new(Mutex::new(Vec::new())); + let instrumented_store = Arc::new(InstrumentedObjectStore::new( + memory_store, + layout, + file_size, + events.clone(), + )); + + let config = SessionConfig::new() + .with_target_partitions(1) + .with_parquet_bloom_filter_pruning(true); + + let ctx = SessionContext::new_with_config(config); + let table_url = Url::parse("mem://").unwrap(); + ctx.register_object_store(&table_url, instrumented_store); + + ctx.register_parquet("t", "mem:///table.parquet", ParquetReadOptions::default()) + .await + .unwrap(); + + // Case 1: Matching row group (a = 10 AND b = 100) + // Loads Bloom filters, verifies match, then reads data pages and produces batch + let trace_matching = run_query_and_collect_events( + &ctx, + "SELECT * FROM t WHERE a = 10 AND b = 100", + &events, + ) + .await; + + // Case 2: Column 'a' does not match (a = 999 AND b = 100) + // On main, all candidate column Bloom filters ('a' and 'b') are loaded upfront. + // In #25811, incremental pruning will evaluate 'a' first, short-circuit, and skip loading 'b'. + let trace_pruned = run_query_and_collect_events( + &ctx, + "SELECT * FROM t WHERE a = 999 AND b = 100", + &events, + ) + .await; + + let combined_trace = format!("{trace_matching}\n{trace_pruned}"); + assert_snapshot!(combined_trace, @r" + --- Query: SELECT * FROM t WHERE a = 10 AND b = 100 (1 rows) --- + Read Parquet footer (offsets 0..678) + Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) + Read Bloom filter for column 'a' (rg 0) (offsets 98..145) + Read Bloom filter for column 'b' (rg 0) (offsets 145..192) + Read data page for column 'a' (rg 0) (offsets 4..51) + Read data page for column 'b' (rg 0) (offsets 51..98) + Produce batch (1 rows) + + --- Query: SELECT * FROM t WHERE a = 999 AND b = 100 (0 rows) --- + Read Parquet footer (offsets 0..678) + Read Bloom filter for column 'a' (rg 0) (offsets 98..145) + Read Bloom filter for column 'b' (rg 0) (offsets 145..192) + "); +} diff --git a/datafusion/core/tests/parquet/bloom_filter_io.rs b/datafusion/core/tests/parquet/io/mod.rs similarity index 65% rename from datafusion/core/tests/parquet/bloom_filter_io.rs rename to datafusion/core/tests/parquet/io/mod.rs index ce90ff6344896..1617f2716804c 100644 --- a/datafusion/core/tests/parquet/bloom_filter_io.rs +++ b/datafusion/core/tests/parquet/io/mod.rs @@ -2,33 +2,32 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -//! Integration tests tracking storage I/O events during Parquet Bloom filter evaluation. +//! Test infrastructure for tracking storage I/O events during Parquet scans. //! -//! Uses an instrumented [`ObjectStore`] to record storage requests (footers, page indexes, -//! bloom filters, data pages) and compares the resulting trace against an `insta` snapshot. +//! Provides an instrumented [`ObjectStore`] that records storage requests and +//! correlates each requested byte range with the Parquet file structure it +//! covers (footer, page index, Bloom filter, data pages), so the resulting +//! trace can be compared against `insta` snapshots. +//! +//! The actual tests live in the submodules. use std::fmt::Write as _; use std::ops::Range; use std::sync::{Arc, Mutex}; -use arrow::array::{ArrayRef, Int32Array, RecordBatch}; -use arrow::datatypes::{DataType, Field, Schema}; use async_trait::async_trait; use bytes::Bytes; -use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use datafusion::prelude::SessionContext; use futures::StreamExt; use futures::stream::BoxStream; -use insta::assert_snapshot; -use object_store::memory::InMemory; use object_store::path::Path; use object_store::{ GetOptions, GetRange, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, }; -use parquet::arrow::ArrowWriter; -use parquet::file::properties::{EnabledStatistics, WriterProperties}; use parquet::file::reader::{FileReader, SerializedFileReader}; -use url::Url; + +mod bloom_filter_io; #[derive(Debug, Clone)] struct ByteRangeDescriptor { @@ -37,12 +36,12 @@ struct ByteRangeDescriptor { } #[derive(Debug, Clone, Default)] -struct ParquetFileLayout { +pub(crate) struct ParquetFileLayout { descriptors: Vec, } impl ParquetFileLayout { - fn new(bytes: &Bytes) -> Self { + pub(crate) fn new(bytes: &Bytes) -> Self { let cursor = bytes.clone(); let reader = SerializedFileReader::new(cursor).expect("valid parquet file"); let metadata = reader.metadata(); @@ -119,7 +118,7 @@ impl ParquetFileLayout { } #[derive(Debug)] -struct InstrumentedObjectStore { +pub(crate) struct InstrumentedObjectStore { inner: Arc, layout: ParquetFileLayout, file_size: u64, @@ -127,7 +126,7 @@ struct InstrumentedObjectStore { } impl InstrumentedObjectStore { - fn new( + pub(crate) fn new( inner: Arc, layout: ParquetFileLayout, file_size: u64, @@ -246,7 +245,7 @@ impl ObjectStore for InstrumentedObjectStore { } } -async fn run_query_and_collect_events( +pub(crate) async fn run_query_and_collect_events( ctx: &SessionContext, sql: &str, events: &Arc>>, @@ -273,99 +272,3 @@ async fn run_query_and_collect_events( } output } - -#[tokio::test] -async fn test_parquet_bloom_filter_io_baseline() { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int32, false), - Field::new("b", DataType::Int32, false), - ])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef, - Arc::new(Int32Array::from(vec![100, 200, 300])) as ArrayRef, - ], - ) - .unwrap(); - - // Disable row group and page statistics so pruning relies entirely on Bloom filters - let props = WriterProperties::builder() - .set_max_row_group_row_count(Some(3)) - .set_bloom_filter_enabled(true) - .set_statistics_enabled(EnabledStatistics::None) - .build(); - - let mut buffer = Vec::new(); - let mut writer = - ArrowWriter::try_new(&mut buffer, batch.schema(), Some(props)).unwrap(); - writer.write(&batch).unwrap(); - writer.close().unwrap(); - - let file_bytes = Bytes::from(buffer); - let file_size = file_bytes.len() as u64; - let layout = ParquetFileLayout::new(&file_bytes); - - let memory_store = Arc::new(InMemory::new()); - let path = Path::from("table.parquet"); - memory_store - .put_opts(&path, PutPayload::from(file_bytes), PutOptions::default()) - .await - .unwrap(); - - let events = Arc::new(Mutex::new(Vec::new())); - let instrumented_store = Arc::new(InstrumentedObjectStore::new( - memory_store, - layout, - file_size, - events.clone(), - )); - - let config = SessionConfig::new() - .with_target_partitions(1) - .with_parquet_bloom_filter_pruning(true); - - let ctx = SessionContext::new_with_config(config); - let table_url = Url::parse("mem://").unwrap(); - ctx.register_object_store(&table_url, instrumented_store); - - ctx.register_parquet("t", "mem:///table.parquet", ParquetReadOptions::default()) - .await - .unwrap(); - - // Case 1: Matching row group (a = 10 AND b = 100) - // Loads Bloom filters, verifies match, then reads data pages and produces batch - let trace_matching = run_query_and_collect_events( - &ctx, - "SELECT * FROM t WHERE a = 10 AND b = 100", - &events, - ) - .await; - - // Case 2: Column 'a' does not match (a = 999 AND b = 100) - // On main, all candidate column Bloom filters ('a' and 'b') are loaded upfront. - // In #25811, incremental pruning will evaluate 'a' first, short-circuit, and skip loading 'b'. - let trace_pruned = run_query_and_collect_events( - &ctx, - "SELECT * FROM t WHERE a = 999 AND b = 100", - &events, - ) - .await; - - let combined_trace = format!("{trace_matching}\n{trace_pruned}"); - assert_snapshot!(combined_trace, @r" - --- Query: SELECT * FROM t WHERE a = 10 AND b = 100 (1 rows) --- - Read Parquet footer (offsets 0..678) - Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) - Read Bloom filter for column 'a' (rg 0) (offsets 98..145) - Read Bloom filter for column 'b' (rg 0) (offsets 145..192) - Read data page for column 'a' (rg 0) (offsets 4..51) - Read data page for column 'b' (rg 0) (offsets 51..98) - Produce batch (1 rows) - - --- Query: SELECT * FROM t WHERE a = 999 AND b = 100 (0 rows) --- - Read Parquet footer (offsets 0..678) - Read Bloom filter for column 'a' (rg 0) (offsets 98..145) - Read Bloom filter for column 'b' (rg 0) (offsets 145..192) - "); -} diff --git a/datafusion/core/tests/parquet/mod.rs b/datafusion/core/tests/parquet/mod.rs index 8fe7184ce6d54..26d0829944c75 100644 --- a/datafusion/core/tests/parquet/mod.rs +++ b/datafusion/core/tests/parquet/mod.rs @@ -31,7 +31,6 @@ use parquet::file::properties::{EnabledStatistics, WriterProperties}; use std::sync::Arc; use tempfile::NamedTempFile; -mod bloom_filter_io; mod content_defined_chunking; mod custom_reader; mod dynamic_row_group_pruning; @@ -41,6 +40,7 @@ mod expr_adapter; mod external_access_plan; mod file_statistics; mod filter_pushdown; +mod io; mod ordering; mod page_pruning; mod row_group_pruning; From a6a9426c13f5b637028daa7820e13288f6094e10 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 2 Nov 2026 15:52:17 -0500 Subject: [PATCH 4/7] Rework tests --- .../core/tests/parquet/io/bloom_filter_io.rs | 133 +++++++----------- datafusion/core/tests/parquet/io/mod.rs | 107 ++++++++++---- 2 files changed, 128 insertions(+), 112 deletions(-) diff --git a/datafusion/core/tests/parquet/io/bloom_filter_io.rs b/datafusion/core/tests/parquet/io/bloom_filter_io.rs index 703b45932c64f..f2536c9483396 100644 --- a/datafusion/core/tests/parquet/io/bloom_filter_io.rs +++ b/datafusion/core/tests/parquet/io/bloom_filter_io.rs @@ -4,38 +4,64 @@ //! Tests tracking storage I/O events during Parquet Bloom filter evaluation. -use std::sync::{Arc, Mutex}; +use std::sync::Arc; use arrow::array::{ArrayRef, Int32Array, RecordBatch}; -use arrow::datatypes::{DataType, Field, Schema}; use bytes::Bytes; -use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use datafusion::prelude::SessionConfig; use insta::assert_snapshot; -use object_store::memory::InMemory; -use object_store::path::Path; -use object_store::{ObjectStore, PutOptions, PutPayload}; use parquet::arrow::ArrowWriter; use parquet::file::properties::{EnabledStatistics, WriterProperties}; -use url::Url; -use super::{InstrumentedObjectStore, ParquetFileLayout, run_query_and_collect_events}; +use super::ParquetIoTest; +/// Matching row group #[tokio::test] -async fn test_parquet_bloom_filter_io_baseline() { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int32, false), - Field::new("b", DataType::Int32, false), - ])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef, - Arc::new(Int32Array::from(vec![100, 200, 300])) as ArrayRef, - ], - ) - .unwrap(); +async fn test_parquet_bloom_filter_io_matching() { + let test = bloom_filter_test().await; + // loads the Bloom filters, finds a match, then reads the data pages and + // produces a batch. + assert_snapshot!(test.query("SELECT * FROM t WHERE a = 10 AND b = 100").await, @r" + --- Query: SELECT * FROM t WHERE a = 10 AND b = 100 (1 rows) --- + Read Parquet footer (offsets 0..678) + Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) + Read Bloom filter for column 'a' (rg 0) (offsets 98..145) + Read Bloom filter for column 'b' (rg 0) (offsets 145..192) + Read data page for column 'a' (rg 0) (offsets 4..51) + Read data page for column 'b' (rg 0) (offsets 51..98) + Produce batch (1 rows) + "); +} + +/// Column `a` does not match, so the row group is pruned. +#[tokio::test] +async fn test_parquet_bloom_filter_io_pruned() { + let test = bloom_filter_test().await; + // Currently the Bloom filters for all candidate columns (`a` and `b`) are + assert_snapshot!(test.query("SELECT * FROM t WHERE a = 999 AND b = 100").await, @r" + --- Query: SELECT * FROM t WHERE a = 999 AND b = 100 (0 rows) --- + Read Parquet footer (offsets 0..678) + Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) + Read Bloom filter for column 'a' (rg 0) (offsets 98..145) + Read Bloom filter for column 'b' (rg 0) (offsets 145..192) + "); +} + +/// Returns a [`ParquetIoTest`] for a file with one row group of 3 rows: +/// +/// | a | b | +/// |----|-----| +/// | 10 | 100 | +/// | 20 | 200 | +/// | 30 | 300 | +/// +/// Both columns have Bloom filters. Statistics are disabled so that row group +/// pruning relies entirely on the Bloom filters. +async fn bloom_filter_test() -> ParquetIoTest { + let a: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30])); + let b: ArrayRef = Arc::new(Int32Array::from(vec![100, 200, 300])); + let batch = RecordBatch::try_from_iter([("a", a), ("b", b)]).unwrap(); - // Disable row group and page statistics so pruning relies entirely on Bloom filters let props = WriterProperties::builder() .set_max_row_group_row_count(Some(3)) .set_bloom_filter_enabled(true) @@ -48,70 +74,9 @@ async fn test_parquet_bloom_filter_io_baseline() { writer.write(&batch).unwrap(); writer.close().unwrap(); - let file_bytes = Bytes::from(buffer); - let file_size = file_bytes.len() as u64; - let layout = ParquetFileLayout::new(&file_bytes); - - let memory_store = Arc::new(InMemory::new()); - let path = Path::from("table.parquet"); - memory_store - .put_opts(&path, PutPayload::from(file_bytes), PutOptions::default()) - .await - .unwrap(); - - let events = Arc::new(Mutex::new(Vec::new())); - let instrumented_store = Arc::new(InstrumentedObjectStore::new( - memory_store, - layout, - file_size, - events.clone(), - )); - let config = SessionConfig::new() .with_target_partitions(1) .with_parquet_bloom_filter_pruning(true); - let ctx = SessionContext::new_with_config(config); - let table_url = Url::parse("mem://").unwrap(); - ctx.register_object_store(&table_url, instrumented_store); - - ctx.register_parquet("t", "mem:///table.parquet", ParquetReadOptions::default()) - .await - .unwrap(); - - // Case 1: Matching row group (a = 10 AND b = 100) - // Loads Bloom filters, verifies match, then reads data pages and produces batch - let trace_matching = run_query_and_collect_events( - &ctx, - "SELECT * FROM t WHERE a = 10 AND b = 100", - &events, - ) - .await; - - // Case 2: Column 'a' does not match (a = 999 AND b = 100) - // On main, all candidate column Bloom filters ('a' and 'b') are loaded upfront. - // In #25811, incremental pruning will evaluate 'a' first, short-circuit, and skip loading 'b'. - let trace_pruned = run_query_and_collect_events( - &ctx, - "SELECT * FROM t WHERE a = 999 AND b = 100", - &events, - ) - .await; - - let combined_trace = format!("{trace_matching}\n{trace_pruned}"); - assert_snapshot!(combined_trace, @r" - --- Query: SELECT * FROM t WHERE a = 10 AND b = 100 (1 rows) --- - Read Parquet footer (offsets 0..678) - Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) - Read Bloom filter for column 'a' (rg 0) (offsets 98..145) - Read Bloom filter for column 'b' (rg 0) (offsets 145..192) - Read data page for column 'a' (rg 0) (offsets 4..51) - Read data page for column 'b' (rg 0) (offsets 51..98) - Produce batch (1 rows) - - --- Query: SELECT * FROM t WHERE a = 999 AND b = 100 (0 rows) --- - Read Parquet footer (offsets 0..678) - Read Bloom filter for column 'a' (rg 0) (offsets 98..145) - Read Bloom filter for column 'b' (rg 0) (offsets 145..192) - "); + ParquetIoTest::new(Bytes::from(buffer), config).await } diff --git a/datafusion/core/tests/parquet/io/mod.rs b/datafusion/core/tests/parquet/io/mod.rs index 1617f2716804c..ce2546ded33c8 100644 --- a/datafusion/core/tests/parquet/io/mod.rs +++ b/datafusion/core/tests/parquet/io/mod.rs @@ -17,15 +17,17 @@ use std::sync::{Arc, Mutex}; use async_trait::async_trait; use bytes::Bytes; -use datafusion::prelude::SessionContext; +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; use futures::StreamExt; use futures::stream::BoxStream; +use object_store::memory::InMemory; use object_store::path::Path; use object_store::{ GetOptions, GetRange, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, }; use parquet::file::reader::{FileReader, SerializedFileReader}; +use url::Url; mod bloom_filter_io; @@ -36,12 +38,12 @@ struct ByteRangeDescriptor { } #[derive(Debug, Clone, Default)] -pub(crate) struct ParquetFileLayout { +struct ParquetFileLayout { descriptors: Vec, } impl ParquetFileLayout { - pub(crate) fn new(bytes: &Bytes) -> Self { + fn new(bytes: &Bytes) -> Self { let cursor = bytes.clone(); let reader = SerializedFileReader::new(cursor).expect("valid parquet file"); let metadata = reader.metadata(); @@ -118,7 +120,7 @@ impl ParquetFileLayout { } #[derive(Debug)] -pub(crate) struct InstrumentedObjectStore { +struct InstrumentedObjectStore { inner: Arc, layout: ParquetFileLayout, file_size: u64, @@ -126,7 +128,7 @@ pub(crate) struct InstrumentedObjectStore { } impl InstrumentedObjectStore { - pub(crate) fn new( + fn new( inner: Arc, layout: ParquetFileLayout, file_size: u64, @@ -245,30 +247,79 @@ impl ObjectStore for InstrumentedObjectStore { } } -pub(crate) async fn run_query_and_collect_events( - ctx: &SessionContext, - sql: &str, - events: &Arc>>, -) -> String { - events.lock().unwrap().clear(); - - let df = ctx.sql(sql).await.unwrap(); - let mut stream = df.execute_stream().await.unwrap(); - - let mut row_count = 0; - while let Some(batch) = stream.next().await { - let batch = batch.unwrap(); - row_count += batch.num_rows(); - events - .lock() - .unwrap() - .push(format!("Produce batch ({} rows)", batch.num_rows())); +/// Test fixture: a single Parquet file registered as table `t`, read through an +/// [`InstrumentedObjectStore`] that records every storage request. +/// +/// Typical use: +/// +/// ```text +/// let test = ParquetIoTest::new(file_bytes, SessionConfig::new()).await; +/// assert_snapshot!(test.query("SELECT * FROM t WHERE ...").await, @"..."); +/// ``` +pub(crate) struct ParquetIoTest { + ctx: SessionContext, + events: Arc>>, +} + +impl ParquetIoTest { + /// Stores `file_bytes` as `mem:///table.parquet` and registers it as table + /// `t` in a new [`SessionContext`] created with `config`. + pub(crate) async fn new(file_bytes: Bytes, config: SessionConfig) -> Self { + let file_size = file_bytes.len() as u64; + let layout = ParquetFileLayout::new(&file_bytes); + + let memory_store = Arc::new(InMemory::new()); + memory_store + .put_opts( + &Path::from("table.parquet"), + PutPayload::from(file_bytes), + PutOptions::default(), + ) + .await + .unwrap(); + + let events = Arc::new(Mutex::new(Vec::new())); + let instrumented_store = Arc::new(InstrumentedObjectStore::new( + memory_store, + layout, + file_size, + Arc::clone(&events), + )); + + let ctx = SessionContext::new_with_config(config); + ctx.register_object_store(&Url::parse("mem://").unwrap(), instrumented_store); + ctx.register_parquet("t", "mem:///table.parquet", ParquetReadOptions::default()) + .await + .unwrap(); + + Self { ctx, events } } - let recorded = events.lock().unwrap().clone(); - let mut output = format!("--- Query: {sql} ({row_count} rows) ---\n"); - for event in recorded { - let _ = writeln!(output, "{event}"); + /// Runs `sql` and returns the storage requests and output batches it + /// produced, formatted for snapshot comparison. + /// + /// Only events recorded while running this query are included. + pub(crate) async fn query(&self, sql: &str) -> String { + self.events.lock().unwrap().clear(); + + let df = self.ctx.sql(sql).await.unwrap(); + let mut stream = df.execute_stream().await.unwrap(); + + let mut row_count = 0; + while let Some(batch) = stream.next().await { + let batch = batch.unwrap(); + row_count += batch.num_rows(); + self.events + .lock() + .unwrap() + .push(format!("Produce batch ({} rows)", batch.num_rows())); + } + + let recorded = self.events.lock().unwrap().clone(); + let mut output = format!("--- Query: {sql} ({row_count} rows) ---\n"); + for event in recorded { + let _ = writeln!(output, "{event}"); + } + output } - output } From 1522d2c1521dbbd44f9a659fa54592338fe4346a Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 2 Nov 2026 15:53:07 -0500 Subject: [PATCH 5/7] test(parquet): complete comment on pruned Bloom filter I/O test Co-Authored-By: Claude Opus 5.5 --- datafusion/core/tests/parquet/io/bloom_filter_io.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/datafusion/core/tests/parquet/io/bloom_filter_io.rs b/datafusion/core/tests/parquet/io/bloom_filter_io.rs index f2536c9483396..bc86e0e699b1b 100644 --- a/datafusion/core/tests/parquet/io/bloom_filter_io.rs +++ b/datafusion/core/tests/parquet/io/bloom_filter_io.rs @@ -19,7 +19,7 @@ use super::ParquetIoTest; #[tokio::test] async fn test_parquet_bloom_filter_io_matching() { let test = bloom_filter_test().await; - // loads the Bloom filters, finds a match, then reads the data pages and + // Loads the Bloom filters, finds a match, then reads the data pages and // produces a batch. assert_snapshot!(test.query("SELECT * FROM t WHERE a = 10 AND b = 100").await, @r" --- Query: SELECT * FROM t WHERE a = 10 AND b = 100 (1 rows) --- @@ -38,6 +38,9 @@ async fn test_parquet_bloom_filter_io_matching() { async fn test_parquet_bloom_filter_io_pruned() { let test = bloom_filter_test().await; // Currently the Bloom filters for all candidate columns (`a` and `b`) are + // loaded upfront. With incremental pruning + // (https://github.com/apache/datafusion/issues/25811), `a` would be + // evaluated first and the Bloom filter for `b` would not be loaded. assert_snapshot!(test.query("SELECT * FROM t WHERE a = 999 AND b = 100").await, @r" --- Query: SELECT * FROM t WHERE a = 999 AND b = 100 (0 rows) --- Read Parquet footer (offsets 0..678) From f1517f66e91670b067188f798da13451a100340f Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 2 Nov 2026 15:54:56 -0500 Subject: [PATCH 6/7] comments --- datafusion/core/tests/parquet/io/mod.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/datafusion/core/tests/parquet/io/mod.rs b/datafusion/core/tests/parquet/io/mod.rs index ce2546ded33c8..d3a0c09c0fddc 100644 --- a/datafusion/core/tests/parquet/io/mod.rs +++ b/datafusion/core/tests/parquet/io/mod.rs @@ -6,10 +6,7 @@ //! //! Provides an instrumented [`ObjectStore`] that records storage requests and //! correlates each requested byte range with the Parquet file structure it -//! covers (footer, page index, Bloom filter, data pages), so the resulting -//! trace can be compared against `insta` snapshots. -//! -//! The actual tests live in the submodules. +//! covers, so the resulting trace can be compared against `insta` snapshots. use std::fmt::Write as _; use std::ops::Range; @@ -29,6 +26,7 @@ use object_store::{ use parquet::file::reader::{FileReader, SerializedFileReader}; use url::Url; +// Tests are in submodules mod bloom_filter_io; #[derive(Debug, Clone)] From 2fa72b06d125581232c93e3de05e03b9d28a930e Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 2 Nov 2026 16:03:25 -0500 Subject: [PATCH 7/7] note where similar tests are --- datafusion/core/tests/datasource/object_store_access.rs | 3 +++ datafusion/core/tests/parquet/io/mod.rs | 3 +++ 2 files changed, 6 insertions(+) diff --git a/datafusion/core/tests/datasource/object_store_access.rs b/datafusion/core/tests/datasource/object_store_access.rs index e44e14df213e8..d9dffa33bda3f 100644 --- a/datafusion/core/tests/datasource/object_store_access.rs +++ b/datafusion/core/tests/datasource/object_store_access.rs @@ -9,6 +9,9 @@ //! various operations (table creation, queries with and without predicates) //! to verify the expected object store access patterns. //! +//! File-level Parquet access tests live in +//! [`datafusion/core/tests/parquet/io`](https://github.com/apache/datafusion/blob/main/datafusion/core/tests/parquet/io/mod.rs). +//! //! [`ListingTable`]: datafusion::datasource::listing::ListingTable use arrow::array::{ArrayRef, Int32Array, RecordBatch}; diff --git a/datafusion/core/tests/parquet/io/mod.rs b/datafusion/core/tests/parquet/io/mod.rs index d3a0c09c0fddc..839d7a981051e 100644 --- a/datafusion/core/tests/parquet/io/mod.rs +++ b/datafusion/core/tests/parquet/io/mod.rs @@ -7,6 +7,9 @@ //! Provides an instrumented [`ObjectStore`] that records storage requests and //! correlates each requested byte range with the Parquet file structure it //! covers, so the resulting trace can be compared against `insta` snapshots. +//! +//! `ListingTable`object store access tests live in +//! [`datafusion/core/tests/datasource/object_store_access.rs`](https://github.com/apache/datafusion/blob/main/datafusion/core/tests/datasource/object_store_access.rs). use std::fmt::Write as _; use std::ops::Range;