Defer dictionary page decode on the skip path

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"]

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);

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

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",