diff --git a/datafusion/core/tests/datasource/object_store_access.rs b/datafusion/core/tests/datasource/object_store_access.rs index e44e14df213e8..d9dffa33bda3f 100644 --- a/datafusion/core/tests/datasource/object_store_access.rs +++ b/datafusion/core/tests/datasource/object_store_access.rs @@ -9,6 +9,9 @@ //! various operations (table creation, queries with and without predicates) //! to verify the expected object store access patterns. //! +//! File-level Parquet access tests live in +//! [`datafusion/core/tests/parquet/io`](https://github.com/apache/datafusion/blob/main/datafusion/core/tests/parquet/io/mod.rs). +//! //! [`ListingTable`]: datafusion::datasource::listing::ListingTable use arrow::array::{ArrayRef, Int32Array, RecordBatch}; 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..bc86e0e699b1b --- /dev/null +++ b/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 +} diff --git a/datafusion/core/tests/parquet/io/mod.rs b/datafusion/core/tests/parquet/io/mod.rs new file mode 100644 index 0000000000000..839d7a981051e --- /dev/null +++ b/datafusion/core/tests/parquet/io/mod.rs @@ -0,0 +1,326 @@ +// 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 + +//! Test infrastructure for tracking storage I/O events during Parquet scans. +//! +//! Provides an instrumented [`ObjectStore`] that records storage requests and +//! 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 +//! [`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, Mutex}; + +use async_trait::async_trait; +use bytes::Bytes; +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; + +// Tests are in submodules +mod bloom_filter_io; + +#[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 + } +} + +/// 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 } + } + + /// 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 + } +} diff --git a/datafusion/core/tests/parquet/mod.rs b/datafusion/core/tests/parquet/mod.rs index be7cce00ddab4..26d0829944c75 100644 --- a/datafusion/core/tests/parquet/mod.rs +++ b/datafusion/core/tests/parquet/mod.rs @@ -40,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;