From c0a70fb733f81d0165b5efb05a72f98e451eb212 Mon Sep 17 00:00:00 2001 From: mohitgurav20 Date: Fri, 30 Oct 2026 00:20:25 +0530 Subject: [PATCH 01/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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; From 895173ff9bcf14349d8717291f342002c234ec32 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 4 Nov 2026 19:01:06 -0300 Subject: [PATCH 12/14] prototype: capture parquet I/O through tracing events The I/O snapshot tests only saw object store requests, so they could not say why a byte range was read. This adds tracing::debug! events with structured fields (offset_start, offset_end, column, row_group) at the parquet read sites, and a test-only tracing Layer that renders the events and the object store's request spans as a snapshot. Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 3 + Cargo.toml | 2 + datafusion/core/Cargo.toml | 2 + .../core/tests/parquet/io/bloom_filter_io.rs | 24 +- .../tests/parquet/io/filter_pushdown_io.rs | 23 +- datafusion/core/tests/parquet/io/mod.rs | 147 ++++++------ .../core/tests/parquet/io/page_pruning_io.rs | 34 ++- .../tests/parquet/io/row_group_pruning_io.rs | 36 ++- datafusion/core/tests/parquet/io/topk_io.rs | 24 +- .../core/tests/parquet/io/trace_capture.rs | 219 ++++++++++++++++++ datafusion/datasource-parquet/Cargo.toml | 1 + .../datasource-parquet/src/file_format.rs | 6 + datafusion/datasource-parquet/src/metadata.rs | 25 ++ .../datasource-parquet/src/opener/mod.rs | 13 ++ .../datasource-parquet/src/push_decoder.rs | 8 + 15 files changed, 472 insertions(+), 95 deletions(-) create mode 100644 datafusion/core/tests/parquet/io/trace_capture.rs diff --git a/Cargo.lock b/Cargo.lock index 6fd56360c54f2..d86d4a6a0dcb8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1788,6 +1788,8 @@ dependencies = [ "tempfile", "test-utils", "tokio", + "tracing", + "tracing-subscriber", "unicode-width 0.2.3", "url", "uuid", @@ -2108,6 +2110,7 @@ dependencies = [ "parquet", "tempfile", "tokio", + "tracing", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index a96d3cbf37e58..9eca5e2bdcf70 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -163,6 +163,8 @@ itertools = "0.15" itoa = "1.0" liblzma = { version = "0.4.6", features = ["static"] } log = "^0.4" +tracing = { version = "0.1.44", default-features = false, features = ["std"] } +tracing-subscriber = { version = "0.3.23", default-features = false, features = ["registry", "std"] } loom = "0.7" memchr = "2.8.1" num-traits = { version = "0.2" } diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 322a6a353fb93..679c058865c53 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -167,6 +167,8 @@ bytes = { workspace = true } env_logger = { workspace = true } glob = { workspace = true } insta = { workspace = true } +tracing = { workspace = true } +tracing-subscriber = { workspace = true } half = { workspace = true } rand = { workspace = true, features = ["small_rng"] } rand_distr = "0.5" diff --git a/datafusion/core/tests/parquet/io/bloom_filter_io.rs b/datafusion/core/tests/parquet/io/bloom_filter_io.rs index 01690a2fd9dd2..ba9f128403876 100644 --- a/datafusion/core/tests/parquet/io/bloom_filter_io.rs +++ b/datafusion/core/tests/parquet/io/bloom_filter_io.rs @@ -37,18 +37,27 @@ use super::ParquetIoTest; #[tokio::test] async fn test_parquet_bloom_filter_io_matching() { let test = bloom_filter_test().await; + let traces = test.capture_traces(); // 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" + test.query("SELECT * FROM t WHERE a = 10 AND b = 100") + .await + .unwrap(); + assert_snapshot!(traces.render(), @" --- Query: SELECT * FROM t WHERE a = 10 AND b = 100 (1 rows) --- Request 1 - Read Parquet footer (offsets 0..678) + HEAD object metadata + Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) Request 2 Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) + Read Bloom filter for column 'a' (rg 0) (offsets 98..145) Request 3 Read Bloom filter for column 'a' (rg 0) (offsets 98..145) + Read Bloom filter for column 'b' (rg 0) (offsets 145..192) Request 4 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 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) @@ -60,18 +69,25 @@ async fn test_parquet_bloom_filter_io_matching() { #[tokio::test] async fn test_parquet_bloom_filter_io_pruned() { let test = bloom_filter_test().await; + let traces = test.capture_traces(); // 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" + test.query("SELECT * FROM t WHERE a = 999 AND b = 100") + .await + .unwrap(); + assert_snapshot!(traces.render(), @" --- Query: SELECT * FROM t WHERE a = 999 AND b = 100 (0 rows) --- Request 1 - Read Parquet footer (offsets 0..678) + HEAD object metadata + Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) Request 2 Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) + Read Bloom filter for column 'a' (rg 0) (offsets 98..145) Request 3 Read Bloom filter for column 'a' (rg 0) (offsets 98..145) + Read Bloom filter for column 'b' (rg 0) (offsets 145..192) 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 95d67cc4d4648..3b12ad00dc975 100644 --- a/datafusion/core/tests/parquet/io/filter_pushdown_io.rs +++ b/datafusion/core/tests/parquet/io/filter_pushdown_io.rs @@ -40,16 +40,24 @@ use super::ParquetIoTest; #[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" + let traces = test.capture_traces(); + test.query("SELECT * FROM t WHERE val = 7").await.unwrap(); + assert_snapshot!(traces.render(), @" --- Query: SELECT * FROM t WHERE val = 7 (1 rows) --- Request 1 - Read Parquet footer (offsets 0..1402) + HEAD object metadata + 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 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) + 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) + 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) @@ -63,18 +71,25 @@ async fn test_filter_pushdown_io_disabled() { #[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" + let traces = test.capture_traces(); + test.query("SELECT * FROM t WHERE val = 7").await.unwrap(); + assert_snapshot!(traces.render(), @" --- Query: SELECT * FROM t WHERE val = 7 (1 rows) --- Request 1 - Read Parquet footer (offsets 0..1402) + HEAD object metadata + Read dictionary page and data page 0 for column 'val' (rg 0) (offsets 51..98) Request 2 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) Request 3 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 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 edf6bd9c34771..4a66c342bdae5 100644 --- a/datafusion/core/tests/parquet/io/mod.rs +++ b/datafusion/core/tests/parquet/io/mod.rs @@ -11,13 +11,13 @@ //! `ListingTable`object store access tests live in //! [`datafusion/core/tests/datasource/object_store_access.rs`](https://github.com/apache/datafusion/blob/main/datafusion/core/tests/datasource/object_store_access.rs). -use std::fmt::Write as _; use std::ops::Range; +use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::{Arc, Mutex}; use async_trait::async_trait; use bytes::Bytes; +use datafusion::error::Result; use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; use futures::StreamExt; use futures::stream::BoxStream; @@ -37,6 +37,9 @@ mod filter_pushdown_io; mod page_pruning_io; mod row_group_pruning_io; mod topk_io; +mod trace_capture; + +use trace_capture::TraceCapture; #[derive(Debug, Clone)] struct ByteRangeDescriptor { @@ -223,9 +226,7 @@ impl ParquetFileLayout { #[derive(Debug)] struct InstrumentedObjectStore { inner: Arc, - layout: ParquetFileLayout, file_size: u64, - events: Arc>>, /// Number of requests made to the store (reset for each query) requests: Arc, } @@ -233,32 +234,48 @@ struct InstrumentedObjectStore { impl InstrumentedObjectStore { fn new( inner: Arc, - 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) { + /// Emits a span for a `HEAD` request, which reads no bytes. + fn record_head_request(&self) { let request = self.requests.fetch_add(1, Ordering::SeqCst) + 1; - self.events - .lock() - .unwrap() - .push(format!("Request {request}")); + let span = tracing::debug_span!( + target: "datafusion::test::object_store", + "object_store_request", + request + ); + span.in_scope(|| { + tracing::debug!(target: "datafusion::test::object_store", "HEAD object metadata"); + }); } - fn record_range(&self, range: &Range) { - let desc = self.layout.describe_range(range, self.file_size); - self.events.lock().unwrap().push(format!(" Read {desc}")); + /// Emits a span for one request to the store, with an event per byte + /// range it reads. + fn record_request(&self, ranges: impl IntoIterator>) { + let request = self.requests.fetch_add(1, Ordering::SeqCst) + 1; + let span = tracing::debug_span!( + target: "datafusion::test::object_store", + "object_store_request", + request + ); + span.in_scope(|| { + for range in ranges { + tracing::debug!( + target: "datafusion::test::object_store", + offset_start = range.start, + offset_end = range.end, + "Read byte range" + ); + } + }); } } @@ -292,23 +309,19 @@ 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) => { - 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)); + if options.head { + self.record_head_request(); + return self.inner.get_opts(location, options).await; } + let range = match &options.range { + Some(GetRange::Bounded(r)) => r.clone(), + Some(GetRange::Offset(offset)) => *offset..self.file_size, + Some(GetRange::Suffix(suffix)) => { + self.file_size.saturating_sub(*suffix)..self.file_size + } + None => 0..self.file_size, + }; + self.record_request([range]); self.inner.get_opts(location, options).await } @@ -317,10 +330,7 @@ impl ObjectStore for InstrumentedObjectStore { location: &Path, ranges: &[Range], ) -> object_store::Result> { - self.record_request(); - for range in ranges { - self.record_range(range); - } + self.record_request(ranges.iter().cloned()); self.inner.get_ranges(location, ranges).await } @@ -364,17 +374,21 @@ impl ObjectStore for InstrumentedObjectStore { } /// Test fixture: a single Parquet file registered as table `t`, read through an -/// [`InstrumentedObjectStore`] that records every storage request. +/// [`InstrumentedObjectStore`] that emits a `tracing` span for every storage +/// request. /// /// Typical use: /// /// ```text /// let test = ParquetIoTest::new(file_bytes, SessionConfig::new()).await; -/// assert_snapshot!(test.query("SELECT * FROM t WHERE ...").await, @"..."); +/// let traces = test.capture_traces(); +/// test.query("SELECT * FROM t WHERE ...").await.unwrap(); +/// assert_snapshot!(traces.render(), @"..."); /// ``` pub(crate) struct ParquetIoTest { ctx: SessionContext, - events: Arc>>, + layout: ParquetFileLayout, + file_size: u64, requests: Arc, } @@ -395,13 +409,10 @@ impl ParquetIoTest { .await .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), )); @@ -413,37 +424,39 @@ impl ParquetIoTest { Self { ctx, - events, + layout, + file_size, requests, } } - /// 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(); - self.requests.store(0, Ordering::SeqCst); - - let df = self.ctx.sql(sql).await.unwrap(); - let mut stream = df.execute_stream().await.unwrap(); + /// Captures the spans and events that DataFusion emits on this thread + /// until the returned [`TraceCapture`] is dropped. + pub(crate) fn capture_traces(&self) -> TraceCapture { + TraceCapture::new(self.layout.clone(), self.file_size) + } + /// Runs `sql` to completion inside a `query` span and returns the number + /// of output rows. Emits a `Produce batch` event for each output batch. + pub(crate) async fn query(&self, sql: &str) -> Result { + self.requests.store(0, Ordering::SeqCst); + let span = tracing::debug_span!( + target: "datafusion::test", + "query", + sql, + rows = tracing::field::Empty + ); + let _entered = span.enter(); + + let df = self.ctx.sql(sql).await?; + let mut stream = df.execute_stream().await?; 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}"); + let rows = batch?.num_rows(); + row_count += rows; + tracing::debug!(target: "datafusion::test", rows, "Produce batch"); } - output + span.record("rows", row_count); + Ok(row_count) } } diff --git a/datafusion/core/tests/parquet/io/page_pruning_io.rs b/datafusion/core/tests/parquet/io/page_pruning_io.rs index 8a5507acc22f2..fa2bb6dfba8a0 100644 --- a/datafusion/core/tests/parquet/io/page_pruning_io.rs +++ b/datafusion/core/tests/parquet/io/page_pruning_io.rs @@ -44,10 +44,14 @@ use super::ParquetIoTest; #[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" + let traces = test.capture_traces(); + test.query("SELECT * FROM t").await.unwrap(); + assert_snapshot!(traces.render(), @" --- Query: SELECT * FROM t (9 rows) --- Request 1 - Read Parquet footer (offsets 0..838) + HEAD object metadata + 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 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) @@ -60,10 +64,14 @@ async fn test_page_pruning_io_no_predicate() { #[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" + let traces = test.capture_traces(); + test.query("SELECT * FROM t WHERE id = 5").await.unwrap(); + assert_snapshot!(traces.render(), @" --- Query: SELECT * FROM t WHERE id = 5 (1 rows) --- Request 1 - Read Parquet footer (offsets 0..838) + HEAD object metadata + 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 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) @@ -76,10 +84,16 @@ async fn test_page_pruning_io_one_page() { #[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" + let traces = test.capture_traces(); + test.query("SELECT * FROM t WHERE id > 3").await.unwrap(); + assert_snapshot!(traces.render(), @" --- Query: SELECT * FROM t WHERE id > 3 (6 rows) --- Request 1 - Read Parquet footer (offsets 0..838) + HEAD object metadata + 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 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) @@ -94,10 +108,14 @@ async fn test_page_pruning_io_two_pages() { #[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" + let traces = test.capture_traces(); + test.query("SELECT * FROM t WHERE id = 5").await.unwrap(); + assert_snapshot!(traces.render(), @" --- Query: SELECT * FROM t WHERE id = 5 (1 rows) --- Request 1 - Read Parquet footer (offsets 0..838) + HEAD object metadata + 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 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) 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 02d72d2be5663..aa58324270e79 100644 --- a/datafusion/core/tests/parquet/io/row_group_pruning_io.rs +++ b/datafusion/core/tests/parquet/io/row_group_pruning_io.rs @@ -38,18 +38,26 @@ use super::ParquetIoTest; #[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" + let traces = test.capture_traces(); + test.query("SELECT * FROM t").await.unwrap(); + assert_snapshot!(traces.render(), @" --- Query: SELECT * FROM t (9 rows) --- Request 1 - Read Parquet footer (offsets 0..1402) + HEAD object metadata + 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 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) @@ -62,10 +70,14 @@ async fn test_row_group_pruning_io_no_predicate() { #[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" + let traces = test.capture_traces(); + test.query("SELECT * FROM t WHERE id = 5").await.unwrap(); + assert_snapshot!(traces.render(), @" --- Query: SELECT * FROM t WHERE id = 5 (1 rows) --- Request 1 - Read Parquet footer (offsets 0..1402) + HEAD object metadata + 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 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) @@ -78,13 +90,19 @@ async fn test_row_group_pruning_io_one_row_group() { #[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" + let traces = test.capture_traces(); + test.query("SELECT * FROM t WHERE id > 5").await.unwrap(); + assert_snapshot!(traces.render(), @" --- Query: SELECT * FROM t WHERE id > 5 (4 rows) --- Request 1 - Read Parquet footer (offsets 0..1402) + HEAD object metadata + 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 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) + 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 2) (offsets 192..239) Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) @@ -96,10 +114,12 @@ async fn test_row_group_pruning_io_two_row_groups() { #[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" + let traces = test.capture_traces(); + test.query("SELECT * FROM t WHERE id > 100").await.unwrap(); + assert_snapshot!(traces.render(), @" --- Query: SELECT * FROM t WHERE id > 100 (0 rows) --- Request 1 - Read Parquet footer (offsets 0..1402) + HEAD object metadata "); } diff --git a/datafusion/core/tests/parquet/io/topk_io.rs b/datafusion/core/tests/parquet/io/topk_io.rs index ce8175116dd31..97e2d8077ecc8 100644 --- a/datafusion/core/tests/parquet/io/topk_io.rs +++ b/datafusion/core/tests/parquet/io/topk_io.rs @@ -44,10 +44,16 @@ use super::ParquetIoTest; #[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" + let traces = test.capture_traces(); + test.query("SELECT * FROM t ORDER BY id DESC LIMIT 2") + .await + .unwrap(); + assert_snapshot!(traces.render(), @" --- Query: SELECT * FROM t ORDER BY id DESC LIMIT 2 (2 rows) --- Request 1 - Read Parquet footer (offsets 0..1402) + HEAD object metadata + 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 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) @@ -60,16 +66,26 @@ async fn test_topk_io_dynamic_filter() { #[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" + let traces = test.capture_traces(); + test.query("SELECT * FROM t ORDER BY id DESC LIMIT 2") + .await + .unwrap(); + assert_snapshot!(traces.render(), @" --- Query: SELECT * FROM t ORDER BY id DESC LIMIT 2 (2 rows) --- Request 1 - Read Parquet footer (offsets 0..1402) + HEAD object metadata + 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 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) + 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) + 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 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) diff --git a/datafusion/core/tests/parquet/io/trace_capture.rs b/datafusion/core/tests/parquet/io/trace_capture.rs new file mode 100644 index 0000000000000..e91005b73874d --- /dev/null +++ b/datafusion/core/tests/parquet/io/trace_capture.rs @@ -0,0 +1,219 @@ +// 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 + +//! Captures `tracing` spans and events emitted by DataFusion during a test and +//! renders them as text for `insta` snapshots. +//! +//! The library emits events with a static message and structured fields +//! (`offset_start`, `offset_end`, `column`, `row_group`, ...). This module owns +//! the human-readable rendering, so the library never formats strings for +//! tests. +//! +//! The subscriber is installed with [`tracing::subscriber::set_default`], which +//! is thread-local. `#[tokio::test]` uses a current-thread runtime, so every +//! task of the query runs on the test thread and reports to this subscriber. + +use std::collections::BTreeMap; +use std::fmt::{self, Write as _}; +use std::sync::{Arc, Mutex}; + +use tracing::field::{Field, Visit}; +use tracing::span::{Attributes, Id, Record as SpanRecord}; +use tracing::subscriber::DefaultGuard; +use tracing::{Event, Subscriber}; +use tracing_subscriber::layer::{Context, Layer, SubscriberExt}; +use tracing_subscriber::registry::{LookupSpan, Registry}; + +use super::ParquetFileLayout; + +/// Only spans and events from these target prefixes are captured. Dependencies +/// such as `h2` or `object_store` can emit their own events. +const TARGET_PREFIX: &str = "datafusion"; + +/// One captured span (rendered when it opens) or event. +#[derive(Debug, Clone)] +struct Record { + /// Nesting level below the outermost captured span + depth: usize, + /// The span name, or the event's `message` field + name: String, + fields: BTreeMap, +} + +#[derive(Default)] +struct FieldVisitor(BTreeMap); + +impl Visit for FieldVisitor { + fn record_str(&mut self, field: &Field, value: &str) { + self.0.insert(field.name().to_string(), value.to_string()); + } + + fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) { + self.0 + .insert(field.name().to_string(), format!("{value:?}")); + } +} + +/// Index into [`CaptureLayer::records`] for a span, plus the depth its children +/// get. Stored in the span's extensions. +struct SpanSlot { + record: usize, + child_depth: usize, +} + +struct CaptureLayer { + records: Arc>>, +} + +impl CaptureLayer { + fn parent_depth(ctx: &Context<'_, S>, parent: Option) -> usize + where + S: Subscriber + for<'a> LookupSpan<'a>, + { + parent + .and_then(|id| ctx.span(&id)) + .and_then(|span| span.extensions().get::().map(|s| s.child_depth)) + .unwrap_or(0) + } +} + +impl Layer for CaptureLayer +where + S: Subscriber + for<'a> LookupSpan<'a>, +{ + fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) { + let metadata = attrs.metadata(); + if !metadata.target().starts_with(TARGET_PREFIX) { + return; + } + let parent = if attrs.is_contextual() { + ctx.current_span().id().cloned() + } else { + attrs.parent().cloned() + }; + let depth = Self::parent_depth(&ctx, parent); + let mut visitor = FieldVisitor::default(); + attrs.record(&mut visitor); + + let mut records = self.records.lock().unwrap(); + records.push(Record { + depth, + name: metadata.name().to_string(), + fields: visitor.0, + }); + // The query span is the frame of the snapshot, so its children are not + // indented. + let child_depth = if metadata.name() == "query" { + depth + } else { + depth + 1 + }; + let span = ctx.span(id).expect("span exists"); + span.extensions_mut().insert(SpanSlot { + record: records.len() - 1, + child_depth, + }); + } + + fn on_record(&self, id: &Id, values: &SpanRecord<'_>, ctx: Context<'_, S>) { + let Some(span) = ctx.span(id) else { return }; + let extensions = span.extensions(); + let Some(slot) = extensions.get::() else { + return; + }; + let mut visitor = FieldVisitor::default(); + values.record(&mut visitor); + self.records.lock().unwrap()[slot.record] + .fields + .extend(visitor.0); + } + + fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) { + if !event.metadata().target().starts_with(TARGET_PREFIX) { + return; + } + let depth = Self::parent_depth(&ctx, ctx.event_span(event).map(|s| s.id())); + let mut visitor = FieldVisitor::default(); + event.record(&mut visitor); + let mut fields = visitor.0; + let name = fields.remove("message").unwrap_or_default(); + self.records.lock().unwrap().push(Record { + depth, + name, + fields, + }); + } +} + +/// Spans and events captured on this thread until it is dropped. +pub(crate) struct TraceCapture { + records: Arc>>, + layout: ParquetFileLayout, + file_size: u64, + _guard: DefaultGuard, +} + +impl TraceCapture { + pub(super) fn new(layout: ParquetFileLayout, file_size: u64) -> Self { + let records = Arc::new(Mutex::new(Vec::new())); + let subscriber = Registry::default().with(CaptureLayer { + records: Arc::clone(&records), + }); + let guard = tracing::subscriber::set_default(subscriber); + Self { + records, + layout, + file_size, + _guard: guard, + } + } + + /// Renders and clears everything captured so far. + pub(crate) fn render(&self) -> String { + let records = std::mem::take(&mut *self.records.lock().unwrap()); + let mut output = String::new(); + for record in records { + let indent = " ".repeat(record.depth); + let _ = writeln!(output, "{indent}{}", self.render_record(&record)); + } + output + } + + fn render_record(&self, record: &Record) -> String { + let field = |name: &str| record.fields.get(name).map(String::as_str); + let offsets = + field("offset_start") + .zip(field("offset_end")) + .map(|(start, end)| { + (start.parse::().unwrap(), end.parse::().unwrap()) + }); + + match (record.name.as_str(), offsets) { + ("query", _) => format!( + "--- Query: {} ({} rows) ---", + field("sql").unwrap_or_default(), + field("rows").unwrap_or("?"), + ), + ("object_store_request", _) => { + format!("Request {}", field("request").unwrap_or_default()) + } + // The event names the structure it reads. + (message, Some((start, end))) if field("column").is_some() => format!( + "{message} for column '{}' (rg {}) (offsets {start}..{end})", + field("column").unwrap(), + field("row_group").unwrap_or("?"), + ), + // The event knows only the byte range. The file layout names the + // structure that the range covers. + (_, Some((start, end))) => format!( + "Read {}", + self.layout.describe_range(&(start..end), self.file_size) + ), + (message, None) => match field("rows") { + Some(rows) => format!("{message} ({rows} rows)"), + None => message.to_string(), + }, + } + } +} diff --git a/datafusion/datasource-parquet/Cargo.toml b/datafusion/datasource-parquet/Cargo.toml index 5adb410cff95d..5631dc402be9f 100644 --- a/datafusion/datasource-parquet/Cargo.toml +++ b/datafusion/datasource-parquet/Cargo.toml @@ -43,6 +43,7 @@ object_store = { workspace = true } parking_lot = { workspace = true } parquet = { workspace = true } tokio = { workspace = true } +tracing = { workspace = true } [dev-dependencies] chrono = { workspace = true } diff --git a/datafusion/datasource-parquet/src/file_format.rs b/datafusion/datasource-parquet/src/file_format.rs index cfced8bbb9903..1a1ec7d3a4ecd 100644 --- a/datafusion/datasource-parquet/src/file_format.rs +++ b/datafusion/datasource-parquet/src/file_format.rs @@ -774,6 +774,12 @@ impl<'a> ObjectStoreFetch<'a> { impl MetadataFetch for ObjectStoreFetch<'_> { fn fetch(&mut self, range: Range) -> BoxFuture<'_, Result> { + tracing::debug!( + target: "datafusion::parquet::io", + offset_start = range.start, + offset_end = range.end, + "Read Parquet metadata" + ); async { self.store .get_range(&self.meta.location, range) diff --git a/datafusion/datasource-parquet/src/metadata.rs b/datafusion/datasource-parquet/src/metadata.rs index 1a3778b31ee72..c4d32aa6bf886 100644 --- a/datafusion/datasource-parquet/src/metadata.rs +++ b/datafusion/datasource-parquet/src/metadata.rs @@ -364,6 +364,12 @@ impl<'a> DFParquetMetadata<'a> { if let Some(hint) = self.metadata_size_hint { let prefetch_start = file_size.saturating_sub(hint as u64); let prefetch_range = prefetch_start..file_size; + tracing::debug!( + target: "datafusion::parquet::io", + offset_start = prefetch_range.start, + offset_end = prefetch_range.end, + "Prefetch Parquet footer" + ); let data = self .store .get_ranges( @@ -381,6 +387,25 @@ impl<'a> DFParquetMetadata<'a> { match decoder.try_decode().map_err(DataFusionError::from)? { DecodeResult::Data(metadata) => break metadata, DecodeResult::NeedsData(ranges) => { + for range in &ranges { + // The decoder asks for the footer at the end of the + // file, and for the page index before it. + if range.end == file_size { + tracing::debug!( + target: "datafusion::parquet::io", + offset_start = range.start, + offset_end = range.end, + "Read Parquet footer" + ); + } else { + tracing::debug!( + target: "datafusion::parquet::io", + offset_start = range.start, + offset_end = range.end, + "Read page index" + ); + } + } let buffers = self .store .get_ranges(&self.object_meta.location, &ranges) diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index e188b4543f284..8fdecc36e61c8 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -1547,6 +1547,19 @@ impl RowGroupsPrunedParquetOpen { for (column_name, column_idx, physical_type, type_length) in &parquet_columns { + let column = builder.metadata().row_group(idx).column(*column_idx); + if let (Some(offset), Some(length)) = + (column.bloom_filter_offset(), column.bloom_filter_length()) + { + tracing::debug!( + target: "datafusion::parquet::io", + column = column_name.as_str(), + row_group = idx, + offset_start = offset, + offset_end = offset + i64::from(length), + "Read Bloom filter" + ); + } let bf: Sbbf = match builder .get_row_group_column_bloom_filter(idx, *column_idx) .await diff --git a/datafusion/datasource-parquet/src/push_decoder.rs b/datafusion/datasource-parquet/src/push_decoder.rs index 57df6190daf87..381b27ac51d2b 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -501,6 +501,14 @@ impl PushDecoderStreamState { Ok(DecodeResult::NeedsData(ranges)) => { // I/O, not compute. timer.stop(); + for range in &ranges { + tracing::debug!( + target: "datafusion::parquet::io", + offset_start = range.start, + offset_end = range.end, + "Read column chunk data" + ); + } let data = self .reader .get_byte_ranges(ranges.clone()) From f6270cb5ed64d430e0130e5640a905c1ab6a2149 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 4 Nov 2026 19:10:17 -0300 Subject: [PATCH 13/14] prototype: trace parquet scan stages and render any recorded fields A fixed set of fields limited what each call site could report. Call sites now record whatever fields fit, and the test layer prints every field. The scan emits spans for metadata load, page index load, Bloom filter load, row group and page pruning, column chunk fetches, and events for row filter evaluation and metadata cache hits. The test object store emits a span for every operation. Co-Authored-By: Claude Opus 5.5 --- .../core/tests/parquet/io/bloom_filter_io.rs | 79 ++++--- .../tests/parquet/io/filter_pushdown_io.rs | 95 ++++---- datafusion/core/tests/parquet/io/mod.rs | 202 +++++++++--------- .../core/tests/parquet/io/page_pruning_io.rs | 111 ++++++---- .../tests/parquet/io/row_group_pruning_io.rs | 118 ++++++---- datafusion/core/tests/parquet/io/topk_io.rs | 66 +++--- .../core/tests/parquet/io/trace_capture.rs | 135 +++++------- .../datasource-parquet/src/file_format.rs | 7 +- datafusion/datasource-parquet/src/metadata.rs | 30 ++- .../datasource-parquet/src/opener/mod.rs | 53 ++++- .../datasource-parquet/src/push_decoder.rs | 14 +- .../datasource-parquet/src/row_filter.rs | 6 + 12 files changed, 507 insertions(+), 409 deletions(-) diff --git a/datafusion/core/tests/parquet/io/bloom_filter_io.rs b/datafusion/core/tests/parquet/io/bloom_filter_io.rs index ba9f128403876..e2c378329752f 100644 --- a/datafusion/core/tests/parquet/io/bloom_filter_io.rs +++ b/datafusion/core/tests/parquet/io/bloom_filter_io.rs @@ -44,24 +44,34 @@ async fn test_parquet_bloom_filter_io_matching() { .await .unwrap(); assert_snapshot!(traces.render(), @" - --- Query: SELECT * FROM t WHERE a = 10 AND b = 100 (1 rows) --- - Request 1 - HEAD object metadata - Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) - Request 2 - Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) - Read Bloom filter for column 'a' (rg 0) (offsets 98..145) - Request 3 - Read Bloom filter for column 'a' (rg 0) (offsets 98..145) - Read Bloom filter for column 'b' (rg 0) (offsets 145..192) - Request 4 - 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 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) + Query (sql=SELECT * FROM t WHERE a = 10 AND b = 100, rows=1) + ObjectStore::head (location=table.parquet) + Metadata cache hit (has_page_index=false) + Load page index into cached metadata + Read Parquet metadata (range=192..213) + ObjectStore::get (location=table.parquet) + Read byte range (range=192..213, parquet=OffsetIndex for column 'a' (rg 0)) + Open Parquet file (file=table.parquet, size=678) + Load Parquet metadata + Metadata cache hit (has_page_index=false) + Prune row groups + Row groups remaining (row_groups=1, row_groups_remaining=1, pruning_predicate=true) + Load Bloom filters + Read Bloom filter (column=a, row_group=0, range=98..145) + ObjectStore::get (location=table.parquet) + Read byte range (range=98..145, parquet=Bloom filter for column 'a' (rg 0)) + Read Bloom filter (column=b, row_group=0, range=145..192) + ObjectStore::get (location=table.parquet) + Read byte range (range=145..192, parquet=Bloom filter for column 'b' (rg 0)) + Prune row groups with Bloom filters + Row groups remaining (row_groups_remaining=1) + Prune pages with page index + Row groups remaining (row_groups_remaining=1) + Fetch column chunk data (ranges=2) + ObjectStore::get_ranges (location=table.parquet, ranges=2) + Read byte range (range=4..51, parquet=dictionary page and data page 0 for column 'a' (rg 0)) + Read byte range (range=51..98, parquet=dictionary page and data page 0 for column 'b' (rg 0)) + Produce batch (rows=1) "); } @@ -78,18 +88,27 @@ async fn test_parquet_bloom_filter_io_pruned() { .await .unwrap(); assert_snapshot!(traces.render(), @" - --- Query: SELECT * FROM t WHERE a = 999 AND b = 100 (0 rows) --- - Request 1 - HEAD object metadata - Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) - Request 2 - Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) - Read Bloom filter for column 'a' (rg 0) (offsets 98..145) - Request 3 - Read Bloom filter for column 'a' (rg 0) (offsets 98..145) - Read Bloom filter for column 'b' (rg 0) (offsets 145..192) - Request 4 - Read Bloom filter for column 'b' (rg 0) (offsets 145..192) + Query (sql=SELECT * FROM t WHERE a = 999 AND b = 100, rows=0) + ObjectStore::head (location=table.parquet) + Metadata cache hit (has_page_index=false) + Load page index into cached metadata + Read Parquet metadata (range=192..213) + ObjectStore::get (location=table.parquet) + Read byte range (range=192..213, parquet=OffsetIndex for column 'a' (rg 0)) + Open Parquet file (file=table.parquet, size=678) + Load Parquet metadata + Metadata cache hit (has_page_index=false) + Prune row groups + Row groups remaining (row_groups=1, row_groups_remaining=1, pruning_predicate=true) + Load Bloom filters + Read Bloom filter (column=a, row_group=0, range=98..145) + ObjectStore::get (location=table.parquet) + Read byte range (range=98..145, parquet=Bloom filter for column 'a' (rg 0)) + Read Bloom filter (column=b, row_group=0, range=145..192) + ObjectStore::get (location=table.parquet) + Read byte range (range=145..192, parquet=Bloom filter for column 'b' (rg 0)) + Prune row groups with Bloom filters + Row groups remaining (row_groups_remaining=0) "); } diff --git a/datafusion/core/tests/parquet/io/filter_pushdown_io.rs b/datafusion/core/tests/parquet/io/filter_pushdown_io.rs index 3b12ad00dc975..9ce98563d72f5 100644 --- a/datafusion/core/tests/parquet/io/filter_pushdown_io.rs +++ b/datafusion/core/tests/parquet/io/filter_pushdown_io.rs @@ -43,25 +43,33 @@ async fn test_filter_pushdown_io_disabled() { let traces = test.capture_traces(); test.query("SELECT * FROM t WHERE val = 7").await.unwrap(); assert_snapshot!(traces.render(), @" - --- Query: SELECT * FROM t WHERE val = 7 (1 rows) --- - Request 1 - HEAD object metadata - 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 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) - 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) - 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 (1 rows) + Query (sql=SELECT * FROM t WHERE val = 7, rows=1) + ObjectStore::head (location=table.parquet) + Metadata cache hit (has_page_index=true) + Open Parquet file (file=table.parquet, size=1402) + Load Parquet metadata + Metadata cache hit (has_page_index=true) + Prune row groups + Row groups remaining (row_groups=3, row_groups_remaining=3, pruning_predicate=true) + Load page index + Load Bloom filters + Prune row groups with Bloom filters + Row groups remaining (row_groups_remaining=3) + Prune pages with page index + Row groups remaining (row_groups_remaining=3) + Fetch column chunk data (ranges=2) + ObjectStore::get_ranges (location=table.parquet, ranges=2) + Read byte range (range=4..51, parquet=dictionary page and data page 0 for column 'id' (rg 0)) + Read byte range (range=51..98, parquet=dictionary page and data page 0 for column 'val' (rg 0)) + Fetch column chunk data (ranges=2) + ObjectStore::get_ranges (location=table.parquet, ranges=2) + Read byte range (range=98..145, parquet=dictionary page and data page 0 for column 'id' (rg 1)) + Read byte range (range=145..192, parquet=dictionary page and data page 0 for column 'val' (rg 1)) + Fetch column chunk data (ranges=2) + ObjectStore::get_ranges (location=table.parquet, ranges=2) + Read byte range (range=192..239, parquet=dictionary page and data page 0 for column 'id' (rg 2)) + Read byte range (range=239..286, parquet=dictionary page and data page 0 for column 'val' (rg 2)) + Produce batch (rows=1) "); } @@ -74,24 +82,37 @@ async fn test_filter_pushdown_io_enabled() { let traces = test.capture_traces(); test.query("SELECT * FROM t WHERE val = 7").await.unwrap(); assert_snapshot!(traces.render(), @" - --- Query: SELECT * FROM t WHERE val = 7 (1 rows) --- - Request 1 - HEAD object metadata - Read dictionary page and data page 0 for column 'val' (rg 0) (offsets 51..98) - Request 2 - 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) - Request 3 - 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 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) + Query (sql=SELECT * FROM t WHERE val = 7, rows=1) + ObjectStore::head (location=table.parquet) + Metadata cache hit (has_page_index=true) + Open Parquet file (file=table.parquet, size=1402) + Load Parquet metadata + Metadata cache hit (has_page_index=true) + Prune row groups + Row groups remaining (row_groups=3, row_groups_remaining=3, pruning_predicate=true) + Load page index + Load Bloom filters + Prune row groups with Bloom filters + Row groups remaining (row_groups_remaining=3) + Prune pages with page index + Row groups remaining (row_groups_remaining=3) + Fetch column chunk data (ranges=1) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=51..98, parquet=dictionary page and data page 0 for column 'val' (rg 0)) + Evaluate row filter (predicate=val@0 = 7, rows=3, rows_matched=0) + Fetch column chunk data (ranges=1) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=145..192, parquet=dictionary page and data page 0 for column 'val' (rg 1)) + Evaluate row filter (predicate=val@0 = 7, rows=3, rows_matched=1) + Fetch column chunk data (ranges=2) + ObjectStore::get_ranges (location=table.parquet, ranges=2) + Read byte range (range=98..124, parquet=dictionary page for column 'id' (rg 1)) + Read byte range (range=124..145, parquet=data page 0 for column 'id' (rg 1)) + Produce batch (rows=1) + Fetch column chunk data (ranges=1) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=239..286, parquet=dictionary page and data page 0 for column 'val' (rg 2)) + Evaluate row filter (predicate=val@0 = 7, rows=3, rows_matched=0) "); } diff --git a/datafusion/core/tests/parquet/io/mod.rs b/datafusion/core/tests/parquet/io/mod.rs index 4a66c342bdae5..7173beb0426f0 100644 --- a/datafusion/core/tests/parquet/io/mod.rs +++ b/datafusion/core/tests/parquet/io/mod.rs @@ -13,7 +13,6 @@ use std::ops::Range; use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; use async_trait::async_trait; use bytes::Bytes; @@ -29,6 +28,7 @@ use object_store::{ }; use parquet::file::reader::{FileReader, SerializedFileReader}; use parquet::file::serialized_reader::ReadOptionsBuilder; +use tracing::Instrument; use url::Url; // Tests are in submodules @@ -197,85 +197,51 @@ impl ParquetFileLayout { } } + /// Names the Parquet structure that `range` covers. 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); + return "footer".to_string(); } - 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 - ); - } + let overlaps = |r: &Range| range.start < r.end && range.end > r.start; + if let Some(desc) = self.descriptors.iter().find(|d| overlaps(&d.range)) { + return desc.description.clone(); } - 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 - ); - } + if let Some(chunk) = self.column_chunks.iter().find(|c| overlaps(&c.range)) { + return chunk.describe(range); } - format!("byte range (offsets {}..{})", range.start, range.end) + "unknown".to_string() } } +/// An [`ObjectStore`] that emits a `tracing` span for every operation. Each +/// byte range read gets an event that names the Parquet structure it covers. #[derive(Debug)] struct InstrumentedObjectStore { inner: Arc, + layout: ParquetFileLayout, file_size: u64, - /// Number of requests made to the store (reset for each query) - requests: Arc, } -impl InstrumentedObjectStore { - fn new( - inner: Arc, - file_size: u64, - requests: Arc, - ) -> Self { - Self { - inner, - file_size, - requests, - } - } +/// Target of the spans and events that this test harness emits. The capture +/// layer only records targets that start with `datafusion`. +const TEST_TARGET: &str = "datafusion::test"; - /// Emits a span for a `HEAD` request, which reads no bytes. - fn record_head_request(&self) { - let request = self.requests.fetch_add(1, Ordering::SeqCst) + 1; - let span = tracing::debug_span!( - target: "datafusion::test::object_store", - "object_store_request", - request - ); - span.in_scope(|| { - tracing::debug!(target: "datafusion::test::object_store", "HEAD object metadata"); - }); - } +macro_rules! store_span { + ($op:literal, $($fields:tt)*) => { + tracing::debug_span!(target: TEST_TARGET, $op, $($fields)*) + }; +} - /// Emits a span for one request to the store, with an event per byte - /// range it reads. - fn record_request(&self, ranges: impl IntoIterator>) { - let request = self.requests.fetch_add(1, Ordering::SeqCst) + 1; - let span = tracing::debug_span!( - target: "datafusion::test::object_store", - "object_store_request", - request - ); - span.in_scope(|| { - for range in ranges { - tracing::debug!( - target: "datafusion::test::object_store", - offset_start = range.start, - offset_end = range.end, - "Read byte range" - ); - } - }); +impl InstrumentedObjectStore { + fn record_ranges(&self, ranges: impl IntoIterator>) { + for range in ranges { + tracing::debug!( + target: TEST_TARGET, + range = ?range, + parquet = %self.layout.describe_range(&range, self.file_size), + "Read byte range" + ); + } } } @@ -293,7 +259,12 @@ impl ObjectStore for InstrumentedObjectStore { payload: PutPayload, opts: PutOptions, ) -> object_store::Result { - self.inner.put_opts(location, payload, opts).await + let span = + store_span!("ObjectStore::put", %location, bytes = payload.content_length()); + self.inner + .put_opts(location, payload, opts) + .instrument(span) + .await } async fn put_multipart_opts( @@ -301,7 +272,11 @@ impl ObjectStore for InstrumentedObjectStore { location: &Path, opts: PutMultipartOptions, ) -> object_store::Result> { - self.inner.put_multipart_opts(location, opts).await + let span = store_span!("ObjectStore::put_multipart", %location); + self.inner + .put_multipart_opts(location, opts) + .instrument(span) + .await } async fn get_opts( @@ -310,8 +285,12 @@ impl ObjectStore for InstrumentedObjectStore { options: GetOptions, ) -> object_store::Result { if options.head { - self.record_head_request(); - return self.inner.get_opts(location, options).await; + let span = store_span!("ObjectStore::head", %location); + return self + .inner + .get_opts(location, options) + .instrument(span) + .await; } let range = match &options.range { Some(GetRange::Bounded(r)) => r.clone(), @@ -321,8 +300,12 @@ impl ObjectStore for InstrumentedObjectStore { } None => 0..self.file_size, }; - self.record_request([range]); - self.inner.get_opts(location, options).await + let span = store_span!("ObjectStore::get", %location); + span.in_scope(|| self.record_ranges([range])); + self.inner + .get_opts(location, options) + .instrument(span) + .await } async fn get_ranges( @@ -330,14 +313,20 @@ impl ObjectStore for InstrumentedObjectStore { location: &Path, ranges: &[Range], ) -> object_store::Result> { - self.record_request(ranges.iter().cloned()); - self.inner.get_ranges(location, ranges).await + let span = + store_span!("ObjectStore::get_ranges", %location, ranges = ranges.len()); + span.in_scope(|| self.record_ranges(ranges.iter().cloned())); + self.inner + .get_ranges(location, ranges) + .instrument(span) + .await } fn delete_stream( &self, locations: BoxStream<'static, object_store::Result>, ) -> BoxStream<'static, object_store::Result> { + tracing::debug!(target: TEST_TARGET, "ObjectStore::delete_stream"); self.inner.delete_stream(locations) } @@ -345,6 +334,7 @@ impl ObjectStore for InstrumentedObjectStore { &self, prefix: Option<&Path>, ) -> BoxStream<'static, object_store::Result> { + tracing::debug!(target: TEST_TARGET, prefix = ?prefix.map(Path::as_ref), "ObjectStore::list"); self.inner.list(prefix) } @@ -353,6 +343,12 @@ impl ObjectStore for InstrumentedObjectStore { prefix: Option<&Path>, offset: &Path, ) -> BoxStream<'static, object_store::Result> { + tracing::debug!( + target: TEST_TARGET, + prefix = ?prefix.map(Path::as_ref), + %offset, + "ObjectStore::list_with_offset" + ); self.inner.list_with_offset(prefix, offset) } @@ -360,7 +356,14 @@ impl ObjectStore for InstrumentedObjectStore { &self, prefix: Option<&Path>, ) -> object_store::Result { - self.inner.list_with_delimiter(prefix).await + let span = store_span!( + "ObjectStore::list_with_delimiter", + prefix = ?prefix.map(Path::as_ref) + ); + self.inner + .list_with_delimiter(prefix) + .instrument(span) + .await } async fn copy_opts( @@ -369,7 +372,11 @@ impl ObjectStore for InstrumentedObjectStore { to: &Path, options: object_store::CopyOptions, ) -> object_store::Result<()> { - self.inner.copy_opts(from, to, options).await + let span = store_span!("ObjectStore::copy", %from, %to); + self.inner + .copy_opts(from, to, options) + .instrument(span) + .await } } @@ -387,9 +394,6 @@ impl ObjectStore for InstrumentedObjectStore { /// ``` pub(crate) struct ParquetIoTest { ctx: SessionContext, - layout: ParquetFileLayout, - file_size: u64, - requests: Arc, } impl ParquetIoTest { @@ -409,12 +413,11 @@ impl ParquetIoTest { .await .unwrap(); - let requests = Arc::new(AtomicUsize::new(0)); - let instrumented_store = Arc::new(InstrumentedObjectStore::new( - memory_store, + let instrumented_store = Arc::new(InstrumentedObjectStore { + inner: memory_store, + layout, file_size, - Arc::clone(&requests), - )); + }); let ctx = SessionContext::new_with_config(config); ctx.register_object_store(&Url::parse("mem://").unwrap(), instrumented_store); @@ -422,40 +425,37 @@ impl ParquetIoTest { .await .unwrap(); - Self { - ctx, - layout, - file_size, - requests, - } + Self { ctx } } /// Captures the spans and events that DataFusion emits on this thread /// until the returned [`TraceCapture`] is dropped. pub(crate) fn capture_traces(&self) -> TraceCapture { - TraceCapture::new(self.layout.clone(), self.file_size) + TraceCapture::new() } /// Runs `sql` to completion inside a `query` span and returns the number /// of output rows. Emits a `Produce batch` event for each output batch. pub(crate) async fn query(&self, sql: &str) -> Result { - self.requests.store(0, Ordering::SeqCst); let span = tracing::debug_span!( - target: "datafusion::test", - "query", + target: TEST_TARGET, + "Query", sql, rows = tracing::field::Empty ); - let _entered = span.enter(); - - let df = self.ctx.sql(sql).await?; - let mut stream = df.execute_stream().await?; - let mut row_count = 0; - while let Some(batch) = stream.next().await { - let rows = batch?.num_rows(); - row_count += rows; - tracing::debug!(target: "datafusion::test", rows, "Produce batch"); + let row_count = async { + let df = self.ctx.sql(sql).await?; + let mut stream = df.execute_stream().await?; + let mut row_count = 0; + while let Some(batch) = stream.next().await { + let rows = batch?.num_rows(); + row_count += rows; + tracing::debug!(target: TEST_TARGET, rows, "Produce batch"); + } + Ok::<_, datafusion::error::DataFusionError>(row_count) } + .instrument(span.clone()) + .await?; span.record("rows", row_count); Ok(row_count) } diff --git a/datafusion/core/tests/parquet/io/page_pruning_io.rs b/datafusion/core/tests/parquet/io/page_pruning_io.rs index fa2bb6dfba8a0..f67acf7c990f9 100644 --- a/datafusion/core/tests/parquet/io/page_pruning_io.rs +++ b/datafusion/core/tests/parquet/io/page_pruning_io.rs @@ -47,15 +47,21 @@ async fn test_page_pruning_io_no_predicate() { let traces = test.capture_traces(); test.query("SELECT * FROM t").await.unwrap(); assert_snapshot!(traces.render(), @" - --- Query: SELECT * FROM t (9 rows) --- - Request 1 - HEAD object metadata - 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 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) + Query (sql=SELECT * FROM t, rows=9) + ObjectStore::head (location=table.parquet) + Metadata cache hit (has_page_index=true) + Open Parquet file (file=table.parquet, size=838) + Load Parquet metadata + Metadata cache hit (has_page_index=true) + Prune row groups + Row groups remaining (row_groups=1, row_groups_remaining=1, pruning_predicate=false) + Load Bloom filters + Prune row groups with Bloom filters + Fetch column chunk data (ranges=2) + ObjectStore::get_ranges (location=table.parquet, ranges=2) + Read byte range (range=4..91, parquet=data pages 0..=2 for column 'id' (rg 0)) + Read byte range (range=91..178, parquet=data pages 0..=2 for column 'val' (rg 0)) + Produce batch (rows=9) "); } @@ -67,15 +73,25 @@ async fn test_page_pruning_io_one_page() { let traces = test.capture_traces(); test.query("SELECT * FROM t WHERE id = 5").await.unwrap(); assert_snapshot!(traces.render(), @" - --- Query: SELECT * FROM t WHERE id = 5 (1 rows) --- - Request 1 - HEAD object metadata - 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 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) + Query (sql=SELECT * FROM t WHERE id = 5, rows=1) + ObjectStore::head (location=table.parquet) + Metadata cache hit (has_page_index=true) + Open Parquet file (file=table.parquet, size=838) + Load Parquet metadata + Metadata cache hit (has_page_index=true) + Prune row groups + Row groups remaining (row_groups=1, row_groups_remaining=1, pruning_predicate=true) + Load page index + Load Bloom filters + Prune row groups with Bloom filters + Row groups remaining (row_groups_remaining=1) + Prune pages with page index + Row groups remaining (row_groups_remaining=1) + Fetch column chunk data (ranges=2) + ObjectStore::get_ranges (location=table.parquet, ranges=2) + Read byte range (range=33..62, parquet=data page 1 for column 'id' (rg 0)) + Read byte range (range=120..149, parquet=data page 1 for column 'val' (rg 0)) + Produce batch (rows=1) "); } @@ -87,19 +103,27 @@ async fn test_page_pruning_io_two_pages() { let traces = test.capture_traces(); test.query("SELECT * FROM t WHERE id > 3").await.unwrap(); assert_snapshot!(traces.render(), @" - --- Query: SELECT * FROM t WHERE id > 3 (6 rows) --- - Request 1 - HEAD object metadata - 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 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) + Query (sql=SELECT * FROM t WHERE id > 3, rows=6) + ObjectStore::head (location=table.parquet) + Metadata cache hit (has_page_index=true) + Open Parquet file (file=table.parquet, size=838) + Load Parquet metadata + Metadata cache hit (has_page_index=true) + Prune row groups + Row groups remaining (row_groups=1, row_groups_remaining=1, pruning_predicate=true) + Load page index + Load Bloom filters + Prune row groups with Bloom filters + Row groups remaining (row_groups_remaining=1) + Prune pages with page index + Row groups remaining (row_groups_remaining=1) + Fetch column chunk data (ranges=4) + ObjectStore::get_ranges (location=table.parquet, ranges=4) + Read byte range (range=33..62, parquet=data page 1 for column 'id' (rg 0)) + Read byte range (range=62..91, parquet=data page 2 for column 'id' (rg 0)) + Read byte range (range=120..149, parquet=data page 1 for column 'val' (rg 0)) + Read byte range (range=149..178, parquet=data page 2 for column 'val' (rg 0)) + Produce batch (rows=6) "); } @@ -111,15 +135,22 @@ async fn test_page_pruning_io_disabled() { let traces = test.capture_traces(); test.query("SELECT * FROM t WHERE id = 5").await.unwrap(); assert_snapshot!(traces.render(), @" - --- Query: SELECT * FROM t WHERE id = 5 (1 rows) --- - Request 1 - HEAD object metadata - 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 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) + Query (sql=SELECT * FROM t WHERE id = 5, rows=1) + ObjectStore::head (location=table.parquet) + Metadata cache hit (has_page_index=true) + Open Parquet file (file=table.parquet, size=838) + Load Parquet metadata + Metadata cache hit (has_page_index=true) + Prune row groups + Row groups remaining (row_groups=1, row_groups_remaining=1, pruning_predicate=true) + Load Bloom filters + Prune row groups with Bloom filters + Row groups remaining (row_groups_remaining=1) + Fetch column chunk data (ranges=2) + ObjectStore::get_ranges (location=table.parquet, ranges=2) + Read byte range (range=4..91, parquet=data pages 0..=2 for column 'id' (rg 0)) + Read byte range (range=91..178, parquet=data pages 0..=2 for column 'val' (rg 0)) + Produce batch (rows=1) "); } 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 aa58324270e79..d0d70b0e06869 100644 --- a/datafusion/core/tests/parquet/io/row_group_pruning_io.rs +++ b/datafusion/core/tests/parquet/io/row_group_pruning_io.rs @@ -41,27 +41,31 @@ async fn test_row_group_pruning_io_no_predicate() { let traces = test.capture_traces(); test.query("SELECT * FROM t").await.unwrap(); assert_snapshot!(traces.render(), @" - --- Query: SELECT * FROM t (9 rows) --- - Request 1 - HEAD object metadata - 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 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) + Query (sql=SELECT * FROM t, rows=9) + ObjectStore::head (location=table.parquet) + Metadata cache hit (has_page_index=true) + Open Parquet file (file=table.parquet, size=1402) + Load Parquet metadata + Metadata cache hit (has_page_index=true) + Prune row groups + Row groups remaining (row_groups=3, row_groups_remaining=3, pruning_predicate=false) + Load Bloom filters + Prune row groups with Bloom filters + Fetch column chunk data (ranges=2) + ObjectStore::get_ranges (location=table.parquet, ranges=2) + Read byte range (range=4..51, parquet=dictionary page and data page 0 for column 'id' (rg 0)) + Read byte range (range=51..98, parquet=dictionary page and data page 0 for column 'val' (rg 0)) + Produce batch (rows=3) + Fetch column chunk data (ranges=2) + ObjectStore::get_ranges (location=table.parquet, ranges=2) + Read byte range (range=98..145, parquet=dictionary page and data page 0 for column 'id' (rg 1)) + Read byte range (range=145..192, parquet=dictionary page and data page 0 for column 'val' (rg 1)) + Produce batch (rows=3) + Fetch column chunk data (ranges=2) + ObjectStore::get_ranges (location=table.parquet, ranges=2) + Read byte range (range=192..239, parquet=dictionary page and data page 0 for column 'id' (rg 2)) + Read byte range (range=239..286, parquet=dictionary page and data page 0 for column 'val' (rg 2)) + Produce batch (rows=3) "); } @@ -73,15 +77,25 @@ async fn test_row_group_pruning_io_one_row_group() { let traces = test.capture_traces(); test.query("SELECT * FROM t WHERE id = 5").await.unwrap(); assert_snapshot!(traces.render(), @" - --- Query: SELECT * FROM t WHERE id = 5 (1 rows) --- - Request 1 - HEAD object metadata - 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 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) + Query (sql=SELECT * FROM t WHERE id = 5, rows=1) + ObjectStore::head (location=table.parquet) + Metadata cache hit (has_page_index=true) + Open Parquet file (file=table.parquet, size=1402) + Load Parquet metadata + Metadata cache hit (has_page_index=true) + Prune row groups + Row groups remaining (row_groups=3, row_groups_remaining=1, pruning_predicate=true) + Load page index + Load Bloom filters + Prune row groups with Bloom filters + Row groups remaining (row_groups_remaining=1) + Prune pages with page index + Row groups remaining (row_groups_remaining=1) + Fetch column chunk data (ranges=2) + ObjectStore::get_ranges (location=table.parquet, ranges=2) + Read byte range (range=98..145, parquet=dictionary page and data page 0 for column 'id' (rg 1)) + Read byte range (range=145..192, parquet=dictionary page and data page 0 for column 'val' (rg 1)) + Produce batch (rows=1) "); } @@ -93,20 +107,29 @@ async fn test_row_group_pruning_io_two_row_groups() { let traces = test.capture_traces(); test.query("SELECT * FROM t WHERE id > 5").await.unwrap(); assert_snapshot!(traces.render(), @" - --- Query: SELECT * FROM t WHERE id > 5 (4 rows) --- - Request 1 - HEAD object metadata - 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 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) - 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 2) (offsets 192..239) - Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) - Produce batch (4 rows) + Query (sql=SELECT * FROM t WHERE id > 5, rows=4) + ObjectStore::head (location=table.parquet) + Metadata cache hit (has_page_index=true) + Open Parquet file (file=table.parquet, size=1402) + Load Parquet metadata + Metadata cache hit (has_page_index=true) + Prune row groups + Row groups remaining (row_groups=3, row_groups_remaining=2, pruning_predicate=true) + Load page index + Load Bloom filters + Prune row groups with Bloom filters + Row groups remaining (row_groups_remaining=2) + Prune pages with page index + Row groups remaining (row_groups_remaining=2) + Fetch column chunk data (ranges=2) + ObjectStore::get_ranges (location=table.parquet, ranges=2) + Read byte range (range=98..145, parquet=dictionary page and data page 0 for column 'id' (rg 1)) + Read byte range (range=145..192, parquet=dictionary page and data page 0 for column 'val' (rg 1)) + Fetch column chunk data (ranges=2) + ObjectStore::get_ranges (location=table.parquet, ranges=2) + Read byte range (range=192..239, parquet=dictionary page and data page 0 for column 'id' (rg 2)) + Read byte range (range=239..286, parquet=dictionary page and data page 0 for column 'val' (rg 2)) + Produce batch (rows=4) "); } @@ -117,9 +140,10 @@ async fn test_row_group_pruning_io_all_pruned() { let traces = test.capture_traces(); test.query("SELECT * FROM t WHERE id > 100").await.unwrap(); assert_snapshot!(traces.render(), @" - --- Query: SELECT * FROM t WHERE id > 100 (0 rows) --- - Request 1 - HEAD object metadata + Query (sql=SELECT * FROM t WHERE id > 100, rows=0) + ObjectStore::head (location=table.parquet) + Metadata cache hit (has_page_index=true) + Open Parquet file (file=table.parquet, size=1402) "); } diff --git a/datafusion/core/tests/parquet/io/topk_io.rs b/datafusion/core/tests/parquet/io/topk_io.rs index 97e2d8077ecc8..b3f259b1ad391 100644 --- a/datafusion/core/tests/parquet/io/topk_io.rs +++ b/datafusion/core/tests/parquet/io/topk_io.rs @@ -49,15 +49,21 @@ async fn test_topk_io_dynamic_filter() { .await .unwrap(); assert_snapshot!(traces.render(), @" - --- Query: SELECT * FROM t ORDER BY id DESC LIMIT 2 (2 rows) --- - Request 1 - HEAD object metadata - 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 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) + Query (sql=SELECT * FROM t ORDER BY id DESC LIMIT 2, rows=2) + ObjectStore::head (location=table.parquet) + Metadata cache hit (has_page_index=true) + Open Parquet file (file=table.parquet, size=1402) + Load Parquet metadata + Metadata cache hit (has_page_index=true) + Prune row groups + Row groups remaining (row_groups=3, row_groups_remaining=3, pruning_predicate=false) + Load Bloom filters + Prune row groups with Bloom filters + Fetch column chunk data (ranges=2) + ObjectStore::get_ranges (location=table.parquet, ranges=2) + Read byte range (range=192..239, parquet=dictionary page and data page 0 for column 'id' (rg 2)) + Read byte range (range=239..286, parquet=dictionary page and data page 0 for column 'val' (rg 2)) + Produce batch (rows=2) "); } @@ -71,25 +77,29 @@ async fn test_topk_io_no_dynamic_filter() { .await .unwrap(); assert_snapshot!(traces.render(), @" - --- Query: SELECT * FROM t ORDER BY id DESC LIMIT 2 (2 rows) --- - Request 1 - HEAD object metadata - 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 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) - 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) - 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 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) + Query (sql=SELECT * FROM t ORDER BY id DESC LIMIT 2, rows=2) + ObjectStore::head (location=table.parquet) + Metadata cache hit (has_page_index=true) + Open Parquet file (file=table.parquet, size=1402) + Load Parquet metadata + Metadata cache hit (has_page_index=true) + Prune row groups + Row groups remaining (row_groups=3, row_groups_remaining=3, pruning_predicate=false) + Load Bloom filters + Prune row groups with Bloom filters + Fetch column chunk data (ranges=2) + ObjectStore::get_ranges (location=table.parquet, ranges=2) + Read byte range (range=192..239, parquet=dictionary page and data page 0 for column 'id' (rg 2)) + Read byte range (range=239..286, parquet=dictionary page and data page 0 for column 'val' (rg 2)) + Fetch column chunk data (ranges=2) + ObjectStore::get_ranges (location=table.parquet, ranges=2) + Read byte range (range=98..145, parquet=dictionary page and data page 0 for column 'id' (rg 1)) + Read byte range (range=145..192, parquet=dictionary page and data page 0 for column 'val' (rg 1)) + Fetch column chunk data (ranges=2) + ObjectStore::get_ranges (location=table.parquet, ranges=2) + Read byte range (range=4..51, parquet=dictionary page and data page 0 for column 'id' (rg 0)) + Read byte range (range=51..98, parquet=dictionary page and data page 0 for column 'val' (rg 0)) + Produce batch (rows=2) "); } diff --git a/datafusion/core/tests/parquet/io/trace_capture.rs b/datafusion/core/tests/parquet/io/trace_capture.rs index e91005b73874d..6370097af41c9 100644 --- a/datafusion/core/tests/parquet/io/trace_capture.rs +++ b/datafusion/core/tests/parquet/io/trace_capture.rs @@ -3,18 +3,22 @@ // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements //! Captures `tracing` spans and events emitted by DataFusion during a test and -//! renders them as text for `insta` snapshots. +//! renders them as an indented tree for `insta` snapshots. //! -//! The library emits events with a static message and structured fields -//! (`offset_start`, `offset_end`, `column`, `row_group`, ...). This module owns -//! the human-readable rendering, so the library never formats strings for -//! tests. +//! Each line is the span name (or the event's `message`) followed by every +//! field the call site recorded, in call-site order: +//! +//! ```text +//! Load Bloom filters +//! Read Bloom filter (column=a, row_group=0, range=98..145) +//! ObjectStore::get_ranges (location=table.parquet, ranges=1) +//! Read byte range (range=98..145, parquet=Bloom filter for column 'a' (rg 0)) +//! ``` //! //! The subscriber is installed with [`tracing::subscriber::set_default`], which //! is thread-local. `#[tokio::test]` uses a current-thread runtime, so every //! task of the query runs on the test thread and reports to this subscriber. -use std::collections::BTreeMap; use std::fmt::{self, Write as _}; use std::sync::{Arc, Mutex}; @@ -25,57 +29,51 @@ use tracing::{Event, Subscriber}; use tracing_subscriber::layer::{Context, Layer, SubscriberExt}; use tracing_subscriber::registry::{LookupSpan, Registry}; -use super::ParquetFileLayout; - -/// Only spans and events from these target prefixes are captured. Dependencies -/// such as `h2` or `object_store` can emit their own events. +/// Only spans and events whose target starts with this prefix are captured. +/// Dependencies such as `h2` emit their own events. const TARGET_PREFIX: &str = "datafusion"; -/// One captured span (rendered when it opens) or event. +/// One captured span (rendered where it opens) or event. #[derive(Debug, Clone)] struct Record { - /// Nesting level below the outermost captured span depth: usize, /// The span name, or the event's `message` field name: String, - fields: BTreeMap, + fields: Vec<(String, String)>, } #[derive(Default)] -struct FieldVisitor(BTreeMap); +struct FieldVisitor(Vec<(String, String)>); impl Visit for FieldVisitor { fn record_str(&mut self, field: &Field, value: &str) { - self.0.insert(field.name().to_string(), value.to_string()); + self.0.push((field.name().to_string(), value.to_string())); } fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) { self.0 - .insert(field.name().to_string(), format!("{value:?}")); + .push((field.name().to_string(), format!("{value:?}"))); } } -/// Index into [`CaptureLayer::records`] for a span, plus the depth its children -/// get. Stored in the span's extensions. +/// Stored in a span's extensions: its index in the captured records. struct SpanSlot { record: usize, - child_depth: usize, + depth: usize, } struct CaptureLayer { records: Arc>>, } -impl CaptureLayer { - fn parent_depth(ctx: &Context<'_, S>, parent: Option) -> usize - where - S: Subscriber + for<'a> LookupSpan<'a>, - { - parent - .and_then(|id| ctx.span(&id)) - .and_then(|span| span.extensions().get::().map(|s| s.child_depth)) - .unwrap_or(0) - } +fn child_depth(ctx: &Context<'_, S>, parent: Option) -> usize +where + S: Subscriber + for<'a> LookupSpan<'a>, +{ + parent + .and_then(|id| ctx.span(&id)) + .and_then(|span| span.extensions().get::().map(|s| s.depth + 1)) + .unwrap_or(0) } impl Layer for CaptureLayer @@ -92,7 +90,7 @@ where } else { attrs.parent().cloned() }; - let depth = Self::parent_depth(&ctx, parent); + let depth = child_depth(&ctx, parent); let mut visitor = FieldVisitor::default(); attrs.record(&mut visitor); @@ -102,17 +100,10 @@ where name: metadata.name().to_string(), fields: visitor.0, }); - // The query span is the frame of the snapshot, so its children are not - // indented. - let child_depth = if metadata.name() == "query" { - depth - } else { - depth + 1 - }; let span = ctx.span(id).expect("span exists"); span.extensions_mut().insert(SpanSlot { record: records.len() - 1, - child_depth, + depth, }); } @@ -133,11 +124,15 @@ where if !event.metadata().target().starts_with(TARGET_PREFIX) { return; } - let depth = Self::parent_depth(&ctx, ctx.event_span(event).map(|s| s.id())); + let depth = child_depth(&ctx, ctx.event_span(event).map(|s| s.id())); let mut visitor = FieldVisitor::default(); event.record(&mut visitor); let mut fields = visitor.0; - let name = fields.remove("message").unwrap_or_default(); + let name = fields + .iter() + .position(|(name, _)| name == "message") + .map(|idx| fields.remove(idx).1) + .unwrap_or_default(); self.records.lock().unwrap().push(Record { depth, name, @@ -149,13 +144,11 @@ where /// Spans and events captured on this thread until it is dropped. pub(crate) struct TraceCapture { records: Arc>>, - layout: ParquetFileLayout, - file_size: u64, _guard: DefaultGuard, } impl TraceCapture { - pub(super) fn new(layout: ParquetFileLayout, file_size: u64) -> Self { + pub(crate) fn new() -> Self { let records = Arc::new(Mutex::new(Vec::new())); let subscriber = Registry::default().with(CaptureLayer { records: Arc::clone(&records), @@ -163,57 +156,27 @@ impl TraceCapture { let guard = tracing::subscriber::set_default(subscriber); Self { records, - layout, - file_size, _guard: guard, } } - /// Renders and clears everything captured so far. + /// Renders everything captured so far as an indented tree. pub(crate) fn render(&self) -> String { - let records = std::mem::take(&mut *self.records.lock().unwrap()); + let records = self.records.lock().unwrap(); let mut output = String::new(); - for record in records { + for record in records.iter() { let indent = " ".repeat(record.depth); - let _ = writeln!(output, "{indent}{}", self.render_record(&record)); - } - output - } - - fn render_record(&self, record: &Record) -> String { - let field = |name: &str| record.fields.get(name).map(String::as_str); - let offsets = - field("offset_start") - .zip(field("offset_end")) - .map(|(start, end)| { - (start.parse::().unwrap(), end.parse::().unwrap()) - }); - - match (record.name.as_str(), offsets) { - ("query", _) => format!( - "--- Query: {} ({} rows) ---", - field("sql").unwrap_or_default(), - field("rows").unwrap_or("?"), - ), - ("object_store_request", _) => { - format!("Request {}", field("request").unwrap_or_default()) + let _ = write!(output, "{indent}{}", record.name); + if !record.fields.is_empty() { + let fields: Vec = record + .fields + .iter() + .map(|(name, value)| format!("{name}={value}")) + .collect(); + let _ = write!(output, " ({})", fields.join(", ")); } - // The event names the structure it reads. - (message, Some((start, end))) if field("column").is_some() => format!( - "{message} for column '{}' (rg {}) (offsets {start}..{end})", - field("column").unwrap(), - field("row_group").unwrap_or("?"), - ), - // The event knows only the byte range. The file layout names the - // structure that the range covers. - (_, Some((start, end))) => format!( - "Read {}", - self.layout.describe_range(&(start..end), self.file_size) - ), - (message, None) => match field("rows") { - Some(rows) => format!("{message} ({rows} rows)"), - None => message.to_string(), - }, + output.push('\n'); } + output } } diff --git a/datafusion/datasource-parquet/src/file_format.rs b/datafusion/datasource-parquet/src/file_format.rs index 1a1ec7d3a4ecd..aac54274374c6 100644 --- a/datafusion/datasource-parquet/src/file_format.rs +++ b/datafusion/datasource-parquet/src/file_format.rs @@ -774,12 +774,7 @@ impl<'a> ObjectStoreFetch<'a> { impl MetadataFetch for ObjectStoreFetch<'_> { fn fetch(&mut self, range: Range) -> BoxFuture<'_, Result> { - tracing::debug!( - target: "datafusion::parquet::io", - offset_start = range.start, - offset_end = range.end, - "Read Parquet metadata" - ); + tracing::debug!(range = ?range, "Read Parquet metadata"); async { self.store .get_range(&self.meta.location, range) diff --git a/datafusion/datasource-parquet/src/metadata.rs b/datafusion/datasource-parquet/src/metadata.rs index c4d32aa6bf886..171767fbc5f91 100644 --- a/datafusion/datasource-parquet/src/metadata.rs +++ b/datafusion/datasource-parquet/src/metadata.rs @@ -40,6 +40,7 @@ use parquet::file::statistics::Statistics as ParquetStatistics; use parquet::schema::types::{ColumnDescriptor, SchemaDescriptor}; use std::any::Any; use std::sync::Arc; +use tracing::Instrument; /// Minimum fraction of row groups that must report NDV statistics for the /// merged result to be `Inexact` rather than `Absent`, as the estimate @@ -276,6 +277,10 @@ impl<'a> DFParquetMetadata<'a> { .downcast_ref::() { let cached_metadata = Arc::clone(cached_parquet.parquet_metadata()); + tracing::debug!( + has_page_index = Self::metadata_has_page_index(cached_metadata.as_ref()), + "Metadata cache hit" + ); // Reuse the cache when it already has page index, or when the caller // asked to skip page index I/O (footer-only metadata is sufficient). if Self::metadata_has_page_index(cached_metadata.as_ref()) @@ -285,6 +290,9 @@ impl<'a> DFParquetMetadata<'a> { } let metadata = Self::load_page_index(self.store, self.object_meta, cached_metadata) + .instrument(tracing::debug_span!( + "Load page index into cached metadata" + )) .await?; if cache_metadata { self.cache_metadata(Arc::clone(&metadata))?; @@ -292,6 +300,7 @@ impl<'a> DFParquetMetadata<'a> { return Ok(metadata); } + tracing::debug!(cache_enabled = cache_metadata, "Metadata cache miss"); let metadata = self.fetch_metadata_from_store(page_index_policy).await?; if cache_metadata { self.cache_metadata(Arc::clone(&metadata))?; @@ -364,12 +373,7 @@ impl<'a> DFParquetMetadata<'a> { if let Some(hint) = self.metadata_size_hint { let prefetch_start = file_size.saturating_sub(hint as u64); let prefetch_range = prefetch_start..file_size; - tracing::debug!( - target: "datafusion::parquet::io", - offset_start = prefetch_range.start, - offset_end = prefetch_range.end, - "Prefetch Parquet footer" - ); + tracing::debug!(range = ?prefetch_range, "Prefetch Parquet footer"); let data = self .store .get_ranges( @@ -391,19 +395,9 @@ impl<'a> DFParquetMetadata<'a> { // The decoder asks for the footer at the end of the // file, and for the page index before it. if range.end == file_size { - tracing::debug!( - target: "datafusion::parquet::io", - offset_start = range.start, - offset_end = range.end, - "Read Parquet footer" - ); + tracing::debug!(range = ?range, "Read Parquet footer"); } else { - tracing::debug!( - target: "datafusion::parquet::io", - offset_start = range.start, - offset_end = range.end, - "Read page index" - ); + tracing::debug!(range = ?range, "Read page index"); } } let buffers = self diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index 8fdecc36e61c8..a112b3e010155 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -76,6 +76,7 @@ use parquet::arrow::parquet_column; use parquet::basic::Type; use parquet::bloom_filter::Sbbf; use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader, RowGroupMetaData}; +use tracing::Instrument; /// Morselizer-level state for virtual columns, precomputed once per scan /// partition so each file skips the validator walks, `null_replacements` @@ -636,7 +637,12 @@ impl ParquetOpenState { let Some(prepared) = (*prepared).prune_file()? else { return Ok(ParquetOpenState::Done); }; - Ok(ParquetOpenState::LoadMetadata(prepared.load().boxed())) + Ok(ParquetOpenState::LoadMetadata( + prepared + .load() + .instrument(tracing::debug_span!("Load Parquet metadata")) + .boxed(), + )) } ParquetOpenState::LoadMetadata(future) => { Ok(ParquetOpenState::LoadMetadata(future)) @@ -654,7 +660,10 @@ impl ParquetOpenState { } if prepared_row_groups.should_load_page_index()? { Ok(ParquetOpenState::LoadPageIndex( - prepared_row_groups.load_page_index().boxed(), + prepared_row_groups + .load_page_index() + .instrument(tracing::debug_span!("Load page index")) + .boxed(), )) } else { if prepared_row_groups @@ -672,7 +681,10 @@ impl ParquetOpenState { ); } Ok(ParquetOpenState::LoadBloomFilters( - prepared_row_groups.load_bloom_filters().boxed(), + prepared_row_groups + .load_bloom_filters() + .instrument(tracing::debug_span!("Load Bloom filters")) + .boxed(), )) } } @@ -797,7 +809,11 @@ impl MorselPlanner for ParquetMorselPlanner { ParquetOpenState::LoadPageIndex(future) => { Ok(Some(Self::schedule_io(async move { Ok(ParquetOpenState::LoadBloomFilters( - future.await?.load_bloom_filters().boxed(), + future + .await? + .load_bloom_filters() + .instrument(tracing::debug_span!("Load Bloom filters")) + .boxed(), )) }))) } @@ -831,6 +847,11 @@ impl ParquetMorselizer { let file_range = partitioned_file.range.clone(); let extensions = partitioned_file.extensions.clone(); let file_name = partitioned_file.object_meta.location.to_string(); + tracing::debug!( + file = %file_name, + size = partitioned_file.object_meta.size, + "Open Parquet file" + ); let file_metrics = ParquetFileMetrics::new(self.partition_index, &file_name, &self.metrics); let baseline_metrics = BaselineMetrics::new(&self.metrics, self.partition_index); @@ -1273,6 +1294,7 @@ impl MetadataLoadedParquetOpen { impl FiltersPreparedParquetOpen { /// Prune row groups using file ranges and parquet metadata. fn prune_row_groups(self) -> Result { + let _span = tracing::debug_span!("Prune row groups").entered(); let loaded = &self.loaded; let prepared = &loaded.prepared; let file_metadata = Arc::clone(loaded.reader_metadata.metadata()); @@ -1354,6 +1376,12 @@ impl FiltersPreparedParquetOpen { .add_matched(remaining); } + tracing::debug!( + row_groups = rg_metadata.len(), + row_groups_remaining = row_groups.remaining_row_group_count(), + pruning_predicate = self.pruning_predicate.is_some(), + "Row groups remaining" + ); Ok(RowGroupsPrunedParquetOpen { prepared: self, row_groups, @@ -1552,11 +1580,9 @@ impl RowGroupsPrunedParquetOpen { (column.bloom_filter_offset(), column.bloom_filter_length()) { tracing::debug!( - target: "datafusion::parquet::io", - column = column_name.as_str(), + column = %column_name, row_group = idx, - offset_start = offset, - offset_end = offset + i64::from(length), + range = ?(offset..offset + i64::from(length)), "Read Bloom filter" ); } @@ -1602,6 +1628,7 @@ impl BloomFiltersLoadedParquetOpen { .bloom_filter_eval_time .clone(); let _timer_guard = bloom_filter_eval_time.timer(); + let _span = tracing::debug_span!("Prune row groups with Bloom filters").entered(); if let Some(predicate) = self .prepared .prepared @@ -1616,6 +1643,11 @@ impl BloomFiltersLoadedParquetOpen { &self.prepared.prepared.loaded.prepared.file_metrics, &self.row_group_bloom_filters, ); + tracing::debug!( + row_groups_remaining = + self.prepared.row_groups.remaining_row_group_count(), + "Row groups remaining" + ); } self.prepared @@ -1661,6 +1693,7 @@ impl RowGroupsPrunedParquetOpen { && !access_plan.is_empty() && let Some(page_pruning_predicate) = page_pruning_predicate { + let _span = tracing::debug_span!("Prune pages with page index").entered(); let page_pruning_result = page_pruning_predicate .prune_plan_with_page_index_and_metrics( access_plan, @@ -1675,6 +1708,10 @@ impl RowGroupsPrunedParquetOpen { }, ); access_plan = page_pruning_result.access_plan; + tracing::debug!( + row_groups_remaining = access_plan.row_group_index_iter().count(), + "Row groups remaining" + ); ParquetFileMetrics::add_page_index_pages_skipped_by_fully_matched( &prepared.metrics, prepared.partition_index, diff --git a/datafusion/datasource-parquet/src/push_decoder.rs b/datafusion/datasource-parquet/src/push_decoder.rs index 381b27ac51d2b..e5a697810ccb9 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -41,6 +41,7 @@ use parquet::arrow::push_decoder::{ ParquetPushDecoder, ParquetPushDecoderBuilder, RowGroupSelection, }; use parquet::file::metadata::ParquetMetaData; +use tracing::Instrument; use datafusion_common::{DataFusionError, Result, internal_err}; use datafusion_physical_expr::expressions::DynamicFilterTracking; @@ -501,17 +502,14 @@ impl PushDecoderStreamState { Ok(DecodeResult::NeedsData(ranges)) => { // I/O, not compute. timer.stop(); - for range in &ranges { - tracing::debug!( - target: "datafusion::parquet::io", - offset_start = range.start, - offset_end = range.end, - "Read column chunk data" - ); - } + let span = tracing::debug_span!( + "Fetch column chunk data", + ranges = ranges.len() + ); let data = self .reader .get_byte_ranges(ranges.clone()) + .instrument(span) .await .map_err(DataFusionError::from); timer.restart(); diff --git a/datafusion/datasource-parquet/src/row_filter.rs b/datafusion/datasource-parquet/src/row_filter.rs index f07e01a365bab..71ccf0760f39e 100644 --- a/datafusion/datasource-parquet/src/row_filter.rs +++ b/datafusion/datasource-parquet/src/row_filter.rs @@ -149,6 +149,12 @@ impl ArrowPredicate for DatafusionArrowPredicate { let bool_arr = as_boolean_array(&array)?.clone(); let num_matched = bool_arr.true_count(); let num_pruned = bool_arr.len() - num_matched; + tracing::debug!( + predicate = %self.physical_expr, + rows = bool_arr.len(), + rows_matched = num_matched, + "Evaluate row filter" + ); self.rows_pruned.add(num_pruned); self.rows_matched.add(num_matched); timer.stop(); From 2e64345594ca9445cee7d403c2cf720a2cc3f681 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 4 Nov 2026 19:52:56 -0300 Subject: [PATCH 14/14] prototype: capture traces with one global subscriber CI dropped spans and events at random when tests ran in parallel. tracing caches callsite interest across threads, so a test thread without a thread-local subscriber disabled callsites for the others. One global subscriber now writes into the capture that is active on the current thread. Also format the Cargo.toml files with taplo. Co-Authored-By: Claude Opus 5.5 --- Cargo.toml | 4 +- .../core/tests/parquet/io/trace_capture.rs | 86 +++++++++++++------ 2 files changed, 60 insertions(+), 30 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 9eca5e2bdcf70..5792595ae0cdb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -163,8 +163,6 @@ itertools = "0.15" itoa = "1.0" liblzma = { version = "0.4.6", features = ["static"] } log = "^0.4" -tracing = { version = "0.1.44", default-features = false, features = ["std"] } -tracing-subscriber = { version = "0.3.23", default-features = false, features = ["registry", "std"] } loom = "0.7" memchr = "2.8.1" num-traits = { version = "0.2" } @@ -180,6 +178,8 @@ pbjson-types = "0.9" percent-encoding = "2.3" pin-project = "1" pin-project-lite = "^0.2.7" +tracing = { version = "0.1.44", default-features = false, features = ["std"] } +tracing-subscriber = { version = "0.3.23", default-features = false, features = ["registry", "std"] } # Should match arrow-flight's version of prost. prost = "0.14.1" rand = "0.9" diff --git a/datafusion/core/tests/parquet/io/trace_capture.rs b/datafusion/core/tests/parquet/io/trace_capture.rs index 6370097af41c9..2cc756014ea1d 100644 --- a/datafusion/core/tests/parquet/io/trace_capture.rs +++ b/datafusion/core/tests/parquet/io/trace_capture.rs @@ -15,16 +15,21 @@ //! Read byte range (range=98..145, parquet=Bloom filter for column 'a' (rg 0)) //! ``` //! -//! The subscriber is installed with [`tracing::subscriber::set_default`], which -//! is thread-local. `#[tokio::test]` uses a current-thread runtime, so every -//! task of the query runs on the test thread and reports to this subscriber. +//! One global subscriber records into the capture that is active on the +//! current thread. `#[tokio::test]` uses a current-thread runtime, so every +//! task of the query runs on the test thread. +//! +//! A thread-local subscriber (`tracing::subscriber::set_default`) is not used. +//! `tracing` caches per-callsite interest across threads. A test thread without +//! a subscriber can then disable a callsite for every other test thread, and +//! events go missing at random when tests run in parallel. +use std::cell::RefCell; use std::fmt::{self, Write as _}; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Mutex, Once}; use tracing::field::{Field, Visit}; use tracing::span::{Attributes, Id, Record as SpanRecord}; -use tracing::subscriber::DefaultGuard; use tracing::{Event, Subscriber}; use tracing_subscriber::layer::{Context, Layer, SubscriberExt}; use tracing_subscriber::registry::{LookupSpan, Registry}; @@ -56,15 +61,25 @@ impl Visit for FieldVisitor { } } -/// Stored in a span's extensions: its index in the captured records. +type Records = Arc>>; + +thread_local! { + /// The records of the capture that is active on this thread + static ACTIVE: RefCell> = const { RefCell::new(None) }; +} + +fn active() -> Option { + ACTIVE.with(|active| active.borrow().clone()) +} + +/// Stored in a span's extensions: where the span's record lives. struct SpanSlot { + records: Records, record: usize, depth: usize, } -struct CaptureLayer { - records: Arc>>, -} +struct CaptureLayer; fn child_depth(ctx: &Context<'_, S>, parent: Option) -> usize where @@ -85,6 +100,7 @@ where if !metadata.target().starts_with(TARGET_PREFIX) { return; } + let Some(records) = active() else { return }; let parent = if attrs.is_contextual() { ctx.current_span().id().cloned() } else { @@ -94,15 +110,19 @@ where let mut visitor = FieldVisitor::default(); attrs.record(&mut visitor); - let mut records = self.records.lock().unwrap(); - records.push(Record { - depth, - name: metadata.name().to_string(), - fields: visitor.0, - }); + let record = { + let mut records = records.lock().unwrap(); + records.push(Record { + depth, + name: metadata.name().to_string(), + fields: visitor.0, + }); + records.len() - 1 + }; let span = ctx.span(id).expect("span exists"); span.extensions_mut().insert(SpanSlot { - record: records.len() - 1, + records, + record, depth, }); } @@ -115,7 +135,7 @@ where }; let mut visitor = FieldVisitor::default(); values.record(&mut visitor); - self.records.lock().unwrap()[slot.record] + slot.records.lock().unwrap()[slot.record] .fields .extend(visitor.0); } @@ -124,6 +144,7 @@ where if !event.metadata().target().starts_with(TARGET_PREFIX) { return; } + let Some(records) = active() else { return }; let depth = child_depth(&ctx, ctx.event_span(event).map(|s| s.id())); let mut visitor = FieldVisitor::default(); event.record(&mut visitor); @@ -133,7 +154,7 @@ where .position(|(name, _)| name == "message") .map(|idx| fields.remove(idx).1) .unwrap_or_default(); - self.records.lock().unwrap().push(Record { + records.lock().unwrap().push(Record { depth, name, fields, @@ -143,21 +164,24 @@ where /// Spans and events captured on this thread until it is dropped. pub(crate) struct TraceCapture { - records: Arc>>, - _guard: DefaultGuard, + records: Records, } impl TraceCapture { pub(crate) fn new() -> Self { - let records = Arc::new(Mutex::new(Vec::new())); - let subscriber = Registry::default().with(CaptureLayer { - records: Arc::clone(&records), + static INIT: Once = Once::new(); + INIT.call_once(|| { + tracing::subscriber::set_global_default( + Registry::default().with(CaptureLayer), + ) + .expect("no other global subscriber in this test binary"); }); - let guard = tracing::subscriber::set_default(subscriber); - Self { - records, - _guard: guard, - } + let records = Records::default(); + ACTIVE.with(|active| { + let previous = active.borrow_mut().replace(Arc::clone(&records)); + assert!(previous.is_none(), "one TraceCapture per thread"); + }); + Self { records } } /// Renders everything captured so far as an indented tree. @@ -180,3 +204,9 @@ impl TraceCapture { output } } + +impl Drop for TraceCapture { + fn drop(&mut self) { + ACTIVE.with(|active| active.borrow_mut().take()); + } +}