diff --git a/Cargo.lock b/Cargo.lock index 053c7dfbd6819..f23fc8f011398 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..5792595ae0cdb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -178,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/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 e87ba4c5a7604..1c69997f1c260 100644 --- a/datafusion/core/tests/parquet/io/bloom_filter_io.rs +++ b/datafusion/core/tests/parquet/io/bloom_filter_io.rs @@ -43,22 +43,41 @@ async fn test_parquet_bloom_filter_io_matching() { // Note: in practice, the default 512 KiB metadata size hint often covers // the OffsetIndex and Bloom filters for small files, so they wouldn't // need separate requests. - 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) --- - Request 1 - Head (object size and metadata) - Request 2 - Read Parquet footer (offsets 213..678) - Request 3 - Read OffsetIndex for columns 'a', 'b' (rg 0) (offsets 192..213) - Request 4 - Read Bloom filter for column 'a' (rg 0) (offsets 98..145) - Request 5 - Read Bloom filter for column 'b' (rg 0) (offsets 145..192) - Request 6 - 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) + let traces = test.capture_traces(); + test.query("SELECT * FROM t WHERE a = 10 AND b = 100") + .await + .unwrap(); + assert_snapshot!(traces.render(), @" + Query (sql=SELECT * FROM t WHERE a = 10 AND b = 100, rows=1) + ObjectStore::head (location=table.parquet) + Metadata cache miss (cache_enabled=true) + Prefetch Parquet footer (range=213..678) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=213..678, parquet=footer) + Read page index (range=192..213) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=192..213, parquet=OffsetIndex for columns 'a', 'b' (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) "); } @@ -70,18 +89,34 @@ async fn test_parquet_bloom_filter_io_pruned() { // 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) --- - Request 1 - Head (object size and metadata) - Request 2 - Read Parquet footer (offsets 213..678) - Request 3 - Read OffsetIndex for columns 'a', 'b' (rg 0) (offsets 192..213) - Request 4 - Read Bloom filter for column 'a' (rg 0) (offsets 98..145) - Request 5 - Read Bloom filter for column 'b' (rg 0) (offsets 145..192) + let traces = test.capture_traces(); + test.query("SELECT * FROM t WHERE a = 999 AND b = 100") + .await + .unwrap(); + assert_snapshot!(traces.render(), @" + Query (sql=SELECT * FROM t WHERE a = 999 AND b = 100, rows=0) + ObjectStore::head (location=table.parquet) + Metadata cache miss (cache_enabled=true) + Prefetch Parquet footer (range=213..678) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=213..678, parquet=footer) + Read page index (range=192..213) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=192..213, parquet=OffsetIndex for columns 'a', 'b' (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 ebc5c8376fe19..720ca5c850a07 100644 --- a/datafusion/core/tests/parquet/io/filter_pushdown_io.rs +++ b/datafusion/core/tests/parquet/io/filter_pushdown_io.rs @@ -43,24 +43,42 @@ async fn test_filter_pushdown_io_disabled() { // Note: in practice, the default 512 KiB metadata size hint often covers // the OffsetIndex and Bloom filters for small files, so they wouldn't // need separate requests. - assert_snapshot!(test.query("SELECT * FROM t WHERE val = 7").await, @r" - --- Query: SELECT * FROM t WHERE val = 7 (1 rows) --- - Request 1 - Head (object size and metadata) - Request 2 - Read Parquet footer (offsets 489..1402) - Request 3 - Read ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0, 1, 2) (offsets 286..489) - 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) - Request 5 - 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 6 - 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) + let traces = test.capture_traces(); + test.query("SELECT * FROM t WHERE val = 7").await.unwrap(); + assert_snapshot!(traces.render(), @" + Query (sql=SELECT * FROM t WHERE val = 7, rows=1) + ObjectStore::head (location=table.parquet) + Metadata cache miss (cache_enabled=true) + Prefetch Parquet footer (range=489..1402) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=489..1402, parquet=footer) + Read page index (range=286..489) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=286..489, parquet=ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0, 1, 2)) + 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) "); } @@ -70,24 +88,46 @@ 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" - --- Query: SELECT * FROM t WHERE val = 7 (1 rows) --- - Request 1 - Head (object size and metadata) - Request 2 - Read Parquet footer (offsets 489..1402) - Request 3 - Read ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0, 1, 2) (offsets 286..489) - Request 4 - Read dictionary page and data page 0 for column 'val' (rg 0) (offsets 51..98) - Request 5 - Read dictionary page and data page 0 for column 'val' (rg 1) (offsets 145..192) - Request 6 - 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) - Request 7 - Read dictionary page and data page 0 for column 'val' (rg 2) (offsets 239..286) + let traces = test.capture_traces(); + test.query("SELECT * FROM t WHERE val = 7").await.unwrap(); + assert_snapshot!(traces.render(), @" + Query (sql=SELECT * FROM t WHERE val = 7, rows=1) + ObjectStore::head (location=table.parquet) + Metadata cache miss (cache_enabled=true) + Prefetch Parquet footer (range=489..1402) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=489..1402, parquet=footer) + Read page index (range=286..489) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=286..489, parquet=ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0, 1, 2)) + 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 0a9e9a47ea289..47a257356d25d 100644 --- a/datafusion/core/tests/parquet/io/mod.rs +++ b/datafusion/core/tests/parquet/io/mod.rs @@ -8,16 +8,15 @@ //! correlates each requested byte range with the Parquet file structure it //! covers, so the resulting trace can be compared against `insta` snapshots. //! -//! `ListingTable` object store access tests live in +//! `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::atomic::{AtomicUsize, Ordering}; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; 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; @@ -31,6 +30,7 @@ use parquet::file::FOOTER_SIZE; use parquet::file::metadata::FooterTail; use parquet::file::reader::{FileReader, SerializedFileReader}; use parquet::file::serialized_reader::ReadOptionsBuilder; +use tracing::Instrument; use url::Url; // Tests are in submodules @@ -39,6 +39,9 @@ mod filter_pushdown_io; mod page_pruning_io; mod row_group_pruning_io; mod topk_io; +mod trace_capture; + +use trace_capture::TraceCapture; /// Kind of per column chunk metadata structure stored outside the footer #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -278,7 +281,7 @@ impl ParquetFileLayout { fn describe_range(&self, range: &Range, file_size: u64) -> String { // The footer is read as a suffix of the file (see `footer_size`) 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(); } let overlaps = |r: &Range| range.start < r.end && range.end > r.start; // Describe the overlapping Bloom filters / page indexes @@ -295,70 +298,42 @@ impl ParquetFileLayout { .map(|chunk| chunk.describe(range)), ); if parts.is_empty() { - parts.push("byte range".to_string()); + parts.push("unknown".to_string()); } - format!( - "{} (offsets {}..{})", - parts.join(", "), - range.start, - range.end - ) + parts.join(", ") } } +/// 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, - events: Arc>>, - /// Number of requests made to the store (reset for each query) - requests: Arc, } -impl InstrumentedObjectStore { - fn new( - inner: Arc, - layout: ParquetFileLayout, - file_size: u64, - events: Arc>>, - requests: Arc, - ) -> Self { - Self { - inner, - layout, - file_size, - events, - 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"; - /// Records a request to the store and the byte ranges it reads. - fn record_request(&self, ranges: impl IntoIterator>) { - let mut events = self.lock_new_request(); +macro_rules! store_span { + ($op:literal, $($fields:tt)*) => { + tracing::debug_span!(target: TEST_TARGET, $op, $($fields)*) + }; +} + +impl InstrumentedObjectStore { + fn record_ranges(&self, ranges: impl IntoIterator>) { for range in ranges { - let desc = self.layout.describe_range(&range, self.file_size); - events.push(format!(" Read {desc}")); + tracing::debug!( + target: TEST_TARGET, + range = ?range, + parquet = %self.layout.describe_range(&range, self.file_size), + "Read byte range" + ); } } - - /// Records a HEAD request (object metadata only, no content) - fn record_head_request(&self) { - self.lock_new_request() - .push(" Head (object size and metadata)".to_string()); - } - - /// Pushes a new `Request N` line and returns the held events lock so - /// the caller can append the lines of that request. - /// - /// Holding the lock while appending keeps the lines of concurrent - /// requests (for example, from multiple partitions) from interleaving. - fn lock_new_request(&self) -> std::sync::MutexGuard<'_, Vec> { - let request = self.requests.fetch_add(1, Ordering::SeqCst) + 1; - let mut events = self.events.lock().unwrap(); - events.push(format!("Request {request}")); - events - } } impl std::fmt::Display for InstrumentedObjectStore { @@ -375,7 +350,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( @@ -383,7 +363,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( @@ -392,19 +376,27 @@ 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.start..r.end, + 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 + 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( @@ -412,14 +404,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) } @@ -427,6 +425,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) } @@ -435,6 +434,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) } @@ -442,7 +447,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( @@ -451,23 +463,28 @@ 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 } } /// 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>>, - requests: Arc, } impl ParquetIoTest { @@ -494,15 +511,11 @@ 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, + let instrumented_store = Arc::new(InstrumentedObjectStore { + inner: memory_store, layout, file_size, - Arc::clone(&events), - Arc::clone(&requests), - )); + }); let ctx = SessionContext::new_with_config(config); ctx.register_object_store(&Url::parse("mem://").unwrap(), instrumented_store); @@ -510,48 +523,47 @@ impl ParquetIoTest { .await .unwrap(); - Self { - ctx, - events, - requests, - } + Self { ctx } } - /// 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. + /// 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() + } + + /// Runs `sql` to completion inside a `query` span and returns the number + /// of output rows. Emits a `Produce batch` event for each output batch. /// /// The file metadata cache is cleared first so each query shows the I/O /// of a cold read, including the footer and page index. - pub(crate) async fn query(&self, sql: &str) -> String { + pub(crate) async fn query(&self, sql: &str) -> Result { self.ctx .runtime_env() .cache_manager .get_file_metadata_cache() .clear(); - 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(); - - 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 span = tracing::debug_span!( + target: TEST_TARGET, + "Query", + sql, + rows = tracing::field::Empty + ); + 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) } - output + .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 cda10224325f0..ed0dd3cfff279 100644 --- a/datafusion/core/tests/parquet/io/page_pruning_io.rs +++ b/datafusion/core/tests/parquet/io/page_pruning_io.rs @@ -46,18 +46,30 @@ async fn test_page_pruning_io_no_predicate() { let test = page_pruning_test(true).await; // Note: in practice, the default 512 KiB metadata size hint would fetch // the page index of a small file like this along with the footer. - assert_snapshot!(test.query("SELECT * FROM t").await, @r" - --- Query: SELECT * FROM t (9 rows) --- - Request 1 - Head (object size and metadata) - Request 2 - Read Parquet footer (offsets 323..838) - Request 3 - Read ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0) (offsets 178..323) - Request 4 - 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) + let traces = test.capture_traces(); + test.query("SELECT * FROM t").await.unwrap(); + assert_snapshot!(traces.render(), @" + Query (sql=SELECT * FROM t, rows=9) + ObjectStore::head (location=table.parquet) + Metadata cache miss (cache_enabled=true) + Prefetch Parquet footer (range=323..838) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=323..838, parquet=footer) + Read page index (range=178..323) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=178..323, parquet=ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0)) + 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) "); } @@ -66,18 +78,34 @@ 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" - --- Query: SELECT * FROM t WHERE id = 5 (1 rows) --- - Request 1 - Head (object size and metadata) - Request 2 - Read Parquet footer (offsets 323..838) - Request 3 - Read ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0) (offsets 178..323) - Request 4 - 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) + let traces = test.capture_traces(); + test.query("SELECT * FROM t WHERE id = 5").await.unwrap(); + assert_snapshot!(traces.render(), @" + Query (sql=SELECT * FROM t WHERE id = 5, rows=1) + ObjectStore::head (location=table.parquet) + Metadata cache miss (cache_enabled=true) + Prefetch Parquet footer (range=323..838) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=323..838, parquet=footer) + Read page index (range=178..323) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=178..323, parquet=ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0)) + 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) "); } @@ -86,20 +114,36 @@ 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" - --- Query: SELECT * FROM t WHERE id > 3 (6 rows) --- - Request 1 - Head (object size and metadata) - Request 2 - Read Parquet footer (offsets 323..838) - Request 3 - Read ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0) (offsets 178..323) - Request 4 - 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) + let traces = test.capture_traces(); + test.query("SELECT * FROM t WHERE id > 3").await.unwrap(); + assert_snapshot!(traces.render(), @" + Query (sql=SELECT * FROM t WHERE id > 3, rows=6) + ObjectStore::head (location=table.parquet) + Metadata cache miss (cache_enabled=true) + Prefetch Parquet footer (range=323..838) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=323..838, parquet=footer) + Read page index (range=178..323) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=178..323, parquet=ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0)) + 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) "); } @@ -109,18 +153,31 @@ 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" - --- Query: SELECT * FROM t WHERE id = 5 (1 rows) --- - Request 1 - Head (object size and metadata) - Request 2 - Read Parquet footer (offsets 323..838) - Request 3 - Read ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0) (offsets 178..323) - Request 4 - 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) + let traces = test.capture_traces(); + test.query("SELECT * FROM t WHERE id = 5").await.unwrap(); + assert_snapshot!(traces.render(), @" + Query (sql=SELECT * FROM t WHERE id = 5, rows=1) + ObjectStore::head (location=table.parquet) + Metadata cache miss (cache_enabled=true) + Prefetch Parquet footer (range=323..838) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=323..838, parquet=footer) + Read page index (range=178..323) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=178..323, parquet=ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0)) + 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 b6ac9cd6c5338..eaeaf8747f341 100644 --- a/datafusion/core/tests/parquet/io/row_group_pruning_io.rs +++ b/datafusion/core/tests/parquet/io/row_group_pruning_io.rs @@ -40,26 +40,40 @@ async fn test_row_group_pruning_io_no_predicate() { let test = row_group_pruning_test().await; // Note: in practice, the default 512 KiB metadata size hint would fetch // the page index of a small file like this along with the footer. - assert_snapshot!(test.query("SELECT * FROM t").await, @r" - --- Query: SELECT * FROM t (9 rows) --- - Request 1 - Head (object size and metadata) - Request 2 - Read Parquet footer (offsets 489..1402) - Request 3 - Read ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0, 1, 2) (offsets 286..489) - 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 (3 rows) - Request 5 - 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) - Request 6 - 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) + let traces = test.capture_traces(); + test.query("SELECT * FROM t").await.unwrap(); + assert_snapshot!(traces.render(), @" + Query (sql=SELECT * FROM t, rows=9) + ObjectStore::head (location=table.parquet) + Metadata cache miss (cache_enabled=true) + Prefetch Parquet footer (range=489..1402) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=489..1402, parquet=footer) + Read page index (range=286..489) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=286..489, parquet=ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0, 1, 2)) + 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) "); } @@ -68,18 +82,34 @@ 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" - --- Query: SELECT * FROM t WHERE id = 5 (1 rows) --- - Request 1 - Head (object size and metadata) - Request 2 - Read Parquet footer (offsets 489..1402) - Request 3 - Read ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0, 1, 2) (offsets 286..489) - Request 4 - 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) + let traces = test.capture_traces(); + test.query("SELECT * FROM t WHERE id = 5").await.unwrap(); + assert_snapshot!(traces.render(), @" + Query (sql=SELECT * FROM t WHERE id = 5, rows=1) + ObjectStore::head (location=table.parquet) + Metadata cache miss (cache_enabled=true) + Prefetch Parquet footer (range=489..1402) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=489..1402, parquet=footer) + Read page index (range=286..489) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=286..489, parquet=ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0, 1, 2)) + 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) "); } @@ -88,21 +118,38 @@ 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" - --- Query: SELECT * FROM t WHERE id > 5 (4 rows) --- - Request 1 - Head (object size and metadata) - Request 2 - Read Parquet footer (offsets 489..1402) - Request 3 - Read ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0, 1, 2) (offsets 286..489) - Request 4 - 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 5 - 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) + let traces = test.capture_traces(); + test.query("SELECT * FROM t WHERE id > 5").await.unwrap(); + assert_snapshot!(traces.render(), @" + Query (sql=SELECT * FROM t WHERE id > 5, rows=4) + ObjectStore::head (location=table.parquet) + Metadata cache miss (cache_enabled=true) + Prefetch Parquet footer (range=489..1402) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=489..1402, parquet=footer) + Read page index (range=286..489) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=286..489, parquet=ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0, 1, 2)) + 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) "); } @@ -111,14 +158,19 @@ 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" - --- Query: SELECT * FROM t WHERE id > 100 (0 rows) --- - Request 1 - Head (object size and metadata) - Request 2 - Read Parquet footer (offsets 489..1402) - Request 3 - Read ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0, 1, 2) (offsets 286..489) + let traces = test.capture_traces(); + test.query("SELECT * FROM t WHERE id > 100").await.unwrap(); + assert_snapshot!(traces.render(), @" + Query (sql=SELECT * FROM t WHERE id > 100, rows=0) + ObjectStore::head (location=table.parquet) + Metadata cache miss (cache_enabled=true) + Prefetch Parquet footer (range=489..1402) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=489..1402, parquet=footer) + Read page index (range=286..489) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=286..489, parquet=ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0, 1, 2)) + 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 3636e71f61a90..4a5ac28413a76 100644 --- a/datafusion/core/tests/parquet/io/topk_io.rs +++ b/datafusion/core/tests/parquet/io/topk_io.rs @@ -46,18 +46,32 @@ async fn test_topk_io_dynamic_filter() { let test = topk_test(true).await; // Note: in practice, the default 512 KiB metadata size hint would fetch // the page index of a small file like this along with the footer. - 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) --- - Request 1 - Head (object size and metadata) - Request 2 - Read Parquet footer (offsets 489..1402) - Request 3 - Read ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0, 1, 2) (offsets 286..489) - 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 (2 rows) + let traces = test.capture_traces(); + test.query("SELECT * FROM t ORDER BY id DESC LIMIT 2") + .await + .unwrap(); + assert_snapshot!(traces.render(), @" + Query (sql=SELECT * FROM t ORDER BY id DESC LIMIT 2, rows=2) + ObjectStore::head (location=table.parquet) + Metadata cache miss (cache_enabled=true) + Prefetch Parquet footer (range=489..1402) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=489..1402, parquet=footer) + Read page index (range=286..489) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=286..489, parquet=ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0, 1, 2)) + 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) "); } @@ -66,24 +80,40 @@ 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" - --- Query: SELECT * FROM t ORDER BY id DESC LIMIT 2 (2 rows) --- - Request 1 - Head (object size and metadata) - Request 2 - Read Parquet footer (offsets 489..1402) - Request 3 - Read ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0, 1, 2) (offsets 286..489) - 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) - Request 5 - 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 6 - 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) + let traces = test.capture_traces(); + test.query("SELECT * FROM t ORDER BY id DESC LIMIT 2") + .await + .unwrap(); + assert_snapshot!(traces.render(), @" + Query (sql=SELECT * FROM t ORDER BY id DESC LIMIT 2, rows=2) + ObjectStore::head (location=table.parquet) + Metadata cache miss (cache_enabled=true) + Prefetch Parquet footer (range=489..1402) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=489..1402, parquet=footer) + Read page index (range=286..489) + ObjectStore::get_ranges (location=table.parquet, ranges=1) + Read byte range (range=286..489, parquet=ColumnIndex and OffsetIndex for columns 'id', 'val' (rg 0, 1, 2)) + 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 new file mode 100644 index 0000000000000..2cc756014ea1d --- /dev/null +++ b/datafusion/core/tests/parquet/io/trace_capture.rs @@ -0,0 +1,212 @@ +// 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 an indented tree for `insta` snapshots. +//! +//! 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)) +//! ``` +//! +//! 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, Once}; + +use tracing::field::{Field, Visit}; +use tracing::span::{Attributes, Id, Record as SpanRecord}; +use tracing::{Event, Subscriber}; +use tracing_subscriber::layer::{Context, Layer, SubscriberExt}; +use tracing_subscriber::registry::{LookupSpan, Registry}; + +/// 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 where it opens) or event. +#[derive(Debug, Clone)] +struct Record { + depth: usize, + /// The span name, or the event's `message` field + name: String, + fields: Vec<(String, String)>, +} + +#[derive(Default)] +struct FieldVisitor(Vec<(String, String)>); + +impl Visit for FieldVisitor { + fn record_str(&mut self, field: &Field, value: &str) { + self.0.push((field.name().to_string(), value.to_string())); + } + + fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) { + self.0 + .push((field.name().to_string(), format!("{value:?}"))); + } +} + +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; + +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 +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 Some(records) = active() else { return }; + let parent = if attrs.is_contextual() { + ctx.current_span().id().cloned() + } else { + attrs.parent().cloned() + }; + let depth = child_depth(&ctx, parent); + let mut visitor = FieldVisitor::default(); + attrs.record(&mut visitor); + + 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 { + records, + record, + 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); + slot.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 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); + let mut fields = visitor.0; + let name = fields + .iter() + .position(|(name, _)| name == "message") + .map(|idx| fields.remove(idx).1) + .unwrap_or_default(); + records.lock().unwrap().push(Record { + depth, + name, + fields, + }); + } +} + +/// Spans and events captured on this thread until it is dropped. +pub(crate) struct TraceCapture { + records: Records, +} + +impl TraceCapture { + pub(crate) fn new() -> Self { + 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 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. + pub(crate) fn render(&self) -> String { + let records = self.records.lock().unwrap(); + let mut output = String::new(); + for record in records.iter() { + let indent = " ".repeat(record.depth); + 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(", ")); + } + output.push('\n'); + } + output + } +} + +impl Drop for TraceCapture { + fn drop(&mut self) { + ACTIVE.with(|active| active.borrow_mut().take()); + } +} 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..aac54274374c6 100644 --- a/datafusion/datasource-parquet/src/file_format.rs +++ b/datafusion/datasource-parquet/src/file_format.rs @@ -774,6 +774,7 @@ impl<'a> ObjectStoreFetch<'a> { impl MetadataFetch for ObjectStoreFetch<'_> { fn fetch(&mut self, range: Range) -> BoxFuture<'_, Result> { + 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 1a3778b31ee72..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,6 +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!(range = ?prefetch_range, "Prefetch Parquet footer"); let data = self .store .get_ranges( @@ -381,6 +391,15 @@ 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!(range = ?range, "Read Parquet footer"); + } else { + tracing::debug!(range = ?range, "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..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, @@ -1547,6 +1575,17 @@ 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!( + column = %column_name, + row_group = idx, + range = ?(offset..offset + i64::from(length)), + "Read Bloom filter" + ); + } let bf: Sbbf = match builder .get_row_group_column_bloom_filter(idx, *column_idx) .await @@ -1589,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 @@ -1603,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 @@ -1648,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, @@ -1662,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 57df6190daf87..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,9 +502,14 @@ impl PushDecoderStreamState { Ok(DecodeResult::NeedsData(ranges)) => { // I/O, not compute. timer.stop(); + 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();