From c0a70fb733f81d0165b5efb05a72f98e451eb212 Mon Sep 17 00:00:00 2001 From: mohitgurav20 Date: Fri, 30 Oct 2026 00:20:25 +0530 Subject: [PATCH 01/11] 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 02/11] 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 03/11] 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 04/11] 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 05/11] 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 06/11] 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 e85a89fc00e575d793fae2929ee360f2f7122ff0 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 2 Nov 2026 16:06:46 -0500 Subject: [PATCH 07/11] test(parquet): add I/O tracking tests for row group pruning, filter pushdown and TopK Co-Authored-By: Claude Opus 5.5 --- .../tests/parquet/io/filter_pushdown_io.rs | 59 ++++++++++++++ datafusion/core/tests/parquet/io/mod.rs | 35 +++++++++ .../tests/parquet/io/row_group_pruning_io.rs | 77 +++++++++++++++++++ datafusion/core/tests/parquet/io/topk_io.rs | 55 +++++++++++++ 4 files changed, 226 insertions(+) create mode 100644 datafusion/core/tests/parquet/io/filter_pushdown_io.rs create mode 100644 datafusion/core/tests/parquet/io/row_group_pruning_io.rs create mode 100644 datafusion/core/tests/parquet/io/topk_io.rs diff --git a/datafusion/core/tests/parquet/io/filter_pushdown_io.rs b/datafusion/core/tests/parquet/io/filter_pushdown_io.rs new file mode 100644 index 0000000000000..5ec6ec6c14834 --- /dev/null +++ b/datafusion/core/tests/parquet/io/filter_pushdown_io.rs @@ -0,0 +1,59 @@ +// 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 with and without Parquet filter pushdown +//! (`datafusion.execution.parquet.pushdown_filters`). +//! +//! Both tests run `val = 7`, which statistics cannot prune (every row group +//! has `val` in `0..=10`), but which only matches a row in row group 1. + +use datafusion::prelude::SessionConfig; +use insta::assert_snapshot; + +use super::{ParquetIoTest, three_row_group_file}; + +/// Without filter pushdown, the scan fetches every column of every row group +/// and the filter is applied after decoding. +#[tokio::test] +async fn test_filter_pushdown_io_disabled() { + let test = filter_pushdown_test(false).await; + assert_snapshot!(test.query("SELECT * FROM t WHERE val = 7").await, @r" + --- Query: SELECT * FROM t WHERE val = 7 (1 rows) --- + Read Parquet footer (offsets 0..1402) + Read data page for column 'id' (rg 0) (offsets 4..51) + Read data page for column 'val' (rg 0) (offsets 51..98) + Read data page for column 'id' (rg 1) (offsets 98..145) + Read data page for column 'val' (rg 1) (offsets 145..192) + Read data page for column 'id' (rg 2) (offsets 192..239) + Read data page for column 'val' (rg 2) (offsets 239..286) + Produce batch (1 rows) + "); +} + +/// With filter pushdown, the scan first fetches and evaluates the filter +/// column (`val`) for each row group, and then fetches the remaining columns +/// (`id`) only for the row groups with matching rows (row group 1). +/// +/// Note that `id` for row group 1 is fetched as two ranges, its dictionary +/// page and its data page, rather than as a single column chunk read. +#[tokio::test] +async fn test_filter_pushdown_io_enabled() { + let test = filter_pushdown_test(true).await; + assert_snapshot!(test.query("SELECT * FROM t WHERE val = 7").await, @r" + --- Query: SELECT * FROM t WHERE val = 7 (1 rows) --- + Read Parquet footer (offsets 0..1402) + Read data page for column 'val' (rg 0) (offsets 51..98) + Read data page for column 'val' (rg 1) (offsets 145..192) + Read data page for column 'id' (rg 1) (offsets 98..124) + Read data page for column 'id' (rg 1) (offsets 124..145) + Produce batch (1 rows) + Read data page for column 'val' (rg 2) (offsets 239..286) + "); +} + +async fn filter_pushdown_test(pushdown_filters: bool) -> ParquetIoTest { + let mut config = SessionConfig::new().with_target_partitions(1); + config.options_mut().execution.parquet.pushdown_filters = pushdown_filters; + ParquetIoTest::new(three_row_group_file(), config).await +} diff --git a/datafusion/core/tests/parquet/io/mod.rs b/datafusion/core/tests/parquet/io/mod.rs index d3a0c09c0fddc..9c48fc69f42a5 100644 --- a/datafusion/core/tests/parquet/io/mod.rs +++ b/datafusion/core/tests/parquet/io/mod.rs @@ -12,6 +12,7 @@ use std::fmt::Write as _; use std::ops::Range; use std::sync::{Arc, Mutex}; +use arrow::array::{ArrayRef, Int32Array, RecordBatch}; use async_trait::async_trait; use bytes::Bytes; use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; @@ -23,11 +24,16 @@ use object_store::{ GetOptions, GetRange, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, }; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::WriterProperties; use parquet::file::reader::{FileReader, SerializedFileReader}; use url::Url; // Tests are in submodules mod bloom_filter_io; +mod filter_pushdown_io; +mod row_group_pruning_io; +mod topk_io; #[derive(Debug, Clone)] struct ByteRangeDescriptor { @@ -321,3 +327,32 @@ impl ParquetIoTest { output } } + +/// Returns a Parquet file with three row groups of three rows each: +/// +/// | row group | id | val | +/// |-----------|---------|-----------| +/// | 0 | 1, 2, 3 | 0, 5, 10 | +/// | 1 | 4, 5, 6 | 0, 7, 10 | +/// | 2 | 7, 8, 9 | 0, 5, 10 | +/// +/// `id` is sorted, so its row group statistics are disjoint and predicates on +/// it can prune row groups. `val` has the same min/max (`0`/`10`) in every row +/// group, so statistics cannot prune row groups for predicates such as +/// `val = 7`, even though only row group 1 contains a match. +pub(crate) fn three_row_group_file() -> Bytes { + let id: ArrayRef = Arc::new(Int32Array::from_iter_values(1..=9)); + let val: ArrayRef = Arc::new(Int32Array::from(vec![0, 5, 10, 0, 7, 10, 0, 5, 10])); + let batch = RecordBatch::try_from_iter([("id", id), ("val", val)]).unwrap(); + + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .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(); + Bytes::from(buffer) +} diff --git a/datafusion/core/tests/parquet/io/row_group_pruning_io.rs b/datafusion/core/tests/parquet/io/row_group_pruning_io.rs new file mode 100644 index 0000000000000..23cecaf0acede --- /dev/null +++ b/datafusion/core/tests/parquet/io/row_group_pruning_io.rs @@ -0,0 +1,77 @@ +// 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 row group pruning based +//! on row group statistics (min/max values). + +use datafusion::prelude::SessionConfig; +use insta::assert_snapshot; + +use super::{ParquetIoTest, three_row_group_file}; + +/// No predicate, so every row group is read. +#[tokio::test] +async fn test_row_group_pruning_io_no_predicate() { + let test = row_group_pruning_test().await; + assert_snapshot!(test.query("SELECT * FROM t").await, @r" + --- Query: SELECT * FROM t (9 rows) --- + Read Parquet footer (offsets 0..1402) + Read data page for column 'id' (rg 0) (offsets 4..51) + Read data page for column 'val' (rg 0) (offsets 51..98) + Produce batch (3 rows) + Read data page for column 'id' (rg 1) (offsets 98..145) + Read data page for column 'val' (rg 1) (offsets 145..192) + Produce batch (3 rows) + Read data page for column 'id' (rg 2) (offsets 192..239) + Read data page for column 'val' (rg 2) (offsets 239..286) + Produce batch (3 rows) + "); +} + +/// Only row group 1 (`id` in `4..=6`) can contain `id = 5`, so the data pages +/// for row groups 0 and 2 are never fetched. +#[tokio::test] +async fn test_row_group_pruning_io_one_row_group() { + let test = row_group_pruning_test().await; + assert_snapshot!(test.query("SELECT * FROM t WHERE id = 5").await, @r" + --- Query: SELECT * FROM t WHERE id = 5 (1 rows) --- + Read Parquet footer (offsets 0..1402) + Read data page for column 'id' (rg 1) (offsets 98..145) + Read data page for column 'val' (rg 1) (offsets 145..192) + Produce batch (1 rows) + "); +} + +/// Row group 0 (`id` in `1..=3`) cannot contain `id > 5`, so only row groups 1 +/// and 2 are fetched. +#[tokio::test] +async fn test_row_group_pruning_io_two_row_groups() { + let test = row_group_pruning_test().await; + assert_snapshot!(test.query("SELECT * FROM t WHERE id > 5").await, @r" + --- Query: SELECT * FROM t WHERE id > 5 (4 rows) --- + Read Parquet footer (offsets 0..1402) + Read data page for column 'id' (rg 1) (offsets 98..145) + Read data page for column 'val' (rg 1) (offsets 145..192) + Read data page for column 'id' (rg 2) (offsets 192..239) + Read data page for column 'val' (rg 2) (offsets 239..286) + Produce batch (4 rows) + "); +} + +/// No row group can contain `id > 100`, so only the footer is read. +#[tokio::test] +async fn test_row_group_pruning_io_all_pruned() { + let test = row_group_pruning_test().await; + assert_snapshot!(test.query("SELECT * FROM t WHERE id > 100").await, @r" + --- Query: SELECT * FROM t WHERE id > 100 (0 rows) --- + Read Parquet footer (offsets 0..1402) + "); +} + +async fn row_group_pruning_test() -> ParquetIoTest { + let config = SessionConfig::new() + .with_target_partitions(1) + .with_parquet_pruning(true); + ParquetIoTest::new(three_row_group_file(), config).await +} diff --git a/datafusion/core/tests/parquet/io/topk_io.rs b/datafusion/core/tests/parquet/io/topk_io.rs new file mode 100644 index 0000000000000..ffe1f343c6395 --- /dev/null +++ b/datafusion/core/tests/parquet/io/topk_io.rs @@ -0,0 +1,55 @@ +// 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 for TopK (`ORDER BY .. LIMIT`) queries. + +use datafusion::prelude::SessionConfig; +use insta::assert_snapshot; + +use super::{ParquetIoTest, three_row_group_file}; + +/// TopK dynamic filter pushdown prunes row groups during the scan. +/// +/// Sort pushdown reorders the row groups using their statistics so row group +/// 2 (largest `id` values) is read first. After it is read, the TopK heap holds +/// `id` 9 and 8, and its dynamic filter (`id > 8`) proves that row groups 1 +/// and 0 cannot contain better rows, so they are pruned without being fetched. +#[tokio::test] +async fn test_topk_io_dynamic_filter() { + let test = topk_test(true).await; + assert_snapshot!(test.query("SELECT * FROM t ORDER BY id DESC LIMIT 2").await, @r" + --- Query: SELECT * FROM t ORDER BY id DESC LIMIT 2 (2 rows) --- + Read Parquet footer (offsets 0..1402) + Read data page for column 'id' (rg 2) (offsets 192..239) + Read data page for column 'val' (rg 2) (offsets 239..286) + Produce batch (2 rows) + "); +} + +/// Without TopK dynamic filter pushdown, the row groups are still read in the +/// optimal order, but all of them are fetched. +#[tokio::test] +async fn test_topk_io_no_dynamic_filter() { + let test = topk_test(false).await; + assert_snapshot!(test.query("SELECT * FROM t ORDER BY id DESC LIMIT 2").await, @r" + --- Query: SELECT * FROM t ORDER BY id DESC LIMIT 2 (2 rows) --- + Read Parquet footer (offsets 0..1402) + Read data page for column 'id' (rg 2) (offsets 192..239) + Read data page for column 'val' (rg 2) (offsets 239..286) + Read data page for column 'id' (rg 1) (offsets 98..145) + Read data page for column 'val' (rg 1) (offsets 145..192) + Read data page for column 'id' (rg 0) (offsets 4..51) + Read data page for column 'val' (rg 0) (offsets 51..98) + Produce batch (2 rows) + "); +} + +async fn topk_test(dynamic_filter_pushdown: bool) -> ParquetIoTest { + let mut config = SessionConfig::new().with_target_partitions(1); + config + .options_mut() + .optimizer + .enable_topk_dynamic_filter_pushdown = dynamic_filter_pushdown; + ParquetIoTest::new(three_row_group_file(), config).await +} From e92655f3917385afc8d1e11ed4d63596b877be61 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 4 Nov 2026 15:36:51 -0500 Subject: [PATCH 08/11] Comments --- .../core/tests/parquet/io/bloom_filter_io.rs | 18 ++++++++++ .../tests/parquet/io/filter_pushdown_io.rs | 33 ++++++++++++++----- .../tests/parquet/io/row_group_pruning_io.rs | 22 +++++++++++-- datafusion/core/tests/parquet/io/topk_io.rs | 30 ++++++++++++++--- 4 files changed, 87 insertions(+), 16 deletions(-) diff --git a/datafusion/core/tests/parquet/io/bloom_filter_io.rs b/datafusion/core/tests/parquet/io/bloom_filter_io.rs index bc86e0e699b1b..7bdcf95c52985 100644 --- a/datafusion/core/tests/parquet/io/bloom_filter_io.rs +++ b/datafusion/core/tests/parquet/io/bloom_filter_io.rs @@ -3,6 +3,24 @@ // 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. +//! +//! The file has a single row group with Bloom filters on both columns and no +//! statistics, so only the Bloom filters are used to prune. +//! +//! ```text +//! row group 0 +//! ┌──────┬───────┐ +//! │ a │ b │ +//! ├──────┼───────┤ +//! │ 10 │ 100 │ +//! │ 20 │ 200 │ +//! │ 30 │ 300 │ +//! └──────┴───────┘ +//! Bloom filter (a) Bloom filter (b) no min/max statistics +//! +//! a = 10 AND b = 100 -> Bloom filters match, row group read +//! a = 999 AND b = 100 -> Bloom filter on `a` misses, row group pruned +//! ``` use std::sync::Arc; diff --git a/datafusion/core/tests/parquet/io/filter_pushdown_io.rs b/datafusion/core/tests/parquet/io/filter_pushdown_io.rs index 5ec6ec6c14834..4feab2b1f10e2 100644 --- a/datafusion/core/tests/parquet/io/filter_pushdown_io.rs +++ b/datafusion/core/tests/parquet/io/filter_pushdown_io.rs @@ -2,11 +2,28 @@ // 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 with and without Parquet filter pushdown -//! (`datafusion.execution.parquet.pushdown_filters`). +//! Tests tracking I/O events with and without Parquet filter pushdown +//! (`datafusion.execution.parquet.pushdown_filters`, see +//! [`ParquetOptions::pushdown_filters`]). //! -//! Both tests run `val = 7`, which statistics cannot prune (every row group -//! has `val` in `0..=10`), but which only matches a row in row group 1. +//! Both tests run `val = 7`, which statistics cannot prune, but which only +//! matches a row in row group 1: +//! +//! ```text +//! row group 0 row group 1 row group 2 +//! ┌──────┬───────┐ ┌──────┬───────┐ ┌──────┬───────┐ +//! │ id │ val │ │ id │ val │ │ id │ val │ +//! ├──────┼───────┤ ├──────┼───────┤ ├──────┼───────┤ +//! │ 1 │ 0 │ │ 4 │ 0 │ │ 7 │ 0 │ +//! │ 2 │ 5 │ │ 5 │ 7 * │ │ 8 │ 5 │ +//! │ 3 │ 10 │ │ 6 │ 10 │ │ 9 │ 10 │ +//! └──────┴───────┘ └──────┴───────┘ └──────┴───────┘ +//! val: min 0 max 10 val: min 0 max 10 val: min 0 max 10 +//! +//! * the only row matching `val = 7` +//! ``` +//! +//! [`ParquetOptions::pushdown_filters`]: datafusion::config::ParquetOptions::pushdown_filters use datafusion::prelude::SessionConfig; use insta::assert_snapshot; @@ -14,7 +31,6 @@ use insta::assert_snapshot; use super::{ParquetIoTest, three_row_group_file}; /// Without filter pushdown, the scan fetches every column of every row group -/// and the filter is applied after decoding. #[tokio::test] async fn test_filter_pushdown_io_disabled() { let test = filter_pushdown_test(false).await; @@ -31,15 +47,14 @@ async fn test_filter_pushdown_io_disabled() { "); } -/// With filter pushdown, the scan first fetches and evaluates the filter +/// **With** filter pushdown, the scan first fetches and evaluates the filter /// column (`val`) for each row group, and then fetches the remaining columns /// (`id`) only for the row groups with matching rows (row group 1). -/// -/// Note that `id` for row group 1 is fetched as two ranges, its dictionary -/// page and its data page, rather than as a single column chunk read. #[tokio::test] async fn test_filter_pushdown_io_enabled() { let test = filter_pushdown_test(true).await; + // Note: `id` for row group 1 is fetched as two ranges, its dictionary + // page and its data page, rather than as a single column chunk read. assert_snapshot!(test.query("SELECT * FROM t WHERE val = 7").await, @r" --- Query: SELECT * FROM t WHERE val = 7 (1 rows) --- Read Parquet footer (offsets 0..1402) diff --git a/datafusion/core/tests/parquet/io/row_group_pruning_io.rs b/datafusion/core/tests/parquet/io/row_group_pruning_io.rs index 23cecaf0acede..c1dd57ddc41aa 100644 --- a/datafusion/core/tests/parquet/io/row_group_pruning_io.rs +++ b/datafusion/core/tests/parquet/io/row_group_pruning_io.rs @@ -2,8 +2,26 @@ // 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 row group pruning based -//! on row group statistics (min/max values). +//! Tests tracking I/O events during Parquet row group pruning +//! +//! All tests filter on `id`, whose values are sorted across row groups so the +//! row group statistics are disjoint and can prune: +//! +//! ```text +//! row group 0 row group 1 row group 2 +//! ┌──────┬───────┐ ┌──────┬───────┐ ┌──────┬───────┐ +//! │ id │ val │ │ id │ val │ │ id │ val │ +//! ├──────┼───────┤ ├──────┼───────┤ ├──────┼───────┤ +//! │ 1 │ 0 │ │ 4 │ 0 │ │ 7 │ 0 │ +//! │ 2 │ 5 │ │ 5 │ 7 │ │ 8 │ 5 │ +//! │ 3 │ 10 │ │ 6 │ 10 │ │ 9 │ 10 │ +//! └──────┴───────┘ └──────┴───────┘ └──────┴───────┘ +//! id: min 1 max 3 id: min 4 max 6 id: min 7 max 9 +//! +//! id = 5 -> row group 1 only +//! id > 5 -> row groups 1 and 2 +//! id > 100 -> no row groups +//! ``` use datafusion::prelude::SessionConfig; use insta::assert_snapshot; diff --git a/datafusion/core/tests/parquet/io/topk_io.rs b/datafusion/core/tests/parquet/io/topk_io.rs index ffe1f343c6395..7d7e9da3bddc7 100644 --- a/datafusion/core/tests/parquet/io/topk_io.rs +++ b/datafusion/core/tests/parquet/io/topk_io.rs @@ -2,7 +2,27 @@ // 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 for TopK (`ORDER BY .. LIMIT`) queries. +//! Tests tracking I/O events for TopK (`ORDER BY .. LIMIT`) queries. +//! +//! Both tests run `ORDER BY id DESC LIMIT 2`. `id` is sorted across row +//! groups, so sort pushdown reads the row groups in reverse order (largest +//! `id` first) and the two matching rows are both in row group 2: +//! +//! ```text +//! row group 0 row group 1 row group 2 +//! ┌──────┬───────┐ ┌──────┬───────┐ ┌──────┬───────┐ +//! │ id │ val │ │ id │ val │ │ id │ val │ +//! ├──────┼───────┤ ├──────┼───────┤ ├──────┼───────┤ +//! │ 1 │ 0 │ │ 4 │ 0 │ │ 7 │ 0 │ +//! │ 2 │ 5 │ │ 5 │ 7 │ │ 8 │ 5 * │ +//! │ 3 │ 10 │ │ 6 │ 10 │ │ 9 │ 10 * │ +//! └──────┴───────┘ └──────┴───────┘ └──────┴───────┘ +//! id: min 1 max 3 id: min 4 max 6 id: min 7 max 9 +//! +//! read order: 3rd read order: 2nd read order: 1st +//! +//! * the two rows returned by `ORDER BY id DESC LIMIT 2` +//! ``` use datafusion::prelude::SessionConfig; use insta::assert_snapshot; @@ -11,10 +31,10 @@ use super::{ParquetIoTest, three_row_group_file}; /// TopK dynamic filter pushdown prunes row groups during the scan. /// -/// Sort pushdown reorders the row groups using their statistics so row group -/// 2 (largest `id` values) is read first. After it is read, the TopK heap holds -/// `id` 9 and 8, and its dynamic filter (`id > 8`) proves that row groups 1 -/// and 0 cannot contain better rows, so they are pruned without being fetched. +/// Sort pushdown reorders the row groups using their statistics so row group 2 +/// (largest `id` values) is read first. After it is read, the TopK filter is +/// `id > 8` which proves that row groups 1 and 0 cannot contain better rows, +/// so they are pruned without being fetched. #[tokio::test] async fn test_topk_io_dynamic_filter() { let test = topk_test(true).await; From 5da1a02885f7c91c1cb75f787a7b339d4d3c58c0 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 4 Nov 2026 15:49:37 -0500 Subject: [PATCH 09/11] Update tests --- .../core/tests/parquet/io/bloom_filter_io.rs | 4 +- .../tests/parquet/io/filter_pushdown_io.rs | 55 +++++-- datafusion/core/tests/parquet/io/mod.rs | 154 +++++++++++++----- .../core/tests/parquet/io/page_pruning_io.rs | 129 +++++++++++++++ .../tests/parquet/io/row_group_pruning_io.rs | 51 ++++-- datafusion/core/tests/parquet/io/topk_io.rs | 43 ++++- 6 files changed, 355 insertions(+), 81 deletions(-) create mode 100644 datafusion/core/tests/parquet/io/page_pruning_io.rs diff --git a/datafusion/core/tests/parquet/io/bloom_filter_io.rs b/datafusion/core/tests/parquet/io/bloom_filter_io.rs index 7bdcf95c52985..b9b7f2213cf1d 100644 --- a/datafusion/core/tests/parquet/io/bloom_filter_io.rs +++ b/datafusion/core/tests/parquet/io/bloom_filter_io.rs @@ -45,8 +45,8 @@ async fn test_parquet_bloom_filter_io_matching() { 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) + Read dictionary page and data page 0 for column 'a' (rg 0) (offsets 4..51) + Read dictionary page and data page 0 for column 'b' (rg 0) (offsets 51..98) Produce batch (1 rows) "); } diff --git a/datafusion/core/tests/parquet/io/filter_pushdown_io.rs b/datafusion/core/tests/parquet/io/filter_pushdown_io.rs index 4feab2b1f10e2..bcf2b3b925985 100644 --- a/datafusion/core/tests/parquet/io/filter_pushdown_io.rs +++ b/datafusion/core/tests/parquet/io/filter_pushdown_io.rs @@ -25,10 +25,16 @@ //! //! [`ParquetOptions::pushdown_filters`]: datafusion::config::ParquetOptions::pushdown_filters +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int32Array, RecordBatch}; +use bytes::Bytes; use datafusion::prelude::SessionConfig; use insta::assert_snapshot; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::WriterProperties; -use super::{ParquetIoTest, three_row_group_file}; +use super::ParquetIoTest; /// Without filter pushdown, the scan fetches every column of every row group #[tokio::test] @@ -37,12 +43,12 @@ async fn test_filter_pushdown_io_disabled() { assert_snapshot!(test.query("SELECT * FROM t WHERE val = 7").await, @r" --- Query: SELECT * FROM t WHERE val = 7 (1 rows) --- Read Parquet footer (offsets 0..1402) - Read data page for column 'id' (rg 0) (offsets 4..51) - Read data page for column 'val' (rg 0) (offsets 51..98) - Read data page for column 'id' (rg 1) (offsets 98..145) - Read data page for column 'val' (rg 1) (offsets 145..192) - Read data page for column 'id' (rg 2) (offsets 192..239) - Read data page for column 'val' (rg 2) (offsets 239..286) + Read dictionary page and data page 0 for column 'id' (rg 0) (offsets 4..51) + Read dictionary page and data page 0 for column 'val' (rg 0) (offsets 51..98) + Read dictionary page and data page 0 for column 'id' (rg 1) (offsets 98..145) + Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) + Read dictionary page and data page 0 for column 'id' (rg 2) (offsets 192..239) + Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) Produce batch (1 rows) "); } @@ -55,15 +61,21 @@ async fn test_filter_pushdown_io_enabled() { let test = filter_pushdown_test(true).await; // Note: `id` for row group 1 is fetched as two ranges, its dictionary // page and its data page, rather than as a single column chunk read. + // Once the filter has produced a row selection, the parquet reader uses + // the offset index to request only the pages needed for the selected + // rows (plus the dictionary page), one range per page, even when the + // ranges are adjacent. The object store's default `get_ranges` + // coalesces adjacent ranges into a single request, so this does not + // necessarily mean two requests to the storage system. assert_snapshot!(test.query("SELECT * FROM t WHERE val = 7").await, @r" --- Query: SELECT * FROM t WHERE val = 7 (1 rows) --- Read Parquet footer (offsets 0..1402) - Read data page for column 'val' (rg 0) (offsets 51..98) - Read data page for column 'val' (rg 1) (offsets 145..192) - Read data page for column 'id' (rg 1) (offsets 98..124) - Read data page for column 'id' (rg 1) (offsets 124..145) + Read dictionary page and data page 0 for column 'val' (rg 0) (offsets 51..98) + Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) + Read dictionary page for column 'id' (rg 1) (offsets 98..124) + Read data page 0 for column 'id' (rg 1) (offsets 124..145) Produce batch (1 rows) - Read data page for column 'val' (rg 2) (offsets 239..286) + Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) "); } @@ -72,3 +84,22 @@ async fn filter_pushdown_test(pushdown_filters: bool) -> ParquetIoTest { config.options_mut().execution.parquet.pushdown_filters = pushdown_filters; ParquetIoTest::new(three_row_group_file(), config).await } + +/// Returns the Parquet file described in the module docs: three row groups of +/// three rows each, with `id` sorted across row groups. +fn three_row_group_file() -> Bytes { + let id: ArrayRef = Arc::new(Int32Array::from_iter_values(1..=9)); + let val: ArrayRef = Arc::new(Int32Array::from(vec![0, 5, 10, 0, 7, 10, 0, 5, 10])); + let batch = RecordBatch::try_from_iter([("id", id), ("val", val)]).unwrap(); + + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .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(); + Bytes::from(buffer) +} diff --git a/datafusion/core/tests/parquet/io/mod.rs b/datafusion/core/tests/parquet/io/mod.rs index eaf3303301054..f1308b4a8f06f 100644 --- a/datafusion/core/tests/parquet/io/mod.rs +++ b/datafusion/core/tests/parquet/io/mod.rs @@ -15,7 +15,6 @@ use std::fmt::Write as _; use std::ops::Range; use std::sync::{Arc, Mutex}; -use arrow::array::{ArrayRef, Int32Array, RecordBatch}; use async_trait::async_trait; use bytes::Bytes; use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; @@ -27,14 +26,14 @@ use object_store::{ GetOptions, GetRange, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, }; -use parquet::arrow::ArrowWriter; -use parquet::file::properties::WriterProperties; use parquet::file::reader::{FileReader, SerializedFileReader}; +use parquet::file::serialized_reader::ReadOptionsBuilder; use url::Url; // Tests are in submodules mod bloom_filter_io; mod filter_pushdown_io; +mod page_pruning_io; mod row_group_pruning_io; mod topk_io; @@ -44,20 +43,79 @@ struct ByteRangeDescriptor { description: String, } +/// Byte ranges of one column chunk: +/// +/// 1. (optional) dictionary page +/// 2. data pages (from the offset index, if present) +#[derive(Debug, Clone)] +struct ColumnChunkLayout { + col_name: String, + /// Row Group Index + rg_idx: usize, + /// The whole column chunk (dictionary page + data pages) + range: Range, + /// Byte range of the dictionary page + dictionary_page: Option>, + /// Byte range of each data page, in page order (empty without an offset index) + data_pages: Vec>, +} + +impl ColumnChunkLayout { + /// Describes `range` in terms of the pages of this column chunk it covers + fn describe(&self, range: &Range) -> String { + let overlaps = |r: &Range| range.start < r.end && range.end > r.start; + let mut parts = vec![]; + if self.dictionary_page.as_ref().is_some_and(overlaps) { + parts.push("dictionary page".to_string()); + } + let pages: Vec = self + .data_pages + .iter() + .enumerate() + .filter(|(_, r)| overlaps(r)) + .map(|(idx, _)| idx) + .collect(); + match pages.as_slice() { + [] if self.data_pages.is_empty() => parts.push("data pages".to_string()), + [] => {} + [page] => parts.push(format!("data page {page}")), + [first, .., last] if last - first + 1 == pages.len() => { + parts.push(format!("data pages {first}..={last}")) + } + pages => { + let pages: Vec = pages.iter().map(|p| p.to_string()).collect(); + parts.push(format!("data pages {}", pages.join(", "))) + } + } + format!( + "{} for column '{}' (rg {})", + parts.join(" and "), + self.col_name, + self.rg_idx + ) + } +} + #[derive(Debug, Clone, Default)] struct ParquetFileLayout { + /// Bloom filter and page index byte ranges descriptors: Vec, + column_chunks: Vec, } impl ParquetFileLayout { fn new(bytes: &Bytes) -> Self { let cursor = bytes.clone(); - let reader = SerializedFileReader::new(cursor).expect("valid parquet file"); + let options = ReadOptionsBuilder::new().with_page_index().build(); + let reader = SerializedFileReader::new_with_options(cursor, options) + .expect("valid parquet file"); let metadata = reader.metadata(); let mut descriptors = Vec::new(); + let mut column_chunks = Vec::new(); for (rg_idx, rg) in metadata.row_groups().iter().enumerate() { - for col in rg.columns() { + let page_index = metadata.page_index_for_row_group(rg_idx); + for (col_idx, col) in rg.columns().iter().enumerate() { let col_name = col.column_descr().name(); // Bloom filter byte range @@ -92,22 +150,47 @@ impl ParquetFileLayout { }); } - // Data pages / column chunk byte range - let data_start = col + // Dictionary page and data pages of the column chunk + let dictionary_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})" - ), + .and_then(|o| u64::try_from(o).ok()); + let data_pages: Vec> = page_index + .offset_index(col_idx) + .map(|offset_index| { + offset_index + .page_locations() + .iter() + .map(|location| { + let start = location.offset as u64; + start..start + location.compressed_page_size as u64 + }) + .collect() + }) + .unwrap_or_default(); + let dictionary_page = dictionary_start.map(|start| { + let end = data_pages + .first() + .map(|page| page.start) + .unwrap_or_else(|| col.data_page_offset() as u64); + start..end + }); + let chunk_start = + dictionary_start.unwrap_or_else(|| col.data_page_offset() as u64); + let chunk_end = chunk_start + col.compressed_size() as u64; + column_chunks.push(ColumnChunkLayout { + col_name: col_name.to_string(), + rg_idx, + range: chunk_start..chunk_end, + dictionary_page, + data_pages, }); } } - Self { descriptors } + Self { + descriptors, + column_chunks, + } } fn describe_range(&self, range: &Range, file_size: u64) -> String { @@ -122,6 +205,16 @@ impl ParquetFileLayout { ); } } + for chunk in &self.column_chunks { + if range.start < chunk.range.end && range.end > chunk.range.start { + return format!( + "{} (offsets {}..{})", + chunk.describe(range), + range.start, + range.end + ); + } + } format!("byte range (offsets {}..{})", range.start, range.end) } } @@ -330,32 +423,3 @@ impl ParquetIoTest { output } } - -/// Returns a Parquet file with three row groups of three rows each: -/// -/// | row group | id | val | -/// |-----------|---------|-----------| -/// | 0 | 1, 2, 3 | 0, 5, 10 | -/// | 1 | 4, 5, 6 | 0, 7, 10 | -/// | 2 | 7, 8, 9 | 0, 5, 10 | -/// -/// `id` is sorted, so its row group statistics are disjoint and predicates on -/// it can prune row groups. `val` has the same min/max (`0`/`10`) in every row -/// group, so statistics cannot prune row groups for predicates such as -/// `val = 7`, even though only row group 1 contains a match. -pub(crate) fn three_row_group_file() -> Bytes { - let id: ArrayRef = Arc::new(Int32Array::from_iter_values(1..=9)); - let val: ArrayRef = Arc::new(Int32Array::from(vec![0, 5, 10, 0, 7, 10, 0, 5, 10])); - let batch = RecordBatch::try_from_iter([("id", id), ("val", val)]).unwrap(); - - let props = WriterProperties::builder() - .set_max_row_group_row_count(Some(3)) - .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(); - Bytes::from(buffer) -} diff --git a/datafusion/core/tests/parquet/io/page_pruning_io.rs b/datafusion/core/tests/parquet/io/page_pruning_io.rs new file mode 100644 index 0000000000000..6d1b037dd66e0 --- /dev/null +++ b/datafusion/core/tests/parquet/io/page_pruning_io.rs @@ -0,0 +1,129 @@ +// 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 I/O events during Parquet page pruning using the page index +//! (`datafusion.execution.parquet.enable_page_index`, see +//! [`ParquetOptions::enable_page_index`]). +//! +//! The file has a single row group, so row group statistics can not prune +//! anything. Each column chunk is split into three data pages of three rows, +//! and `id` is sorted, so the page index (per page min/max of `id`) can +//! prune pages within the row group: +//! +//! ```text +//! row group 0 +//! column 'id' column 'val' +//! ┌──────────────┐ ┌──────────────┐ +//! page 0 │ 1, 2, 3 │ │ 0, 5, 10 │ rows 0..3 +//! ├──────────────┤ ├──────────────┤ +//! page 1 │ 4, 5, 6 │ │ 0, 7, 10 │ rows 3..6 +//! ├──────────────┤ ├──────────────┤ +//! page 2 │ 7, 8, 9 │ │ 0, 5, 10 │ rows 6..9 +//! └──────────────┘ └──────────────┘ +//! id page min/max: page 0: 1/3 page 1: 4/6 page 2: 7/9 +//! +//! id = 5 -> rows 3..6 -> page 1 of each column +//! id > 3 -> rows 3..9 -> pages 1 and 2 of each column +//! ``` +//! +//! [`ParquetOptions::enable_page_index`]: datafusion::config::ParquetOptions::enable_page_index + +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int32Array, RecordBatch}; +use bytes::Bytes; +use datafusion::prelude::SessionConfig; +use insta::assert_snapshot; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::WriterProperties; + +use super::ParquetIoTest; + +/// No predicate, so every page of every column is read as one column chunk. +#[tokio::test] +async fn test_page_pruning_io_no_predicate() { + let test = page_pruning_test(true).await; + assert_snapshot!(test.query("SELECT * FROM t").await, @r" + --- Query: SELECT * FROM t (9 rows) --- + Read Parquet footer (offsets 0..838) + Read data pages 0..=2 for column 'id' (rg 0) (offsets 4..91) + Read data pages 0..=2 for column 'val' (rg 0) (offsets 91..178) + Produce batch (9 rows) + "); +} + +/// Only page 1 of `id` (`4..=6`) can contain `id = 5`, so only page 1 of +/// each column is fetched. +#[tokio::test] +async fn test_page_pruning_io_one_page() { + let test = page_pruning_test(true).await; + assert_snapshot!(test.query("SELECT * FROM t WHERE id = 5").await, @r" + --- Query: SELECT * FROM t WHERE id = 5 (1 rows) --- + Read Parquet footer (offsets 0..838) + Read data page 1 for column 'id' (rg 0) (offsets 33..62) + Read data page 1 for column 'val' (rg 0) (offsets 120..149) + Produce batch (1 rows) + "); +} + +/// Page 0 of `id` (`1..=3`) cannot contain `id > 3`, so pages 1 and 2 of +/// each column are fetched. +#[tokio::test] +async fn test_page_pruning_io_two_pages() { + let test = page_pruning_test(true).await; + assert_snapshot!(test.query("SELECT * FROM t WHERE id > 3").await, @r" + --- Query: SELECT * FROM t WHERE id > 3 (6 rows) --- + Read Parquet footer (offsets 0..838) + Read data page 1 for column 'id' (rg 0) (offsets 33..62) + Read data page 2 for column 'id' (rg 0) (offsets 62..91) + Read data page 1 for column 'val' (rg 0) (offsets 120..149) + Read data page 2 for column 'val' (rg 0) (offsets 149..178) + Produce batch (6 rows) + "); +} + +/// With page index pruning disabled, the whole column chunk of every column +/// is read even though only page 1 can match. +#[tokio::test] +async fn test_page_pruning_io_disabled() { + let test = page_pruning_test(false).await; + assert_snapshot!(test.query("SELECT * FROM t WHERE id = 5").await, @r" + --- Query: SELECT * FROM t WHERE id = 5 (1 rows) --- + Read Parquet footer (offsets 0..838) + Read data pages 0..=2 for column 'id' (rg 0) (offsets 4..91) + Read data pages 0..=2 for column 'val' (rg 0) (offsets 91..178) + Produce batch (1 rows) + "); +} + +/// Returns a [`ParquetIoTest`] for the file described in the module docs +async fn page_pruning_test(page_index_pruning: bool) -> ParquetIoTest { + let config = SessionConfig::new() + .with_target_partitions(1) + .with_parquet_page_index_pruning(page_index_pruning); + ParquetIoTest::new(three_page_file(), config).await +} + +/// Returns a Parquet file with one row group of nine rows, where each column +/// chunk is written as three data pages of three rows each (see the module +/// docs). Dictionary encoding is disabled so each column chunk contains only +/// data pages. +fn three_page_file() -> Bytes { + let id: ArrayRef = Arc::new(Int32Array::from_iter_values(1..=9)); + let val: ArrayRef = Arc::new(Int32Array::from(vec![0, 5, 10, 0, 7, 10, 0, 5, 10])); + let batch = RecordBatch::try_from_iter([("id", id), ("val", val)]).unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_row_count_limit(3) + .set_write_batch_size(3) + .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(); + Bytes::from(buffer) +} diff --git a/datafusion/core/tests/parquet/io/row_group_pruning_io.rs b/datafusion/core/tests/parquet/io/row_group_pruning_io.rs index c1dd57ddc41aa..ef889a169a013 100644 --- a/datafusion/core/tests/parquet/io/row_group_pruning_io.rs +++ b/datafusion/core/tests/parquet/io/row_group_pruning_io.rs @@ -23,10 +23,16 @@ //! id > 100 -> no row groups //! ``` +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int32Array, RecordBatch}; +use bytes::Bytes; use datafusion::prelude::SessionConfig; use insta::assert_snapshot; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::WriterProperties; -use super::{ParquetIoTest, three_row_group_file}; +use super::ParquetIoTest; /// No predicate, so every row group is read. #[tokio::test] @@ -35,14 +41,14 @@ async fn test_row_group_pruning_io_no_predicate() { assert_snapshot!(test.query("SELECT * FROM t").await, @r" --- Query: SELECT * FROM t (9 rows) --- Read Parquet footer (offsets 0..1402) - Read data page for column 'id' (rg 0) (offsets 4..51) - Read data page for column 'val' (rg 0) (offsets 51..98) + Read dictionary page and data page 0 for column 'id' (rg 0) (offsets 4..51) + Read dictionary page and data page 0 for column 'val' (rg 0) (offsets 51..98) Produce batch (3 rows) - Read data page for column 'id' (rg 1) (offsets 98..145) - Read data page for column 'val' (rg 1) (offsets 145..192) + Read dictionary page and data page 0 for column 'id' (rg 1) (offsets 98..145) + Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) Produce batch (3 rows) - Read data page for column 'id' (rg 2) (offsets 192..239) - Read data page for column 'val' (rg 2) (offsets 239..286) + Read dictionary page and data page 0 for column 'id' (rg 2) (offsets 192..239) + Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) Produce batch (3 rows) "); } @@ -55,8 +61,8 @@ async fn test_row_group_pruning_io_one_row_group() { assert_snapshot!(test.query("SELECT * FROM t WHERE id = 5").await, @r" --- Query: SELECT * FROM t WHERE id = 5 (1 rows) --- Read Parquet footer (offsets 0..1402) - Read data page for column 'id' (rg 1) (offsets 98..145) - Read data page for column 'val' (rg 1) (offsets 145..192) + Read dictionary page and data page 0 for column 'id' (rg 1) (offsets 98..145) + Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) Produce batch (1 rows) "); } @@ -69,10 +75,10 @@ async fn test_row_group_pruning_io_two_row_groups() { assert_snapshot!(test.query("SELECT * FROM t WHERE id > 5").await, @r" --- Query: SELECT * FROM t WHERE id > 5 (4 rows) --- Read Parquet footer (offsets 0..1402) - Read data page for column 'id' (rg 1) (offsets 98..145) - Read data page for column 'val' (rg 1) (offsets 145..192) - Read data page for column 'id' (rg 2) (offsets 192..239) - Read data page for column 'val' (rg 2) (offsets 239..286) + Read dictionary page and data page 0 for column 'id' (rg 1) (offsets 98..145) + Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) + Read dictionary page and data page 0 for column 'id' (rg 2) (offsets 192..239) + Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) Produce batch (4 rows) "); } @@ -93,3 +99,22 @@ async fn row_group_pruning_test() -> ParquetIoTest { .with_parquet_pruning(true); ParquetIoTest::new(three_row_group_file(), config).await } + +/// Returns the Parquet file described in the module docs: three row groups of +/// three rows each, with `id` sorted across row groups. +fn three_row_group_file() -> Bytes { + let id: ArrayRef = Arc::new(Int32Array::from_iter_values(1..=9)); + let val: ArrayRef = Arc::new(Int32Array::from(vec![0, 5, 10, 0, 7, 10, 0, 5, 10])); + let batch = RecordBatch::try_from_iter([("id", id), ("val", val)]).unwrap(); + + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .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(); + Bytes::from(buffer) +} diff --git a/datafusion/core/tests/parquet/io/topk_io.rs b/datafusion/core/tests/parquet/io/topk_io.rs index 7d7e9da3bddc7..945f2ca3e2fbe 100644 --- a/datafusion/core/tests/parquet/io/topk_io.rs +++ b/datafusion/core/tests/parquet/io/topk_io.rs @@ -24,10 +24,16 @@ //! * the two rows returned by `ORDER BY id DESC LIMIT 2` //! ``` +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int32Array, RecordBatch}; +use bytes::Bytes; use datafusion::prelude::SessionConfig; use insta::assert_snapshot; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::WriterProperties; -use super::{ParquetIoTest, three_row_group_file}; +use super::ParquetIoTest; /// TopK dynamic filter pushdown prunes row groups during the scan. /// @@ -41,8 +47,8 @@ async fn test_topk_io_dynamic_filter() { assert_snapshot!(test.query("SELECT * FROM t ORDER BY id DESC LIMIT 2").await, @r" --- Query: SELECT * FROM t ORDER BY id DESC LIMIT 2 (2 rows) --- Read Parquet footer (offsets 0..1402) - Read data page for column 'id' (rg 2) (offsets 192..239) - Read data page for column 'val' (rg 2) (offsets 239..286) + Read dictionary page and data page 0 for column 'id' (rg 2) (offsets 192..239) + Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) Produce batch (2 rows) "); } @@ -55,12 +61,12 @@ async fn test_topk_io_no_dynamic_filter() { assert_snapshot!(test.query("SELECT * FROM t ORDER BY id DESC LIMIT 2").await, @r" --- Query: SELECT * FROM t ORDER BY id DESC LIMIT 2 (2 rows) --- Read Parquet footer (offsets 0..1402) - Read data page for column 'id' (rg 2) (offsets 192..239) - Read data page for column 'val' (rg 2) (offsets 239..286) - Read data page for column 'id' (rg 1) (offsets 98..145) - Read data page for column 'val' (rg 1) (offsets 145..192) - Read data page for column 'id' (rg 0) (offsets 4..51) - Read data page for column 'val' (rg 0) (offsets 51..98) + Read dictionary page and data page 0 for column 'id' (rg 2) (offsets 192..239) + Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) + Read dictionary page and data page 0 for column 'id' (rg 1) (offsets 98..145) + Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) + Read dictionary page and data page 0 for column 'id' (rg 0) (offsets 4..51) + Read dictionary page and data page 0 for column 'val' (rg 0) (offsets 51..98) Produce batch (2 rows) "); } @@ -73,3 +79,22 @@ async fn topk_test(dynamic_filter_pushdown: bool) -> ParquetIoTest { .enable_topk_dynamic_filter_pushdown = dynamic_filter_pushdown; ParquetIoTest::new(three_row_group_file(), config).await } + +/// Returns the Parquet file described in the module docs: three row groups of +/// three rows each, with `id` sorted across row groups. +fn three_row_group_file() -> Bytes { + let id: ArrayRef = Arc::new(Int32Array::from_iter_values(1..=9)); + let val: ArrayRef = Arc::new(Int32Array::from(vec![0, 5, 10, 0, 7, 10, 0, 5, 10])); + let batch = RecordBatch::try_from_iter([("id", id), ("val", val)]).unwrap(); + + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .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(); + Bytes::from(buffer) +} From 9a1bfd106acc7c2aa085d6abcb4ab3c2d4446196 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 4 Nov 2026 15:57:54 -0500 Subject: [PATCH 10/11] show requests --- .../core/tests/parquet/io/bloom_filter_io.rs | 29 ++++++++----- .../tests/parquet/io/filter_pushdown_io.rs | 43 ++++++++++--------- datafusion/core/tests/parquet/io/mod.rs | 28 +++++++++++- .../core/tests/parquet/io/page_pruning_io.rs | 36 ++++++++++------ .../tests/parquet/io/row_group_pruning_io.rs | 42 +++++++++++------- datafusion/core/tests/parquet/io/topk_io.rs | 26 ++++++----- 6 files changed, 131 insertions(+), 73 deletions(-) diff --git a/datafusion/core/tests/parquet/io/bloom_filter_io.rs b/datafusion/core/tests/parquet/io/bloom_filter_io.rs index b9b7f2213cf1d..01690a2fd9dd2 100644 --- a/datafusion/core/tests/parquet/io/bloom_filter_io.rs +++ b/datafusion/core/tests/parquet/io/bloom_filter_io.rs @@ -41,12 +41,17 @@ async fn test_parquet_bloom_filter_io_matching() { // 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 dictionary page and data page 0 for column 'a' (rg 0) (offsets 4..51) - Read dictionary page and data page 0 for column 'b' (rg 0) (offsets 51..98) + Request 1 + Read Parquet footer (offsets 0..678) + Request 2 + Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) + Request 3 + Read Bloom filter for column 'a' (rg 0) (offsets 98..145) + Request 4 + Read Bloom filter for column 'b' (rg 0) (offsets 145..192) + Request 5 + Read dictionary page and data page 0 for column 'a' (rg 0) (offsets 4..51) + Read dictionary page and data page 0 for column 'b' (rg 0) (offsets 51..98) Produce batch (1 rows) "); } @@ -61,10 +66,14 @@ async fn test_parquet_bloom_filter_io_pruned() { // 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) - 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) + Request 1 + Read Parquet footer (offsets 0..678) + Request 2 + Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) + Request 3 + Read Bloom filter for column 'a' (rg 0) (offsets 98..145) + Request 4 + Read Bloom filter for column 'b' (rg 0) (offsets 145..192) "); } diff --git a/datafusion/core/tests/parquet/io/filter_pushdown_io.rs b/datafusion/core/tests/parquet/io/filter_pushdown_io.rs index bcf2b3b925985..95d67cc4d4648 100644 --- a/datafusion/core/tests/parquet/io/filter_pushdown_io.rs +++ b/datafusion/core/tests/parquet/io/filter_pushdown_io.rs @@ -42,13 +42,17 @@ async fn test_filter_pushdown_io_disabled() { let test = filter_pushdown_test(false).await; assert_snapshot!(test.query("SELECT * FROM t WHERE val = 7").await, @r" --- Query: SELECT * FROM t WHERE val = 7 (1 rows) --- - Read Parquet footer (offsets 0..1402) - Read dictionary page and data page 0 for column 'id' (rg 0) (offsets 4..51) - Read dictionary page and data page 0 for column 'val' (rg 0) (offsets 51..98) - Read dictionary page and data page 0 for column 'id' (rg 1) (offsets 98..145) - Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) - Read dictionary page and data page 0 for column 'id' (rg 2) (offsets 192..239) - Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) + Request 1 + Read Parquet footer (offsets 0..1402) + Request 2 + Read dictionary page and data page 0 for column 'id' (rg 0) (offsets 4..51) + Read dictionary page and data page 0 for column 'val' (rg 0) (offsets 51..98) + Request 3 + Read dictionary page and data page 0 for column 'id' (rg 1) (offsets 98..145) + Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) + Request 4 + Read dictionary page and data page 0 for column 'id' (rg 2) (offsets 192..239) + Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) Produce batch (1 rows) "); } @@ -59,23 +63,20 @@ async fn test_filter_pushdown_io_disabled() { #[tokio::test] async fn test_filter_pushdown_io_enabled() { let test = filter_pushdown_test(true).await; - // Note: `id` for row group 1 is fetched as two ranges, its dictionary - // page and its data page, rather than as a single column chunk read. - // Once the filter has produced a row selection, the parquet reader uses - // the offset index to request only the pages needed for the selected - // rows (plus the dictionary page), one range per page, even when the - // ranges are adjacent. The object store's default `get_ranges` - // coalesces adjacent ranges into a single request, so this does not - // necessarily mean two requests to the storage system. assert_snapshot!(test.query("SELECT * FROM t WHERE val = 7").await, @r" --- Query: SELECT * FROM t WHERE val = 7 (1 rows) --- - Read Parquet footer (offsets 0..1402) - Read dictionary page and data page 0 for column 'val' (rg 0) (offsets 51..98) - Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) - Read dictionary page for column 'id' (rg 1) (offsets 98..124) - Read data page 0 for column 'id' (rg 1) (offsets 124..145) + Request 1 + Read Parquet footer (offsets 0..1402) + Request 2 + Read dictionary page and data page 0 for column 'val' (rg 0) (offsets 51..98) + Request 3 + Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) + Request 4 + Read dictionary page for column 'id' (rg 1) (offsets 98..124) + Read data page 0 for column 'id' (rg 1) (offsets 124..145) Produce batch (1 rows) - Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) + Request 5 + Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) "); } diff --git a/datafusion/core/tests/parquet/io/mod.rs b/datafusion/core/tests/parquet/io/mod.rs index f1308b4a8f06f..ea47a8b986d3d 100644 --- a/datafusion/core/tests/parquet/io/mod.rs +++ b/datafusion/core/tests/parquet/io/mod.rs @@ -13,6 +13,7 @@ use std::fmt::Write as _; use std::ops::Range; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use async_trait::async_trait; @@ -225,6 +226,8 @@ struct InstrumentedObjectStore { layout: ParquetFileLayout, file_size: u64, events: Arc>>, + /// Number of requests made to the store (reset for each query) + requests: Arc, } impl InstrumentedObjectStore { @@ -233,18 +236,29 @@ impl InstrumentedObjectStore { layout: ParquetFileLayout, file_size: u64, events: Arc>>, + requests: Arc, ) -> Self { Self { inner, layout, file_size, events, + requests, } } + + /// Records the start of a new request to the store + fn record_request(&self) { + let request = self.requests.fetch_add(1, Ordering::SeqCst) + 1; + self.events + .lock() + .unwrap() + .push(format!("Request {request}")); + } fn record_range(&self, range: &Range) { let desc = self.layout.describe_range(range, self.file_size); - self.events.lock().unwrap().push(format!("Read {desc}")); + self.events.lock().unwrap().push(format!(" Read {desc}")); } } @@ -278,6 +292,7 @@ impl ObjectStore for InstrumentedObjectStore { location: &Path, options: GetOptions, ) -> object_store::Result { + self.record_request(); if let Some(range) = &options.range { match range { GetRange::Bounded(r) => { @@ -302,6 +317,7 @@ impl ObjectStore for InstrumentedObjectStore { location: &Path, ranges: &[Range], ) -> object_store::Result> { + self.record_request(); for range in ranges { self.record_range(range); } @@ -359,6 +375,7 @@ impl ObjectStore for InstrumentedObjectStore { pub(crate) struct ParquetIoTest { ctx: SessionContext, events: Arc>>, + requests: Arc, } impl ParquetIoTest { @@ -379,11 +396,13 @@ impl ParquetIoTest { .unwrap(); let events = Arc::new(Mutex::new(Vec::new())); + let requests = Arc::new(AtomicUsize::new(0)); let instrumented_store = Arc::new(InstrumentedObjectStore::new( memory_store, layout, file_size, Arc::clone(&events), + Arc::clone(&requests), )); let ctx = SessionContext::new_with_config(config); @@ -392,7 +411,11 @@ impl ParquetIoTest { .await .unwrap(); - Self { ctx, events } + Self { + ctx, + events, + requests, + } } /// Runs `sql` and returns the storage requests and output batches it @@ -401,6 +424,7 @@ impl ParquetIoTest { /// Only events recorded while running this query are included. pub(crate) async fn query(&self, sql: &str) -> String { self.events.lock().unwrap().clear(); + self.requests.store(0, Ordering::SeqCst); let df = self.ctx.sql(sql).await.unwrap(); let mut stream = df.execute_stream().await.unwrap(); diff --git a/datafusion/core/tests/parquet/io/page_pruning_io.rs b/datafusion/core/tests/parquet/io/page_pruning_io.rs index 6d1b037dd66e0..8a5507acc22f2 100644 --- a/datafusion/core/tests/parquet/io/page_pruning_io.rs +++ b/datafusion/core/tests/parquet/io/page_pruning_io.rs @@ -46,9 +46,11 @@ async fn test_page_pruning_io_no_predicate() { let test = page_pruning_test(true).await; assert_snapshot!(test.query("SELECT * FROM t").await, @r" --- Query: SELECT * FROM t (9 rows) --- - Read Parquet footer (offsets 0..838) - Read data pages 0..=2 for column 'id' (rg 0) (offsets 4..91) - Read data pages 0..=2 for column 'val' (rg 0) (offsets 91..178) + Request 1 + Read Parquet footer (offsets 0..838) + Request 2 + Read data pages 0..=2 for column 'id' (rg 0) (offsets 4..91) + Read data pages 0..=2 for column 'val' (rg 0) (offsets 91..178) Produce batch (9 rows) "); } @@ -60,9 +62,11 @@ async fn test_page_pruning_io_one_page() { let test = page_pruning_test(true).await; assert_snapshot!(test.query("SELECT * FROM t WHERE id = 5").await, @r" --- Query: SELECT * FROM t WHERE id = 5 (1 rows) --- - Read Parquet footer (offsets 0..838) - Read data page 1 for column 'id' (rg 0) (offsets 33..62) - Read data page 1 for column 'val' (rg 0) (offsets 120..149) + Request 1 + Read Parquet footer (offsets 0..838) + Request 2 + Read data page 1 for column 'id' (rg 0) (offsets 33..62) + Read data page 1 for column 'val' (rg 0) (offsets 120..149) Produce batch (1 rows) "); } @@ -74,11 +78,13 @@ async fn test_page_pruning_io_two_pages() { let test = page_pruning_test(true).await; assert_snapshot!(test.query("SELECT * FROM t WHERE id > 3").await, @r" --- Query: SELECT * FROM t WHERE id > 3 (6 rows) --- - Read Parquet footer (offsets 0..838) - Read data page 1 for column 'id' (rg 0) (offsets 33..62) - Read data page 2 for column 'id' (rg 0) (offsets 62..91) - Read data page 1 for column 'val' (rg 0) (offsets 120..149) - Read data page 2 for column 'val' (rg 0) (offsets 149..178) + Request 1 + Read Parquet footer (offsets 0..838) + Request 2 + Read data page 1 for column 'id' (rg 0) (offsets 33..62) + Read data page 2 for column 'id' (rg 0) (offsets 62..91) + Read data page 1 for column 'val' (rg 0) (offsets 120..149) + Read data page 2 for column 'val' (rg 0) (offsets 149..178) Produce batch (6 rows) "); } @@ -90,9 +96,11 @@ async fn test_page_pruning_io_disabled() { let test = page_pruning_test(false).await; assert_snapshot!(test.query("SELECT * FROM t WHERE id = 5").await, @r" --- Query: SELECT * FROM t WHERE id = 5 (1 rows) --- - Read Parquet footer (offsets 0..838) - Read data pages 0..=2 for column 'id' (rg 0) (offsets 4..91) - Read data pages 0..=2 for column 'val' (rg 0) (offsets 91..178) + Request 1 + Read Parquet footer (offsets 0..838) + Request 2 + Read data pages 0..=2 for column 'id' (rg 0) (offsets 4..91) + Read data pages 0..=2 for column 'val' (rg 0) (offsets 91..178) Produce batch (1 rows) "); } diff --git a/datafusion/core/tests/parquet/io/row_group_pruning_io.rs b/datafusion/core/tests/parquet/io/row_group_pruning_io.rs index ef889a169a013..02d72d2be5663 100644 --- a/datafusion/core/tests/parquet/io/row_group_pruning_io.rs +++ b/datafusion/core/tests/parquet/io/row_group_pruning_io.rs @@ -40,15 +40,19 @@ async fn test_row_group_pruning_io_no_predicate() { let test = row_group_pruning_test().await; assert_snapshot!(test.query("SELECT * FROM t").await, @r" --- Query: SELECT * FROM t (9 rows) --- - Read Parquet footer (offsets 0..1402) - Read dictionary page and data page 0 for column 'id' (rg 0) (offsets 4..51) - Read dictionary page and data page 0 for column 'val' (rg 0) (offsets 51..98) + Request 1 + Read Parquet footer (offsets 0..1402) + Request 2 + Read dictionary page and data page 0 for column 'id' (rg 0) (offsets 4..51) + Read dictionary page and data page 0 for column 'val' (rg 0) (offsets 51..98) Produce batch (3 rows) - Read dictionary page and data page 0 for column 'id' (rg 1) (offsets 98..145) - Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) + Request 3 + Read dictionary page and data page 0 for column 'id' (rg 1) (offsets 98..145) + Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) Produce batch (3 rows) - Read dictionary page and data page 0 for column 'id' (rg 2) (offsets 192..239) - Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) + Request 4 + Read dictionary page and data page 0 for column 'id' (rg 2) (offsets 192..239) + Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) Produce batch (3 rows) "); } @@ -60,9 +64,11 @@ async fn test_row_group_pruning_io_one_row_group() { let test = row_group_pruning_test().await; assert_snapshot!(test.query("SELECT * FROM t WHERE id = 5").await, @r" --- Query: SELECT * FROM t WHERE id = 5 (1 rows) --- - Read Parquet footer (offsets 0..1402) - Read dictionary page and data page 0 for column 'id' (rg 1) (offsets 98..145) - Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) + Request 1 + Read Parquet footer (offsets 0..1402) + Request 2 + Read dictionary page and data page 0 for column 'id' (rg 1) (offsets 98..145) + Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) Produce batch (1 rows) "); } @@ -74,11 +80,14 @@ async fn test_row_group_pruning_io_two_row_groups() { let test = row_group_pruning_test().await; assert_snapshot!(test.query("SELECT * FROM t WHERE id > 5").await, @r" --- Query: SELECT * FROM t WHERE id > 5 (4 rows) --- - Read Parquet footer (offsets 0..1402) - Read dictionary page and data page 0 for column 'id' (rg 1) (offsets 98..145) - Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) - Read dictionary page and data page 0 for column 'id' (rg 2) (offsets 192..239) - Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) + Request 1 + Read Parquet footer (offsets 0..1402) + Request 2 + Read dictionary page and data page 0 for column 'id' (rg 1) (offsets 98..145) + Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) + Request 3 + Read dictionary page and data page 0 for column 'id' (rg 2) (offsets 192..239) + Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) Produce batch (4 rows) "); } @@ -89,7 +98,8 @@ async fn test_row_group_pruning_io_all_pruned() { let test = row_group_pruning_test().await; assert_snapshot!(test.query("SELECT * FROM t WHERE id > 100").await, @r" --- Query: SELECT * FROM t WHERE id > 100 (0 rows) --- - Read Parquet footer (offsets 0..1402) + Request 1 + Read Parquet footer (offsets 0..1402) "); } diff --git a/datafusion/core/tests/parquet/io/topk_io.rs b/datafusion/core/tests/parquet/io/topk_io.rs index 945f2ca3e2fbe..ce8175116dd31 100644 --- a/datafusion/core/tests/parquet/io/topk_io.rs +++ b/datafusion/core/tests/parquet/io/topk_io.rs @@ -46,9 +46,11 @@ async fn test_topk_io_dynamic_filter() { let test = topk_test(true).await; assert_snapshot!(test.query("SELECT * FROM t ORDER BY id DESC LIMIT 2").await, @r" --- Query: SELECT * FROM t ORDER BY id DESC LIMIT 2 (2 rows) --- - Read Parquet footer (offsets 0..1402) - Read dictionary page and data page 0 for column 'id' (rg 2) (offsets 192..239) - Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) + Request 1 + Read Parquet footer (offsets 0..1402) + Request 2 + Read dictionary page and data page 0 for column 'id' (rg 2) (offsets 192..239) + Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) Produce batch (2 rows) "); } @@ -60,13 +62,17 @@ async fn test_topk_io_no_dynamic_filter() { let test = topk_test(false).await; assert_snapshot!(test.query("SELECT * FROM t ORDER BY id DESC LIMIT 2").await, @r" --- Query: SELECT * FROM t ORDER BY id DESC LIMIT 2 (2 rows) --- - Read Parquet footer (offsets 0..1402) - Read dictionary page and data page 0 for column 'id' (rg 2) (offsets 192..239) - Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) - Read dictionary page and data page 0 for column 'id' (rg 1) (offsets 98..145) - Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) - Read dictionary page and data page 0 for column 'id' (rg 0) (offsets 4..51) - Read dictionary page and data page 0 for column 'val' (rg 0) (offsets 51..98) + Request 1 + Read Parquet footer (offsets 0..1402) + Request 2 + Read dictionary page and data page 0 for column 'id' (rg 2) (offsets 192..239) + Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) + Request 3 + Read dictionary page and data page 0 for column 'id' (rg 1) (offsets 98..145) + Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) + Request 4 + Read dictionary page and data page 0 for column 'id' (rg 0) (offsets 4..51) + Read dictionary page and data page 0 for column 'val' (rg 0) (offsets 51..98) Produce batch (2 rows) "); } From e04f166b238e3ea163757f303ee32ca24b73d1f2 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 4 Nov 2026 15:59:19 -0500 Subject: [PATCH 11/11] ft --- datafusion/core/tests/parquet/io/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/core/tests/parquet/io/mod.rs b/datafusion/core/tests/parquet/io/mod.rs index ea47a8b986d3d..edf6bd9c34771 100644 --- a/datafusion/core/tests/parquet/io/mod.rs +++ b/datafusion/core/tests/parquet/io/mod.rs @@ -246,7 +246,7 @@ impl InstrumentedObjectStore { requests, } } - + /// Records the start of a new request to the store fn record_request(&self) { let request = self.requests.fetch_add(1, Ordering::SeqCst) + 1;