zhuqi-lucas commented on code in PR #11168:
URL: https://github.com/apache/arrow-rs/pull/11168#discussion_r4278893887
##########
parquet/src/column/page.rs:
##########
@@ -404,6 +404,31 @@ 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`]. The default preserves compatibility with
existing
Review Comment:
Thanks, removed it.
##########
parquet/src/column/page.rs:
##########
@@ -404,6 +404,39 @@ pub trait PageReader: Iterator<Item = Result<Page>> + Send
{
/// column index information
fn skip_next_page(&mut self) -> Result<()>;
+ /// Decodes and returns a dictionary page this reader has previously
+ /// skipped past, if any.
+ ///
+ /// [`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.
Review Comment:
Thanks, reworded it to say it returns the decoded dictionary `Page`.
##########
parquet/src/file/serialized_reader.rs:
##########
@@ -2911,4 +3020,439 @@ 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]);
+ }
+
+ /// Prints the cost the deferral avoids, on a chunk shaped like a real
Review Comment:
Good point, thanks. Moved it to a criterion bench
(`benches/dictionary_skip.rs`). Skipping a whole chunk takes ~614µs on main and
~1.1µs with this PR.
##########
parquet/tests/arrow_reader/io/sync_reader.rs:
##########
@@ -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]",
Review Comment:
Thanks! The dictionary page is now fetched only when a data page actually
needs it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
