prototype: capture Parquet scan operations with tracing in I/O tests

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"

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"

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)
     ");
 }
 

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)
     ");
 }