This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11168-40766f3f7db68f706c532b9b132b3563e26615f5 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit 8c9d69472000e9237683600a5eaf9fc46163ef34 Author: Qi Zhu <[email protected]> AuthorDate: Thu Oct 15 12:57:42 2026 +0000 Defer dictionary page decode on the skip path (#11168) # Which issue does this PR close? - Closes #11154. # Rationale for this change `GenericColumnReader::skip_records` decompresses a column chunk's dictionary page as soon as it reaches it, even when every remaining page is then skipped and no value from the chunk is ever decoded. Skipping values needs only the RLE index cursor and skipping whole pages needs nothing, so dictionary contents are only ever needed by value decoding. On zstd-compressed data with large dictionaries this is a large, avoidable cost, and it is most visible with `pushdown_filters`/`RowSelection` and with predicate caching, where output columns are consumed largely by skipping. A local timing test on a zstd column whose dictionary holds 40k distinct strings (14,800 compressed bytes): | | | | --- | --- | | whole-chunk skip, dictionary never decoded | **84 µs** | | partial skip that pays the deferred dictionary + first page | 4.76 ms | No regression on the plain read path. The read path gains one `matches!` on the page encoding and one always-`None` `take()` per data page; `arrow_reader`'s dictionary-encoded read benchmarks are all within the noise threshold against `main`: | dictionary-encoded read | change vs main | | --- | --- | | Int32, mandatory | +0.43% | | Int32, optional | +0.58% | | Int32, optional, half NULLs | +0.72% | | UInt8, mandatory | +0.28% | | UInt8, optional | +0.35% | (Run locally with `--features arrow,test_common,experimental`; CodSpeed does not run on fork PRs.) # What changes are included in this PR? The deferral is created on the skip path and consumed on the read path; that crossing is the whole change: ```mermaid flowchart TD A["skip_records reaches a dictionary page"] --> G{"supports_deferred_dictionary()?"} G -->|"false (trait default)"| H["read_dictionary_page()<br/>eager, unchanged from main"] G -->|"true"| B["record data_start + header<br/>no decompress"] B -.->|"deferred, may never fire"| C["read_records: a data page needs values"] C --> D["take_deferred_dictionary()"] D --> E["get_bytes(data_start)<br/>decode_page, set_dict"] B ==>|"chunk ends first"| F["never decompressed"] ``` - The skip path records the dictionary page's location (`DeferredDictionaryPage`) instead of decoding it. A data page that turns out to need the dictionary installs it then, paying the decompression exactly once; a chunk skipped end to end never pays it. - `PageReader::supports_deferred_dictionary`, a new trait method defaulting to `false`, decides whether the skip path may skip past a dictionary at all. A reader that does not opt in keeps the eager path exactly as before — skipping past a dictionary it cannot hand back would simply lose it. `SerializedPageReader` answers per state: it can always defer while walking page headers, and declines in the offset-index state when the metadata recorded no `dictionary_page_offset`, because there a dictionary inlined ahead of the first data page is indistinguishable from a data page location. - `PageReader::take_deferred_dictionary`, a new trait method defaulting to `Ok(None)`, returns a dictionary the reader skipped past. - Whether a page is a dictionary is judged by its actual header type, not the `dictionary_page_offset` metadata flag, so files that inline the dictionary without recording an offset (older parquet-mr writers, e.g. the one behind `test_read_nested_lists`) are handled correctly. **On the re-read.** The deferred install is a `ChunkReader::get_read` at the recorded offset, which is already part of the trait. It never adds I/O relative to `main`: these are the same bytes the eager path would have read, either read later or, when the chunk is skipped end to end, not at all. On the async path it is not a fetch at all — the column chunk is already buffered, so the read is an in-memory slice. # Are these changes tested? Yes. The deferral itself: - whole-chunk skip never decodes the dictionary (corrupting the dictionary body leaves the skip green); - a partial skip pays it (the same corruption then fails); - skip-then-read returns correct values; - alternating skip and read — the shape a `RowSelection` produces — triggers the deferral more than once per chunk; - the plain read path is unchanged; - a timing test prints the avoided cost. Page readers that cannot defer: - a reader forwarding only the pre-existing trait methods still gets its dictionary, through both a partial and a whole-chunk skip. Removing the `supports_deferred_dictionary` gate makes it panic in `set_data`; - a reader forwarding `take_deferred_dictionary` but not `supports_deferred_dictionary` — the easy mistake when adding both to an existing wrapper — stays correct by falling back to the eager path. The probe's answer in all three `SerializedPageReader` states, including the offset-index state without a `dictionary_page_offset`, where it declines. The existing `test_read_nested_lists` (nested column + `RowSelection`) caught the offset-less-dictionary case during development and still covers it. # Are there any user-facing changes? No format change, and no behavior change for existing `PageReader` implementations: deferral is opt-in through `supports_deferred_dictionary`, which defaults to `false`, so a reader that does not implement it keeps the eager dictionary read it had before. `PageReader` gains two defaulted trait methods; an implementation that wraps another `PageReader` should forward both to benefit, and stays correct if it forwards neither. Forwarding only `take_deferred_dictionary` without `supports_deferred_dictionary` is also safe — the skip path simply stays eager. One observable difference for a reader that does opt in: a **corrupt dictionary page is no longer guaranteed to raise an error**. Today the page is decompressed the moment the skip path reaches it, so a corrupt one fails even for a column chunk that goes on to decode nothing. With the deferral, a chunk skipped end to end never touches those bytes and the error never surfaces. Any chunk that decodes a value still hits it, just later. This is not a correctness change — no wrong data is returned either way — but it does mean corruption in a never-read dictionary now goes unreported. `skipping_a_whole_chunk_never_decodes_the_dictionary` asserts exactly this, using a deliberately corrupted dictionary body as the probe for "was it decompressed". --------- Co-authored-by: Ed Seidl <[email protected]> Co-authored-by: Andrew Lamb <[email protected]> --- parquet/Cargo.toml | 5 + parquet/benches/dictionary_skip.rs | 99 ++++++ parquet/src/column/page.rs | 24 ++ parquet/src/column/reader.rs | 48 ++- parquet/src/file/serialized_reader.rs | 491 ++++++++++++++++++++++++++- parquet/tests/arrow_reader/io/sync_reader.rs | 27 +- 6 files changed, 676 insertions(+), 18 deletions(-) diff --git a/parquet/Cargo.toml b/parquet/Cargo.toml index e34565c260..c36791f571 100644 --- a/parquet/Cargo.toml +++ b/parquet/Cargo.toml @@ -295,6 +295,11 @@ name = "compression" required-features = ["experimental", "default"] harness = false +[[bench]] +name = "dictionary_skip" +required-features = ["default"] +harness = false + [[bench]] name = "encoding" required-features = ["experimental", "default"] diff --git a/parquet/benches/dictionary_skip.rs b/parquet/benches/dictionary_skip.rs new file mode 100644 index 0000000000..103d908ecb --- /dev/null +++ b/parquet/benches/dictionary_skip.rs @@ -0,0 +1,99 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Skipping records in a dictionary-encoded column chunk. +//! +//! The chunk is shaped like a real serving file: a zstd-compressed BYTE_ARRAY +//! column whose dictionary holds 40k distinct strings. Skipping the whole +//! chunk never needs the dictionary; skipping part of it and then reading +//! does. + +use std::hint; +use std::sync::Arc; + +use bytes::Bytes; +use criterion::*; +use parquet::basic::Compression; +use parquet::column::reader::ColumnReaderImpl; +use parquet::data_type::{ByteArray, ByteArrayType}; +use parquet::file::properties::WriterProperties; +use parquet::file::reader::{FileReader, SerializedFileReader}; +use parquet::file::serialized_reader::SerializedPageReader; +use parquet::file::writer::SerializedFileWriter; +use parquet::schema::parser::parse_message_type; + +const NUM_ROWS: usize = 80_000; + +fn write_chunk() -> Bytes { + let schema = + Arc::new(parse_message_type("message test_schema { REQUIRED BYTE_ARRAY v; }").unwrap()); + let props = Arc::new( + WriterProperties::builder() + .set_compression(Compression::ZSTD(Default::default())) + .build(), + ); + let mut out = Vec::new(); + let mut writer = SerializedFileWriter::new(&mut out, schema, props).unwrap(); + let mut row_group = writer.next_row_group().unwrap(); + let mut column = row_group.next_column().unwrap().unwrap(); + let values: Vec<ByteArray> = (0..NUM_ROWS as u32) + .map(|i| ByteArray::from(format!("ticker-{:016}", i / 2).into_bytes())) + .collect(); + column + .typed::<ByteArrayType>() + .write_batch(&values, None, None) + .unwrap(); + column.close().unwrap(); + row_group.close().unwrap(); + writer.close().unwrap(); + Bytes::from(out) +} + +fn dictionary_skip_benchmark(c: &mut Criterion) { + let bytes = write_chunk(); + let reader = SerializedFileReader::new(bytes.clone()).unwrap(); + let metadata = reader.metadata(); + let descr = metadata.file_metadata().schema_descr().column(0); + let column_metadata = metadata.row_group(0).column(0).clone(); + + let new_reader = || { + let pages = + SerializedPageReader::new(Arc::new(bytes.clone()), &column_metadata, NUM_ROWS, None) + .unwrap(); + ColumnReaderImpl::<ByteArrayType>::new(descr.clone(), Box::new(pages)) + }; + + let mut group = c.benchmark_group("dictionary_skip"); + group.bench_function("skip whole chunk", |b| { + b.iter(|| { + let mut reader = new_reader(); + hint::black_box(reader.skip_records(NUM_ROWS).unwrap()) + }) + }); + group.bench_function("skip one record then read", |b| { + b.iter(|| { + let mut reader = new_reader(); + reader.skip_records(1).unwrap(); + let mut values = Vec::new(); + hint::black_box(reader.read_records(1, None, None, &mut values).unwrap()) + }) + }); + group.finish(); +} + +criterion_group!(benches, dictionary_skip_benchmark); +criterion_main!(benches); diff --git a/parquet/src/column/page.rs b/parquet/src/column/page.rs index b0e60b9f0d..05f4f478e9 100644 --- a/parquet/src/column/page.rs +++ b/parquet/src/column/page.rs @@ -404,6 +404,30 @@ pub trait PageReader: Iterator<Item = Result<Page>> + Send { /// column index information fn skip_next_page(&mut self) -> Result<()>; + /// Returns whether skipped dictionary pages can later be recovered with + /// [`Self::take_deferred_dictionary`]. + /// + /// Implementations returning `true` must retain dictionary pages consumed by + /// [`Self::skip_next_page`]. + fn supports_deferred_dictionary(&self) -> bool { + false + } + + /// Returns the decoded dictionary [`Page`] this reader has previously + /// skipped past, if any, removing it from the reader. + /// + /// [`Self::skip_next_page`] may skip a dictionary page without decoding + /// it, since skipping rows never needs dictionary contents. If decoding + /// later reaches a dictionary-encoded data page, the reader is asked for + /// the dictionary through this method, which pays the deferred + /// decompression exactly once. A chunk skipped end to end never pays it. + /// + /// The default implementation returns `Ok(None)`, meaning the reader + /// never defers. + fn take_deferred_dictionary(&mut self) -> Result<Option<Page>> { + Ok(None) + } + /// Returns `true` if the next page can be assumed to contain the start of a new record /// /// Prior to parquet V2 the specification was ambiguous as to whether a single record diff --git a/parquet/src/column/reader.rs b/parquet/src/column/reader.rs index e6fe0b6a27..c122c491af 100644 --- a/parquet/src/column/reader.rs +++ b/parquet/src/column/reader.rs @@ -330,9 +330,14 @@ where return Ok(num_records - remaining_records); }; - // If dictionary, we must read it + // Defer dictionary decoding when the page reader can recover it later. + // Otherwise, preserve the eager behavior required by existing readers. if metadata.is_dict { - self.read_dictionary_page()?; + if self.page_reader.supports_deferred_dictionary() { + self.page_reader.skip_next_page()?; + } else { + self.read_dictionary_page()?; + } continue; } @@ -449,8 +454,12 @@ where Ok(()) } - /// Read the next page as a dictionary page. If the next page is not a dictionary page, - /// this will return an error. + /// Reads the next page as a dictionary page and installs it. + /// + /// Used on the skip path for readers that do not retain skipped + /// dictionaries, and so must consume the page through `get_next_page` + /// while it is still available. Returns an error if the next page is not + /// a dictionary page. fn read_dictionary_page(&mut self) -> Result<()> { match self.page_reader.get_next_page()? { Some(Page::DictionaryPage { @@ -467,6 +476,25 @@ where } } + /// Installs a dictionary previously deferred by the page reader. + /// A no-op when nothing was deferred. + fn install_deferred_dictionary(&mut self) -> Result<()> { + match self.page_reader.take_deferred_dictionary()? { + Some(Page::DictionaryPage { + buf, + num_values, + encoding, + is_sorted, + }) => self + .values_decoder + .set_dict(buf, num_values, encoding, is_sorted), + Some(_) => Err(ParquetError::General( + "Invalid page. Expecting dictionary page".to_string(), + )), + None => Ok(()), + } + } + /// Reads a new page and set up the decoders for levels, values or dictionary. /// Returns false if there's no page left. fn read_new_page(&mut self) -> Result<bool> { @@ -495,6 +523,12 @@ where rep_level_encoding, statistics: _, } => { + if matches!( + encoding, + Encoding::PLAIN_DICTIONARY | Encoding::RLE_DICTIONARY + ) { + self.install_deferred_dictionary()?; + } self.num_buffered_values = num_values as _; self.num_decoded_values = 0; @@ -580,6 +614,12 @@ where is_compressed: _, statistics: _, } => { + if matches!( + encoding, + Encoding::PLAIN_DICTIONARY | Encoding::RLE_DICTIONARY + ) { + self.install_deferred_dictionary()?; + } if num_nulls > num_values { return Err(general_err!( "more nulls than values in page, contained {} values and {} nulls", diff --git a/parquet/src/file/serialized_reader.rs b/parquet/src/file/serialized_reader.rs index aab7bc8ab3..1f6924c4b0 100644 --- a/parquet/src/file/serialized_reader.rs +++ b/parquet/src/file/serialized_reader.rs @@ -531,6 +531,39 @@ pub(crate) fn decode_page( Ok(result) } +/// Location of a dictionary page skipped without decoding. +/// +/// ```text +/// skip, and the reader can defer first data page that +/// (record location, zero decompress) needs the dictionary +/// NotSeen ─────────────────────────────▶ Deferred ─────────────────────▶ Decoded +/// │ ▲ +/// ├── read hits the dictionary page (eager path, unchanged) ─────────────┤ +/// │ │ +/// └── skip, but the reader cannot defer ─────────────────────────────────┘ +/// (`supports_deferred_dictionary()` is false: read it eagerly, +/// because a skipped page this reader cannot hand back is lost) +/// ``` +/// +/// A chunk skipped end to end stays `Deferred`: its dictionary is never +/// decompressed, which is the point. Skipping *values* needs only the RLE +/// index cursor, and skipping whole pages needs nothing at all, so the only +/// consumer of dictionary *contents* is value decoding. Random re-read is +/// what makes the deferral safe: [`ChunkReader`] hands out readers at any +/// offset, and the async path's in-memory column chunk already holds the +/// page's bytes. +enum DeferredDictionaryPage { + /// `Values` state: the header was already parsed on the skip path, only + /// the page body was left untouched. + Header { + data_start: u64, + header: Box<PageHeader>, + }, + /// `Pages` state: the offset index records the page location; header and + /// body are both still unread. + Location(PageLocation), +} + enum SerializedPageReaderState { Values { /// The current byte offset in the reader @@ -576,6 +609,10 @@ pub struct SerializedPageReader<R: ChunkReader> { /// The chunk reader reader: Arc<R>, + /// A dictionary page skipped past without decoding. See + /// [`DeferredDictionaryPage`] for the state diagram. + deferred_dictionary: Option<DeferredDictionaryPage>, + /// The compression codec for this column chunk. Only set for non-PLAIN codec. decompressor: Option<Box<dyn Codec>>, @@ -683,6 +720,7 @@ impl<R: ChunkReader> SerializedPageReader<R> { } Ok(Self { reader, + deferred_dictionary: None, decompressor, state, physical_type: meta.column_type(), @@ -1114,7 +1152,9 @@ impl<R: ChunkReader> PageReader for SerializedPageReader<R> { } fn skip_next_page(&mut self) -> Result<()> { - match &mut self.state { + // Retain skipped dictionary pages so they can be decoded on demand. + // See [`DeferredDictionaryPage`] for the state diagram. + let deferred = match &mut self.state { SerializedPageReaderState::Values { offset, remaining_bytes, @@ -1122,15 +1162,24 @@ impl<R: ChunkReader> PageReader for SerializedPageReader<R> { page_index, require_dictionary, } => { + let mut deferred = None; if let Some(buffered_header) = next_page_header.take() { verify_page_size( buffered_header.compressed_page_size, buffered_header.uncompressed_page_size, *remaining_bytes, )?; - // The next page header has already been peeked, so just advance the offset - *offset += buffered_header.compressed_page_size as u64; - *remaining_bytes -= buffered_header.compressed_page_size as u64; + // The next page header has already been peeked, so the + // offset already points at the page body. + let body_len = buffered_header.compressed_page_size as u64; + if buffered_header.r#type == PageType::DICTIONARY_PAGE { + deferred = Some(DeferredDictionaryPage::Header { + data_start: *offset, + header: buffered_header, + }); + } + *offset += body_len; + *remaining_bytes -= body_len; } else { let mut read = self.reader.get_read(*offset)?; let (header_len, header) = Self::read_page_header_len( @@ -1146,14 +1195,22 @@ impl<R: ChunkReader> PageReader for SerializedPageReader<R> { *remaining_bytes, )?; let data_page_size = header.compressed_page_size as u64; + let is_dict = header.r#type == PageType::DICTIONARY_PAGE; + if is_dict { + deferred = Some(DeferredDictionaryPage::Header { + data_start: *offset + header_len as u64, + header: Box::new(header), + }); + } *offset += header_len as u64 + data_page_size; *remaining_bytes -= header_len as u64 + data_page_size; } - if *require_dictionary { + if deferred.is_some() { *require_dictionary = false; } else { *page_index += 1; } + deferred } SerializedPageReaderState::Pages { page_locations, @@ -1162,19 +1219,70 @@ impl<R: ChunkReader> PageReader for SerializedPageReader<R> { .. } => { if dictionary_page.is_some() { - // If a dictionary page exists, consume it by taking it (sets to None) - dictionary_page.take(); + // Consume the dictionary page, keeping its location. + dictionary_page.take().map(DeferredDictionaryPage::Location) } else { // If no dictionary page exists, simply pop the data page from page_locations if page_locations.pop_front().is_some() { *page_index += 1; } + None } } + }; + if deferred.is_some() { + self.deferred_dictionary = deferred; } Ok(()) } + fn supports_deferred_dictionary(&self) -> bool { + match &self.state { + // Walking headers identifies dictionaries independently of metadata. + SerializedPageReaderState::Values { .. } => true, + // With an offset index, a dictionary is recoverable only when its + // location can be inferred before the first data-page location, in + // which case `dictionary_page` will be `Some`. + SerializedPageReaderState::Pages { + dictionary_page, .. + } => dictionary_page.is_some(), + } + } + + fn take_deferred_dictionary(&mut self) -> Result<Option<Page>> { + let Some(deferred) = self.deferred_dictionary.take() else { + return Ok(None); + }; + let page = match deferred { + DeferredDictionaryPage::Header { data_start, header } => { + let data_len = header.compressed_page_size as usize; + let buffer = self.reader.get_bytes(data_start, data_len)?; + let buffer = self.context.decrypt_page_data(buffer, 0, true)?; + decode_page( + *header, + buffer, + self.physical_type, + self.decompressor.as_mut(), + )? + } + DeferredDictionaryPage::Location(front) => { + let page_len = usize::try_from(front.compressed_page_size)?; + let buffer = self.reader.get_bytes(front.offset as u64, page_len)?; + let (offset, header) = + read_page_header_len_from_bytes(&self.context, buffer.as_ref(), 0, true)?; + let bytes = buffer.slice(offset..); + let bytes = self.context.decrypt_page_data(bytes, 0, true)?; + decode_page( + header, + bytes, + self.physical_type, + self.decompressor.as_mut(), + )? + } + }; + Ok(Some(page)) + } + fn at_record_boundary(&mut self) -> Result<bool> { match &mut self.state { SerializedPageReaderState::Values { .. } => match self.peek_next_page()? { @@ -1199,12 +1307,13 @@ mod tests { }; use crate::file::properties::{EnabledStatistics, WriterProperties}; - use crate::basic::{self, BoundaryOrder, ColumnOrder, Encoding, SortOrder}; + use crate::basic::{self, BoundaryOrder, ColumnOrder, Compression, Encoding, SortOrder}; use crate::column::reader::ColumnReader; use crate::data_type::private::ParquetValueType; use crate::data_type::{AsBytes, FixedLenByteArrayType, Int32Type}; use crate::file::metadata::thrift::DataPageHeaderV2; use crate::file::writer::SerializedFileWriter; + use crate::parquet_thrift::{ReadThrift, ThriftSliceInputProtocol}; use crate::record::RowAccessor; use crate::schema::parser::parse_message_type; use crate::util::test_common::file_util::{get_test_file, get_test_path}; @@ -2911,4 +3020,370 @@ mod tests { } assert_eq!(num_rows, reader.metadata().file_metadata().num_rows()); } + + /// Builds an in-memory single-row-group file with one snappy-compressed, + /// dictionary-encoded INT32 column holding `1..=5`, and returns its bytes + /// together with the byte range of the dictionary page *body* (the + /// compressed payload after the page header), so tests can corrupt + /// exactly the part lazy skipping promises never to touch. + fn dict_column_file() -> (Bytes, std::ops::Range<usize>) { + let message_type = " + message test_schema { + REQUIRED INT32 a; + } + "; + let schema = Arc::new(parse_message_type(message_type).unwrap()); + let props = Arc::new( + WriterProperties::builder() + .set_compression(Compression::SNAPPY) + .build(), + ); + let mut out = Vec::new(); + let mut writer = SerializedFileWriter::new(&mut out, schema, props).unwrap(); + let mut rg = writer.next_row_group().unwrap(); + let mut col = rg.next_column().unwrap().unwrap(); + col.typed::<Int32Type>() + .write_batch(&[1, 2, 3, 4, 5], None, None) + .unwrap(); + col.close().unwrap(); + rg.close().unwrap(); + writer.close().unwrap(); + let bytes = Bytes::from(out); + + let reader = SerializedFileReader::new(bytes.clone()).unwrap(); + let col_meta = reader.metadata().row_group(0).column(0).clone(); + let dict_offset = col_meta.dictionary_page_offset().unwrap() as usize; + let data_offset = col_meta.data_page_offset() as usize; + // Parse the dictionary page header to find where its body starts. + let mut prot = ThriftSliceInputProtocol::new(&bytes[dict_offset..data_offset]); + let _header = PageHeader::read_thrift(&mut prot).unwrap(); + let header_len = (data_offset - dict_offset) - prot.as_slice().len(); + (bytes, dict_offset + header_len..data_offset) + } + + /// Zeroes the dictionary page body: an all-zero snappy stream declares an + /// uncompressed length of zero, so decoding it fails deterministically on + /// the decompressed-size check. Any code path that decompresses the + /// dictionary therefore errors; a path that genuinely never touches it + /// succeeds. + fn corrupt_dictionary_body(bytes: &Bytes, body: &std::ops::Range<usize>) -> Bytes { + let mut v = bytes.to_vec(); + v[body.clone()].fill(0); + Bytes::from(v) + } + + fn int32_reader(bytes: Bytes) -> crate::column::reader::ColumnReaderImpl<Int32Type> { + let reader = SerializedFileReader::new(bytes.clone()).unwrap(); + let meta = reader.metadata(); + let descr = meta.file_metadata().schema_descr().column(0); + let col_meta = meta.row_group(0).column(0).clone(); + let total_rows = meta.row_group(0).num_rows() as usize; + let pages = + SerializedPageReader::new(Arc::new(bytes), &col_meta, total_rows, None).unwrap(); + crate::column::reader::ColumnReaderImpl::<Int32Type>::new(descr, Box::new(pages)) + } + + /// Skipping the whole chunk never decompresses the dictionary: with the + /// dictionary body corrupted so that any decompression errors, the skip + /// still succeeds. Reverting the lazy branch in + /// `GenericColumnReader::skip_records` turns this test red. + #[test] + fn skipping_a_whole_chunk_never_decodes_the_dictionary() { + let (bytes, body) = dict_column_file(); + let corrupted = corrupt_dictionary_body(&bytes, &body); + let mut reader = int32_reader(corrupted); + let skipped = reader.skip_records(5).unwrap(); + assert_eq!(skipped, 5); + } + + /// The deferred dictionary is genuinely decoded the moment a data page + /// needs it: the same corruption that whole-chunk skipping never notices + /// fails a skip that has to open the data page. + #[test] + fn a_partial_skip_pays_the_deferred_dictionary() { + let (bytes, body) = dict_column_file(); + let corrupted = corrupt_dictionary_body(&bytes, &body); + let mut reader = int32_reader(corrupted); + let err = reader.skip_records(2).unwrap_err().to_string(); + assert!( + err.contains("corrupt") || err.contains("size"), + "expected the deferred dictionary decode to fail on \ + decompression, got: {err}" + ); + } + + /// Skip-then-read returns the right values: the dictionary deferred by + /// the skip is installed before the first value decodes. + #[test] + fn skip_then_read_decodes_correct_values() { + let (bytes, _) = dict_column_file(); + let mut reader = int32_reader(bytes); + assert_eq!(reader.skip_records(2).unwrap(), 2); + let mut values = Vec::new(); + let (records, _, read) = reader.read_records(3, None, None, &mut values).unwrap(); + assert_eq!(records, 3); + assert_eq!(read, 3); + assert_eq!(values, vec![3, 4, 5]); + } + + /// The eager path is untouched: reading from the start still installs the + /// dictionary through `get_next_page` and decodes every value. + #[test] + fn plain_read_is_unchanged() { + let (bytes, _) = dict_column_file(); + let mut reader = int32_reader(bytes); + let mut values = Vec::new(); + let (records, _, read) = reader.read_records(5, None, None, &mut values).unwrap(); + assert_eq!((records, read), (5, 5)); + assert_eq!(values, vec![1, 2, 3, 4, 5]); + } + + /// A [`PageReader`] that implements only what existed before deferral: it + /// forwards the required trait methods and takes the defaults for + /// `supports_deferred_dictionary` (`false`) and `take_deferred_dictionary` + /// (`Ok(None)`). This stands in for every third-party implementation, none + /// of which can retain a dictionary that `skip_next_page` skips past. + struct NonDeferringPageReader(SerializedPageReader<Bytes>); + + impl Iterator for NonDeferringPageReader { + type Item = Result<Page>; + + fn next(&mut self) -> Option<Self::Item> { + self.get_next_page().transpose() + } + } + + impl PageReader for NonDeferringPageReader { + fn get_next_page(&mut self) -> Result<Option<Page>> { + self.0.get_next_page() + } + + fn peek_next_page(&mut self) -> Result<Option<PageMetadata>> { + self.0.peek_next_page() + } + + fn skip_next_page(&mut self) -> Result<()> { + self.0.skip_next_page() + } + } + + fn non_deferring_int32_reader( + bytes: Bytes, + ) -> crate::column::reader::ColumnReaderImpl<Int32Type> { + let reader = SerializedFileReader::new(bytes.clone()).unwrap(); + let meta = reader.metadata(); + let descr = meta.file_metadata().schema_descr().column(0); + let col_meta = meta.row_group(0).column(0).clone(); + let total_rows = meta.row_group(0).num_rows() as usize; + let pages = + SerializedPageReader::new(Arc::new(bytes), &col_meta, total_rows, None).unwrap(); + crate::column::reader::ColumnReaderImpl::<Int32Type>::new( + descr, + Box::new(NonDeferringPageReader(pages)), + ) + } + + /// A page reader that cannot give a skipped dictionary back must still get + /// the eager path. Skipping past its dictionary would drop it, and the + /// first dictionary-encoded data page decoded afterwards would have no + /// dictionary to decode against — a panic in `set_data`, not an error. + /// + /// Removing the `supports_deferred_dictionary` gate in + /// `GenericColumnReader::skip_records` turns this test red. + #[test] + fn a_reader_that_cannot_defer_still_gets_its_dictionary() { + let (bytes, _) = dict_column_file(); + let mut reader = non_deferring_int32_reader(bytes); + + assert_eq!(reader.skip_records(2).unwrap(), 2); + + let mut values = Vec::new(); + let (records, _, read) = reader.read_records(3, None, None, &mut values).unwrap(); + assert_eq!(records, 3); + assert_eq!(read, 3); + assert_eq!(values, vec![3, 4, 5]); + } + + /// The same reader must also survive a whole-chunk skip. Nothing is + /// decoded, so nothing needs the dictionary, but the eager read still has + /// to leave the page cursor in a consistent place. + #[test] + fn a_reader_that_cannot_defer_survives_a_whole_chunk_skip() { + let (bytes, _) = dict_column_file(); + let mut reader = non_deferring_int32_reader(bytes); + + assert_eq!(reader.skip_records(5).unwrap(), 5); + + let mut values = Vec::new(); + let (records, _, _) = reader.read_records(1, None, None, &mut values).unwrap(); + assert_eq!(records, 0, "chunk is exhausted"); + assert!(values.is_empty()); + } + + /// The probe has to describe what the reader can actually do. Without page + /// locations the reader walks page headers, so it recognises a dictionary + /// by its header type and can always defer. + #[test] + fn values_state_supports_deferral() { + let (bytes, _) = dict_column_file(); + let reader = SerializedFileReader::new(bytes.clone()).unwrap(); + let meta = reader.metadata(); + let col_meta = meta.row_group(0).column(0).clone(); + let total_rows = meta.row_group(0).num_rows() as usize; + let pages = + SerializedPageReader::new(Arc::new(bytes), &col_meta, total_rows, None).unwrap(); + + assert!(pages.supports_deferred_dictionary()); + } + + /// With page locations, a dictionary is only representable when the + /// metadata recorded an offset for it. `dict_column_file` writes one, so + /// the synthesised `dictionary_page` is present and deferral is safe. + #[test] + fn offset_index_state_with_a_dictionary_offset_supports_deferral() { + let (bytes, _) = dict_column_file(); + let reader = SerializedFileReader::new(bytes.clone()).unwrap(); + let meta = reader.metadata(); + let col_meta = meta.row_group(0).column(0).clone(); + let total_rows = meta.row_group(0).num_rows() as usize; + let (dict_offset, data_offset) = ( + col_meta.dictionary_page_offset().unwrap(), + col_meta.data_page_offset(), + ); + assert_ne!( + dict_offset, data_offset, + "fixture must record a dictionary offset" + ); + + let locations = vec![PageLocation { + offset: data_offset, + compressed_page_size: 0, + first_row_index: 0, + }]; + let pages = + SerializedPageReader::new(Arc::new(bytes), &col_meta, total_rows, Some(locations)) + .unwrap(); + + assert!(pages.supports_deferred_dictionary()); + } + + /// When the metadata records no dictionary offset, `byte_range()` starts + /// at the first data page, so the synthesised `dictionary_page` is `None` + /// and a dictionary inlined ahead of that page is indistinguishable from a + /// data page location. Skipping it would drop it, so the probe must say + /// so and let the column reader read it eagerly instead. + #[test] + fn offset_index_state_without_a_dictionary_offset_declines_deferral() { + let (bytes, _) = dict_column_file(); + let reader = SerializedFileReader::new(bytes.clone()).unwrap(); + let meta = reader.metadata(); + let total_rows = meta.row_group(0).num_rows() as usize; + + // Drop the dictionary offset, reproducing the older parquet-mr writers + // that inline a dictionary without recording where it is. + let col_meta = meta.row_group(0).column(0).clone(); + let stripped = col_meta + .clone() + .into_builder() + .set_dictionary_page_offset(None) + .build() + .unwrap(); + + let locations = vec![PageLocation { + offset: stripped.data_page_offset(), + compressed_page_size: 0, + first_row_index: 0, + }]; + let pages = + SerializedPageReader::new(Arc::new(bytes), &stripped, total_rows, Some(locations)) + .unwrap(); + + assert!( + !pages.supports_deferred_dictionary(), + "an offset-less dictionary cannot be recovered from the offset-index state" + ); + // Note this combination is unlikely to occur in practice: a writer old + // enough to omit `dictionary_page_offset` predates the offset index, + // so it would not emit page locations either. The probe answers + // defensively rather than relying on that. + } + + /// Alternating skip and read is the shape a `RowSelection` produces, and + /// it exercises the deferral more than once per chunk: the first skip + /// defers, the first read installs, and every later skip has to cope with + /// a dictionary that is already installed. + #[test] + fn alternating_skip_and_read_decodes_correct_values() { + let (bytes, _) = dict_column_file(); + let mut reader = int32_reader(bytes); + + assert_eq!(reader.skip_records(1).unwrap(), 1); + let mut first = Vec::new(); + let (records, _, _) = reader.read_records(2, None, None, &mut first).unwrap(); + assert_eq!(records, 2); + assert_eq!(first, vec![2, 3]); + + assert_eq!(reader.skip_records(1).unwrap(), 1); + let mut second = Vec::new(); + let (records, _, _) = reader.read_records(1, None, None, &mut second).unwrap(); + assert_eq!(records, 1); + assert_eq!(second, vec![5]); + } + + /// A wrapper that forwards `take_deferred_dictionary` but not + /// `supports_deferred_dictionary`, which is the easy mistake to make when + /// adding the new methods to an existing wrapper. + struct HalfForwardingPageReader(SerializedPageReader<Bytes>); + + impl Iterator for HalfForwardingPageReader { + type Item = Result<Page>; + + fn next(&mut self) -> Option<Self::Item> { + self.get_next_page().transpose() + } + } + + impl PageReader for HalfForwardingPageReader { + fn get_next_page(&mut self) -> Result<Option<Page>> { + self.0.get_next_page() + } + + fn peek_next_page(&mut self) -> Result<Option<PageMetadata>> { + self.0.peek_next_page() + } + + fn skip_next_page(&mut self) -> Result<()> { + self.0.skip_next_page() + } + + fn take_deferred_dictionary(&mut self) -> Result<Option<Page>> { + self.0.take_deferred_dictionary() + } + } + + /// Forwarding only half the pair must stay correct: the probe still + /// reports `false`, so the skip path reads the dictionary eagerly and the + /// forwarded `take_deferred_dictionary` simply never has anything to + /// return. Slower than forwarding both, never wrong. + #[test] + fn forwarding_only_take_deferred_dictionary_is_still_correct() { + let (bytes, _) = dict_column_file(); + let reader = SerializedFileReader::new(bytes.clone()).unwrap(); + let meta = reader.metadata(); + let descr = meta.file_metadata().schema_descr().column(0); + let col_meta = meta.row_group(0).column(0).clone(); + let total_rows = meta.row_group(0).num_rows() as usize; + let pages = + SerializedPageReader::new(Arc::new(bytes), &col_meta, total_rows, None).unwrap(); + let mut reader = crate::column::reader::ColumnReaderImpl::<Int32Type>::new( + descr, + Box::new(HalfForwardingPageReader(pages)), + ); + + assert_eq!(reader.skip_records(2).unwrap(), 2); + let mut values = Vec::new(); + let (records, _, _) = reader.read_records(3, None, None, &mut values).unwrap(); + assert_eq!(records, 3); + assert_eq!(values, vec![3, 4, 5]); + } } diff --git a/parquet/tests/arrow_reader/io/sync_reader.rs b/parquet/tests/arrow_reader/io/sync_reader.rs index 835029d5c6..1c1c1f8dd1 100644 --- a/parquet/tests/arrow_reader/io/sync_reader.rs +++ b/parquet/tests/arrow_reader/io/sync_reader.rs @@ -156,6 +156,11 @@ fn test_read_single_column_no_page_index() { #[test] fn test_read_row_selection() { + // Dictionary pages are read on first decode of a data page that needs + // them, so skipping a page does not pull its dictionary in. That is why + // the trace below shows the dictionary *after* the data page wherever a + // skip precedes it, and *before* the data page where none does. + // See apache/arrow-rs#11154. // There are 400 total rows spread across 4 data pages (100 rows each) // select rows 175..225 (i.e. DataPage(1) of row group 0 and DataPage(0) of row group 1) let test_file = test_file(); @@ -181,10 +186,10 @@ fn test_read_row_selection() { "UNKNOWN: 22230..22877 (maybe Page Index)", "Event: Builder Configured", "Event: Reader Built", - "Row Group 0, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", "Row Group 0, column 'a': DataPage(1) (126 bytes , 1 requests) [data]", - "Row Group 0, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 0, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", "Row Group 0, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 0, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", "Row Group 1, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", "Row Group 1, column 'a': DataPage(0) (113 bytes , 1 requests) [data]", "Row Group 1, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", @@ -221,6 +226,11 @@ fn test_read_limit() { #[test] fn test_read_single_row_filter() { + // Dictionary pages are read on first decode of a data page that needs + // them, so skipping a page does not pull its dictionary in. That is why + // the trace below shows the dictionary *after* the data page wherever a + // skip precedes it, and *before* the data page where none does. + // See apache/arrow-rs#11154. // Values from column "b" range 400..799 // filter "b" > 575 and < 625 // (last data page in Row Group 0 and first DataPage in Row Group 1) @@ -255,10 +265,10 @@ fn test_read_single_row_filter() { "Row Group 1, column 'b': DataPage(0) (113 bytes , 1 requests) [data]", "Row Group 1, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", "Event: Reader Built", - "Row Group 0, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", "Row Group 0, column 'a': DataPage(1) (126 bytes , 1 requests) [data]", - "Row Group 0, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", + "Row Group 0, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", "Row Group 0, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 0, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", "Row Group 1, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", "Row Group 1, column 'a': DataPage(0) (113 bytes , 1 requests) [data]", "Row Group 1, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", @@ -269,6 +279,11 @@ fn test_read_single_row_filter() { #[test] fn test_read_multiple_row_filter() { + // Dictionary pages are read on first decode of a data page that needs + // them, so skipping a page does not pull its dictionary in. That is why + // the trace below shows the dictionary *after* the data page wherever a + // skip precedes it, and *before* the data page where none does. + // See apache/arrow-rs#11154. // Values in column "a" range 0..399 // Values in column "b" range 400..799 // First filter: "a" > 175 (last data page in Row Group 0) @@ -305,14 +320,14 @@ fn test_read_multiple_row_filter() { "Row Group 1, column 'a': DictionaryPage (1617 bytes, 1 requests) [data]", "Row Group 1, column 'a': DataPage(0) (113 bytes , 1 requests) [data]", "Row Group 1, column 'a': DataPage(1) (126 bytes , 1 requests) [data]", - "Row Group 0, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", "Row Group 0, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 0, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", "Row Group 1, column 'b': DictionaryPage (1617 bytes, 1 requests) [data]", "Row Group 1, column 'b': DataPage(0) (113 bytes , 1 requests) [data]", "Row Group 1, column 'b': DataPage(1) (126 bytes , 1 requests) [data]", "Event: Reader Built", - "Row Group 0, column 'c': DictionaryPage (7107 bytes, 1 requests) [data]", "Row Group 0, column 'c': DataPage(1) (126 bytes , 1 requests) [data]", + "Row Group 0, column 'c': DictionaryPage (7107 bytes, 1 requests) [data]", "Row Group 1, column 'c': DictionaryPage (7217 bytes, 1 requests) [data]", "Row Group 1, column 'c': DataPage(0) (113 bytes , 1 requests) [data]", ]
