diff --git a/datafusion/core/tests/parquet/io/bloom_filter_io.rs b/datafusion/core/tests/parquet/io/bloom_filter_io.rs index bc86e0e699b1b..e87ba4c5a7604 100644 --- a/datafusion/core/tests/parquet/io/bloom_filter_io.rs +++ b/datafusion/core/tests/parquet/io/bloom_filter_io.rs @@ -3,6 +3,24 @@ // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements //! Tests tracking storage I/O events during Parquet Bloom filter evaluation. +//! +//! The file has a single row group with Bloom filters on both columns and no +//! statistics, so only the Bloom filters are used to prune. +//! +//! ```text +//! row group 0 +//! ┌──────┬───────┐ +//! │ a │ b │ +//! ├──────┼───────┤ +//! │ 10 │ 100 │ +//! │ 20 │ 200 │ +//! │ 30 │ 300 │ +//! └──────┴───────┘ +//! Bloom filter (a) Bloom filter (b) no min/max statistics +//! +//! a = 10 AND b = 100 -> Bloom filters match, row group read +//! a = 999 AND b = 100 -> Bloom filter on `a` misses, row group pruned +//! ``` use std::sync::Arc; @@ -21,14 +39,25 @@ 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. + // + // 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) --- - 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) + 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) "); } @@ -43,10 +72,16 @@ async fn test_parquet_bloom_filter_io_pruned() { // evaluated first and the Bloom filter for `b` would not be loaded. assert_snapshot!(test.query("SELECT * FROM t WHERE a = 999 AND b = 100").await, @r" --- Query: SELECT * FROM t WHERE a = 999 AND b = 100 (0 rows) --- - Read Parquet footer (offsets 0..678) - Read OffsetIndex for column 'a' (rg 0) (offsets 192..213) - Read Bloom filter for column 'a' (rg 0) (offsets 98..145) - Read Bloom filter for column 'b' (rg 0) (offsets 145..192) + Request 1 + 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) "); } diff --git a/datafusion/core/tests/parquet/io/filter_pushdown_io.rs b/datafusion/core/tests/parquet/io/filter_pushdown_io.rs new file mode 100644 index 0000000000000..ebc5c8376fe19 --- /dev/null +++ b/datafusion/core/tests/parquet/io/filter_pushdown_io.rs @@ -0,0 +1,118 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! Tests tracking I/O events with and without Parquet filter pushdown +//! (`datafusion.execution.parquet.pushdown_filters`, see +//! [`ParquetOptions::pushdown_filters`]). +//! +//! Both tests run `val = 7`, which statistics cannot prune, but which only +//! matches a row in row group 1: +//! +//! ```text +//! row group 0 row group 1 row group 2 +//! ┌──────┬───────┐ ┌──────┬───────┐ ┌──────┬───────┐ +//! │ id │ val │ │ id │ val │ │ id │ val │ +//! ├──────┼───────┤ ├──────┼───────┤ ├──────┼───────┤ +//! │ 1 │ 0 │ │ 4 │ 0 │ │ 7 │ 0 │ +//! │ 2 │ 5 │ │ 5 │ 7 * │ │ 8 │ 5 │ +//! │ 3 │ 10 │ │ 6 │ 10 │ │ 9 │ 10 │ +//! └──────┴───────┘ └──────┴───────┘ └──────┴───────┘ +//! val: min 0 max 10 val: min 0 max 10 val: min 0 max 10 +//! +//! * the only row matching `val = 7` +//! ``` +//! +//! [`ParquetOptions::pushdown_filters`]: datafusion::config::ParquetOptions::pushdown_filters + +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int32Array, RecordBatch}; +use bytes::Bytes; +use datafusion::prelude::SessionConfig; +use insta::assert_snapshot; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::WriterProperties; + +use super::ParquetIoTest; + +/// Without filter pushdown, the scan fetches every column of every row group. +#[tokio::test] +async fn test_filter_pushdown_io_disabled() { + let test = filter_pushdown_test(false).await; + // 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) + "); +} + +/// **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). +#[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) + "); +} + +/// Returns a [`ParquetIoTest`] for the file described in the module docs +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 +} + +/// Returns the Parquet file described in the module docs: three row groups of +/// three rows each, with `id` sorted across row groups. +fn three_row_group_file() -> Bytes { + let id: ArrayRef = Arc::new(Int32Array::from_iter_values(1..=9)); + let val: ArrayRef = Arc::new(Int32Array::from(vec![0, 5, 10, 0, 7, 10, 0, 5, 10])); + let batch = RecordBatch::try_from_iter([("id", id), ("val", val)]).unwrap(); + + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .build(); + + let mut buffer = Vec::new(); + let mut writer = + ArrowWriter::try_new(&mut buffer, batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + Bytes::from(buffer) +} diff --git a/datafusion/core/tests/parquet/io/mod.rs b/datafusion/core/tests/parquet/io/mod.rs index 839d7a981051e..0a9e9a47ea289 100644 --- a/datafusion/core/tests/parquet/io/mod.rs +++ b/datafusion/core/tests/parquet/io/mod.rs @@ -8,11 +8,12 @@ //! 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 async_trait::async_trait; @@ -26,32 +27,173 @@ use object_store::{ GetOptions, GetRange, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, }; +use parquet::file::FOOTER_SIZE; +use parquet::file::metadata::FooterTail; use parquet::file::reader::{FileReader, SerializedFileReader}; +use parquet::file::serialized_reader::ReadOptionsBuilder; use url::Url; // Tests are in submodules mod bloom_filter_io; +mod filter_pushdown_io; +mod page_pruning_io; +mod row_group_pruning_io; +mod topk_io; + +/// Kind of per column chunk metadata structure stored outside the footer +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum MetadataKind { + BloomFilter, + ColumnIndex, + OffsetIndex, +} +impl std::fmt::Display for MetadataKind { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::BloomFilter => write!(f, "Bloom filter"), + Self::ColumnIndex => write!(f, "ColumnIndex"), + Self::OffsetIndex => write!(f, "OffsetIndex"), + } + } +} + +/// Byte range of a Bloom filter, ColumnIndex, or OffsetIndex of one column chunk #[derive(Debug, Clone)] struct ByteRangeDescriptor { range: Range, - description: String, + kind: MetadataKind, + col_name: String, + rg_idx: usize, +} + +/// Describes the Bloom filters / page indexes in `descriptors` as +/// ` for column(s) '', '' (rg , )`. +/// +/// Kinds that cover the same columns and row groups are merged into one +/// clause (for example `ColumnIndex and OffsetIndex for columns 'id', 'val' +/// (rg 0, 1, 2)`). +fn describe_descriptors(descriptors: &[&ByteRangeDescriptor]) -> Vec { + // (columns and row groups covered) -> kinds, in order of first appearance + let mut groups: Vec<(String, Vec)> = vec![]; + for kind in [ + MetadataKind::BloomFilter, + MetadataKind::ColumnIndex, + MetadataKind::OffsetIndex, + ] { + let of_kind: Vec<&ByteRangeDescriptor> = descriptors + .iter() + .copied() + .filter(|d| d.kind == kind) + .collect(); + if of_kind.is_empty() { + continue; + } + // Column names in file order, without repeating them for each row group + let mut col_names: Vec<&str> = vec![]; + for d in &of_kind { + if !col_names.contains(&d.col_name.as_str()) { + col_names.push(&d.col_name); + } + } + let mut rg_idxs: Vec = of_kind.iter().map(|d| d.rg_idx).collect(); + rg_idxs.sort_unstable(); + rg_idxs.dedup(); + let columns: Vec = col_names.iter().map(|c| format!("'{c}'")).collect(); + let rgs: Vec = rg_idxs.iter().map(|rg| rg.to_string()).collect(); + let covered = format!( + "for column{} {} (rg {})", + if columns.len() == 1 { "" } else { "s" }, + columns.join(", "), + rgs.join(", ") + ); + match groups.iter_mut().find(|(c, _)| *c == covered) { + Some((_, kinds)) => kinds.push(kind), + None => groups.push((covered, vec![kind])), + } + } + groups + .into_iter() + .map(|(covered, kinds)| { + let kinds: Vec = kinds.iter().map(|k| k.to_string()).collect(); + format!("{} {covered}", kinds.join(" and ")) + }) + .collect() +} + +/// Byte ranges of one column chunk: +/// +/// 1. (optional) dictionary page +/// 2. data pages (from the offset index, if present) +#[derive(Debug, Clone)] +struct ColumnChunkLayout { + col_name: String, + /// Row Group Index + rg_idx: usize, + /// The whole column chunk (dictionary page + data pages) + range: Range, + /// Byte range of the dictionary page + dictionary_page: Option>, + /// Byte range of each data page, in page order (empty without an offset index) + data_pages: Vec>, +} + +impl ColumnChunkLayout { + /// Describes `range` in terms of the pages of this column chunk it covers + fn describe(&self, range: &Range) -> String { + let overlaps = |r: &Range| range.start < r.end && range.end > r.start; + let mut parts = vec![]; + if self.dictionary_page.as_ref().is_some_and(overlaps) { + parts.push("dictionary page".to_string()); + } + let pages: Vec = self + .data_pages + .iter() + .enumerate() + .filter(|(_, r)| overlaps(r)) + .map(|(idx, _)| idx) + .collect(); + match pages.as_slice() { + [] if self.data_pages.is_empty() => parts.push("data pages".to_string()), + [] => {} + [page] => parts.push(format!("data page {page}")), + [first, .., last] if last - first + 1 == pages.len() => { + parts.push(format!("data pages {first}..={last}")) + } + pages => { + let pages: Vec = pages.iter().map(|p| p.to_string()).collect(); + parts.push(format!("data pages {}", pages.join(", "))) + } + } + format!( + "{} for column '{}' (rg {})", + parts.join(" and "), + self.col_name, + self.rg_idx + ) + } } #[derive(Debug, Clone, Default)] struct ParquetFileLayout { + /// Bloom filter and page index byte ranges descriptors: Vec, + column_chunks: Vec, } impl ParquetFileLayout { fn new(bytes: &Bytes) -> Self { let cursor = bytes.clone(); - let reader = SerializedFileReader::new(cursor).expect("valid parquet file"); + let options = ReadOptionsBuilder::new().with_page_index().build(); + let reader = SerializedFileReader::new_with_options(cursor, options) + .expect("valid parquet file"); let metadata = reader.metadata(); let mut descriptors = Vec::new(); + let mut column_chunks = Vec::new(); for (rg_idx, rg) in metadata.row_groups().iter().enumerate() { - for col in rg.columns() { + let page_index = metadata.page_index_for_row_group(rg_idx); + for (col_idx, col) in rg.columns().iter().enumerate() { let col_name = col.column_descr().name(); // Bloom filter byte range @@ -62,9 +204,9 @@ impl ParquetFileLayout { let end = start + bf_len as u64; descriptors.push(ByteRangeDescriptor { range: start..end, - description: format!( - "Bloom filter for column '{col_name}' (rg {rg_idx})" - ), + kind: MetadataKind::BloomFilter, + col_name: col_name.to_string(), + rg_idx, }); } @@ -72,51 +214,95 @@ impl ParquetFileLayout { if let Some(r) = col.column_index_range() { descriptors.push(ByteRangeDescriptor { range: r, - description: format!( - "ColumnIndex for column '{col_name}' (rg {rg_idx})" - ), + kind: MetadataKind::ColumnIndex, + col_name: col_name.to_string(), + 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})" - ), + kind: MetadataKind::OffsetIndex, + col_name: col_name.to_string(), + rg_idx, }); } - // Data pages / column chunk byte range - let data_start = col + // Dictionary page and data pages of the column chunk + let dictionary_start = col .dictionary_page_offset() - .and_then(|o| u64::try_from(o).ok()) - .unwrap_or_else(|| col.data_page_offset() as u64); - let data_end = data_start + col.compressed_size() as u64; - descriptors.push(ByteRangeDescriptor { - range: data_start..data_end, - description: format!( - "data page for column '{col_name}' (rg {rg_idx})" - ), + .and_then(|o| u64::try_from(o).ok()); + let data_pages: Vec> = page_index + .offset_index(col_idx) + .map(|offset_index| { + offset_index + .page_locations() + .iter() + .map(|location| { + let start = location.offset as u64; + start..start + location.compressed_page_size as u64 + }) + .collect() + }) + .unwrap_or_default(); + let dictionary_page = dictionary_start.map(|start| { + let end = data_pages + .first() + .map(|page| page.start) + .unwrap_or_else(|| col.data_page_offset() as u64); + start..end + }); + let chunk_start = + dictionary_start.unwrap_or_else(|| col.data_page_offset() as u64); + let chunk_end = chunk_start + col.compressed_size() as u64; + column_chunks.push(ColumnChunkLayout { + col_name: col_name.to_string(), + rg_idx, + range: chunk_start..chunk_end, + dictionary_page, + data_pages, }); } } - Self { descriptors } + Self { + descriptors, + column_chunks, + } } + /// Describes `range` in terms of the parts of the Parquet file it covers. + /// + /// A single request may span several items (for example, a coalesced read + /// of adjacent column chunks), so all overlapping items are listed. 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); } - for desc in &self.descriptors { - if range.start < desc.range.end && range.end > desc.range.start { - return format!( - "{} (offsets {}..{})", - desc.description, range.start, range.end - ); - } + let overlaps = |r: &Range| range.start < r.end && range.end > r.start; + // Describe the overlapping Bloom filters / page indexes + let descriptors: Vec<&ByteRangeDescriptor> = self + .descriptors + .iter() + .filter(|desc| overlaps(&desc.range)) + .collect(); + let mut parts = describe_descriptors(&descriptors); + parts.extend( + self.column_chunks + .iter() + .filter(|chunk| overlaps(&chunk.range)) + .map(|chunk| chunk.describe(range)), + ); + if parts.is_empty() { + parts.push("byte range".to_string()); } - format!("byte range (offsets {}..{})", range.start, range.end) + format!( + "{} (offsets {}..{})", + parts.join(", "), + range.start, + range.end + ) } } @@ -126,6 +312,8 @@ struct InstrumentedObjectStore { layout: ParquetFileLayout, file_size: u64, events: Arc>>, + /// Number of requests made to the store (reset for each query) + requests: Arc, } impl InstrumentedObjectStore { @@ -134,18 +322,42 @@ impl InstrumentedObjectStore { layout: ParquetFileLayout, file_size: u64, events: Arc>>, + requests: Arc, ) -> Self { Self { inner, layout, file_size, events, + requests, } } - fn record_range(&self, range: &Range) { - let desc = self.layout.describe_range(range, self.file_size); - self.events.lock().unwrap().push(format!("Read {desc}")); + /// 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(); + for range in ranges { + let desc = self.layout.describe_range(&range, self.file_size); + events.push(format!(" Read {desc}")); + } + } + + /// 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 } } @@ -179,22 +391,19 @@ impl ObjectStore for InstrumentedObjectStore { 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)); + if options.head { + self.record_head_request(); + return self.inner.get_opts(location, options).await; } + let range = match &options.range { + Some(GetRange::Bounded(r)) => r.start..r.end, + 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 } @@ -203,9 +412,7 @@ impl ObjectStore for InstrumentedObjectStore { location: &Path, ranges: &[Range], ) -> object_store::Result> { - for range in ranges { - self.record_range(range); - } + self.record_request(ranges.iter().cloned()); self.inner.get_ranges(location, ranges).await } @@ -260,14 +467,22 @@ impl ObjectStore for InstrumentedObjectStore { pub(crate) struct ParquetIoTest { ctx: SessionContext, events: Arc>>, + requests: 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 { + /// + /// Note: sets the metadata size hint to the exact size of the footer so + /// the first request reads only the footer metadata. Otherwise the default + /// hint (512 KiB) covers the (small) test files entirely, hiding requests + /// for the page index and Bloom filters. + pub(crate) async fn new(file_bytes: Bytes, mut config: SessionConfig) -> Self { let file_size = file_bytes.len() as u64; let layout = ParquetFileLayout::new(&file_bytes); + config.options_mut().execution.parquet.metadata_size_hint = + Some(footer_size(&file_bytes)); let memory_store = Arc::new(InMemory::new()); memory_store @@ -280,11 +495,13 @@ impl ParquetIoTest { .unwrap(); let events = Arc::new(Mutex::new(Vec::new())); + let requests = Arc::new(AtomicUsize::new(0)); let instrumented_store = Arc::new(InstrumentedObjectStore::new( memory_store, layout, file_size, Arc::clone(&events), + Arc::clone(&requests), )); let ctx = SessionContext::new_with_config(config); @@ -293,15 +510,28 @@ impl ParquetIoTest { .await .unwrap(); - Self { ctx, events } + Self { + ctx, + events, + requests, + } } /// Runs `sql` and returns the storage requests and output batches it /// produced, formatted for snapshot comparison. /// /// Only events recorded while running this query are included. + /// + /// 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 { + 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(); @@ -324,3 +554,12 @@ impl ParquetIoTest { output } } + +/// Returns the size in bytes of the Parquet footer (the serialized file +/// metadata plus the 8 byte footer tail) of `file_bytes` +fn footer_size(file_bytes: &Bytes) -> usize { + let tail: &[u8; FOOTER_SIZE] = file_bytes[file_bytes.len() - FOOTER_SIZE..] + .try_into() + .unwrap(); + FooterTail::try_new(tail).unwrap().metadata_length() + FOOTER_SIZE +} diff --git a/datafusion/core/tests/parquet/io/page_pruning_io.rs b/datafusion/core/tests/parquet/io/page_pruning_io.rs new file mode 100644 index 0000000000000..cda10224325f0 --- /dev/null +++ b/datafusion/core/tests/parquet/io/page_pruning_io.rs @@ -0,0 +1,156 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! Tests tracking I/O events during Parquet page pruning using the page index +//! (`datafusion.execution.parquet.enable_page_index`, see +//! [`ParquetOptions::enable_page_index`]). +//! +//! The file has a single row group, so row group statistics can not prune +//! anything. Each column chunk is split into three data pages of three rows, +//! and `id` is sorted, so the page index (per page min/max of `id`) can +//! prune pages within the row group: +//! +//! ```text +//! row group 0 +//! column 'id' column 'val' +//! ┌──────────────┐ ┌──────────────┐ +//! page 0 │ 1, 2, 3 │ │ 0, 5, 10 │ rows 0..3 +//! ├──────────────┤ ├──────────────┤ +//! page 1 │ 4, 5, 6 │ │ 0, 7, 10 │ rows 3..6 +//! ├──────────────┤ ├──────────────┤ +//! page 2 │ 7, 8, 9 │ │ 0, 5, 10 │ rows 6..9 +//! └──────────────┘ └──────────────┘ +//! id page min/max: page 0: 1/3 page 1: 4/6 page 2: 7/9 +//! +//! id = 5 -> rows 3..6 -> page 1 of each column +//! id > 3 -> rows 3..9 -> pages 1 and 2 of each column +//! ``` +//! +//! [`ParquetOptions::enable_page_index`]: datafusion::config::ParquetOptions::enable_page_index + +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int32Array, RecordBatch}; +use bytes::Bytes; +use datafusion::prelude::SessionConfig; +use insta::assert_snapshot; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::WriterProperties; + +use super::ParquetIoTest; + +/// No predicate, so every page of every column is read as one column chunk. +#[tokio::test] +async fn test_page_pruning_io_no_predicate() { + let test = page_pruning_test(true).await; + // 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) + "); +} + +/// Only page 1 of `id` (`4..=6`) can contain `id = 5`, so only page 1 of +/// each column is fetched. +#[tokio::test] +async fn test_page_pruning_io_one_page() { + let test = page_pruning_test(true).await; + assert_snapshot!(test.query("SELECT * FROM t WHERE id = 5").await, @r" + --- Query: SELECT * FROM t WHERE id = 5 (1 rows) --- + 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) + "); +} + +/// Page 0 of `id` (`1..=3`) cannot contain `id > 3`, so pages 1 and 2 of +/// each column are fetched. +#[tokio::test] +async fn test_page_pruning_io_two_pages() { + let test = page_pruning_test(true).await; + assert_snapshot!(test.query("SELECT * FROM t WHERE id > 3").await, @r" + --- Query: SELECT * FROM t WHERE id > 3 (6 rows) --- + 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) + "); +} + +/// With page index pruning disabled, the whole column chunk of every column +/// is read even though only page 1 can match. Note the page index itself is +/// still read, it is just not used to prune. +#[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) + "); +} + +/// Returns a [`ParquetIoTest`] for the file described in the module docs +async fn page_pruning_test(page_index_pruning: bool) -> ParquetIoTest { + let config = SessionConfig::new() + .with_target_partitions(1) + .with_parquet_page_index_pruning(page_index_pruning); + ParquetIoTest::new(three_page_file(), config).await +} + +/// Returns a Parquet file with one row group of nine rows, where each column +/// chunk is written as three data pages of three rows each (see the module +/// docs). Dictionary encoding is disabled so each column chunk contains only +/// data pages. +fn three_page_file() -> Bytes { + let id: ArrayRef = Arc::new(Int32Array::from_iter_values(1..=9)); + let val: ArrayRef = Arc::new(Int32Array::from(vec![0, 5, 10, 0, 7, 10, 0, 5, 10])); + let batch = RecordBatch::try_from_iter([("id", id), ("val", val)]).unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_data_page_row_count_limit(3) + .set_write_batch_size(3) + .build(); + + let mut buffer = Vec::new(); + let mut writer = + ArrowWriter::try_new(&mut buffer, batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + Bytes::from(buffer) +} diff --git a/datafusion/core/tests/parquet/io/row_group_pruning_io.rs b/datafusion/core/tests/parquet/io/row_group_pruning_io.rs new file mode 100644 index 0000000000000..b6ac9cd6c5338 --- /dev/null +++ b/datafusion/core/tests/parquet/io/row_group_pruning_io.rs @@ -0,0 +1,150 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! Tests tracking I/O events during Parquet row group pruning. +//! +//! All tests filter on `id`, whose values are sorted across row groups so the +//! row group statistics are disjoint and can prune: +//! +//! ```text +//! row group 0 row group 1 row group 2 +//! ┌──────┬───────┐ ┌──────┬───────┐ ┌──────┬───────┐ +//! │ id │ val │ │ id │ val │ │ id │ val │ +//! ├──────┼───────┤ ├──────┼───────┤ ├──────┼───────┤ +//! │ 1 │ 0 │ │ 4 │ 0 │ │ 7 │ 0 │ +//! │ 2 │ 5 │ │ 5 │ 7 │ │ 8 │ 5 │ +//! │ 3 │ 10 │ │ 6 │ 10 │ │ 9 │ 10 │ +//! └──────┴───────┘ └──────┴───────┘ └──────┴───────┘ +//! id: min 1 max 3 id: min 4 max 6 id: min 7 max 9 +//! +//! id = 5 -> row group 1 only +//! id > 5 -> row groups 1 and 2 +//! id > 100 -> no row groups +//! ``` + +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int32Array, RecordBatch}; +use bytes::Bytes; +use datafusion::prelude::SessionConfig; +use insta::assert_snapshot; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::WriterProperties; + +use super::ParquetIoTest; + +/// No predicate, so every row group is read. +#[tokio::test] +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) + "); +} + +/// Only row group 1 (`id` in `4..=6`) can contain `id = 5`, so the data pages +/// for row groups 0 and 2 are never fetched. +#[tokio::test] +async fn test_row_group_pruning_io_one_row_group() { + let test = row_group_pruning_test().await; + assert_snapshot!(test.query("SELECT * FROM t WHERE id = 5").await, @r" + --- Query: SELECT * FROM t WHERE id = 5 (1 rows) --- + 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) + "); +} + +/// Row group 0 (`id` in `1..=3`) cannot contain `id > 5`, so only row groups 1 +/// and 2 are fetched. +#[tokio::test] +async fn test_row_group_pruning_io_two_row_groups() { + let test = row_group_pruning_test().await; + assert_snapshot!(test.query("SELECT * FROM t WHERE id > 5").await, @r" + --- Query: SELECT * FROM t WHERE id > 5 (4 rows) --- + 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) + "); +} + +/// No row group can contain `id > 100`, so only the footer and page index +/// are read. +#[tokio::test] +async fn test_row_group_pruning_io_all_pruned() { + let test = row_group_pruning_test().await; + assert_snapshot!(test.query("SELECT * FROM t WHERE id > 100").await, @r" + --- Query: SELECT * FROM t WHERE id > 100 (0 rows) --- + 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) + "); +} + +/// Returns a [`ParquetIoTest`] for the file described in the module docs +async fn row_group_pruning_test() -> ParquetIoTest { + let config = SessionConfig::new() + .with_target_partitions(1) + .with_parquet_pruning(true); + ParquetIoTest::new(three_row_group_file(), config).await +} + +/// Returns the Parquet file described in the module docs: three row groups of +/// three rows each, with `id` sorted across row groups. +fn three_row_group_file() -> Bytes { + let id: ArrayRef = Arc::new(Int32Array::from_iter_values(1..=9)); + let val: ArrayRef = Arc::new(Int32Array::from(vec![0, 5, 10, 0, 7, 10, 0, 5, 10])); + let batch = RecordBatch::try_from_iter([("id", id), ("val", val)]).unwrap(); + + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .build(); + + let mut buffer = Vec::new(); + let mut writer = + ArrowWriter::try_new(&mut buffer, batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + Bytes::from(buffer) +} diff --git a/datafusion/core/tests/parquet/io/topk_io.rs b/datafusion/core/tests/parquet/io/topk_io.rs new file mode 100644 index 0000000000000..3636e71f61a90 --- /dev/null +++ b/datafusion/core/tests/parquet/io/topk_io.rs @@ -0,0 +1,117 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! Tests tracking I/O events for TopK (`ORDER BY .. LIMIT`) queries. +//! +//! Both tests run `ORDER BY id DESC LIMIT 2`. `id` is sorted across row +//! groups, so sort pushdown reads the row groups in reverse order (largest +//! `id` first) and the two matching rows are both in row group 2: +//! +//! ```text +//! row group 0 row group 1 row group 2 +//! ┌──────┬───────┐ ┌──────┬───────┐ ┌──────┬───────┐ +//! │ id │ val │ │ id │ val │ │ id │ val │ +//! ├──────┼───────┤ ├──────┼───────┤ ├──────┼───────┤ +//! │ 1 │ 0 │ │ 4 │ 0 │ │ 7 │ 0 │ +//! │ 2 │ 5 │ │ 5 │ 7 │ │ 8 │ 5 * │ +//! │ 3 │ 10 │ │ 6 │ 10 │ │ 9 │ 10 * │ +//! └──────┴───────┘ └──────┴───────┘ └──────┴───────┘ +//! id: min 1 max 3 id: min 4 max 6 id: min 7 max 9 +//! +//! read order: 3rd read order: 2nd read order: 1st +//! +//! * the two rows returned by `ORDER BY id DESC LIMIT 2` +//! ``` + +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int32Array, RecordBatch}; +use bytes::Bytes; +use datafusion::prelude::SessionConfig; +use insta::assert_snapshot; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::WriterProperties; + +use super::ParquetIoTest; + +/// TopK dynamic filter pushdown prunes row groups during the scan. +/// +/// Sort pushdown reorders the row groups using their statistics so row group 2 +/// (largest `id` values) is read first. After it is read, the TopK filter is +/// `id > 8` which proves that row groups 1 and 0 cannot contain better rows, +/// so they are pruned without being fetched. +#[tokio::test] +async fn test_topk_io_dynamic_filter() { + let test = topk_test(true).await; + // 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) + "); +} + +/// Without TopK dynamic filter pushdown, the row groups are still read in the +/// optimal order, but all of them are fetched. +#[tokio::test] +async fn test_topk_io_no_dynamic_filter() { + let test = topk_test(false).await; + assert_snapshot!(test.query("SELECT * FROM t ORDER BY id DESC LIMIT 2").await, @r" + --- Query: SELECT * FROM t ORDER BY id DESC LIMIT 2 (2 rows) --- + 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) + "); +} + +/// Returns a [`ParquetIoTest`] for the file described in the module docs +async fn topk_test(dynamic_filter_pushdown: bool) -> ParquetIoTest { + let mut config = SessionConfig::new().with_target_partitions(1); + config + .options_mut() + .optimizer + .enable_topk_dynamic_filter_pushdown = dynamic_filter_pushdown; + ParquetIoTest::new(three_row_group_file(), config).await +} + +/// Returns the Parquet file described in the module docs: three row groups of +/// three rows each, with `id` sorted across row groups. +fn three_row_group_file() -> Bytes { + let id: ArrayRef = Arc::new(Int32Array::from_iter_values(1..=9)); + let val: ArrayRef = Arc::new(Int32Array::from(vec![0, 5, 10, 0, 7, 10, 0, 5, 10])); + let batch = RecordBatch::try_from_iter([("id", id), ("val", val)]).unwrap(); + + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .build(); + + let mut buffer = Vec::new(); + let mut writer = + ArrowWriter::try_new(&mut buffer, batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + Bytes::from(buffer) +}