test(parquet): add I/O tracking tests for row group pruning, filter pushdown and TopK

datafusion/core/tests/parquet/io/bloom_filter_io.rs

@@ -0,0 +1,85 @@
+// 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;
+
+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::{EnabledStatistics, WriterProperties};
+
+use super::ParquetIoTest;
+
+/// Matching row group
+#[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
+    // 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
+    // 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)
+    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();
+
+    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 config = SessionConfig::new()
+        .with_target_partitions(1)
+        .with_parquet_bloom_filter_pruning(true);
+
+    ParquetIoTest::new(Bytes::from(buffer), config).await
+}

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
+}