From 8c9538d2fe0d9d4869a453940edd1c49c1512834 Mon Sep 17 00:00:00 2001 From: "Oleg V. Kozlyuk" Date: Fri, 24 Jul 2026 00:16:15 +0200 Subject: [PATCH 1/4] feat(parquet): decode dictionary pages independent of the array reader Add sync and async APIs for decoding a BYTE_ARRAY column chunk dictionary page without materializing the full column, enabling exact row-group membership pruning for fully dictionary-encoded chunks. Reuse the regular page reader header, size validation, and decryption path so standalone dictionary decoding handles encrypted, truncated, and malformed pages safely. See #9010. Co-Authored-By: Claude Sonnet 5 --- parquet/src/arrow/array_reader/mod.rs | 3 + parquet/src/arrow/async_reader/mod.rs | 64 +++- parquet/src/arrow/mod.rs | 4 + parquet/src/file/metadata/dictionary.rs | 383 ++++++++++++++++++++++++ parquet/src/file/metadata/mod.rs | 41 +++ parquet/src/file/metadata/reader.rs | 81 +++++ parquet/src/file/serialized_reader.rs | 52 ++-- 7 files changed, 607 insertions(+), 21 deletions(-) create mode 100644 parquet/src/file/metadata/dictionary.rs diff --git a/parquet/src/arrow/array_reader/mod.rs b/parquet/src/arrow/array_reader/mod.rs index 32fb90d2e13d..2d9b982379b4 100644 --- a/parquet/src/arrow/array_reader/mod.rs +++ b/parquet/src/arrow/array_reader/mod.rs @@ -55,6 +55,9 @@ pub(crate) mod test_util; use crate::file::metadata::RowGroupMetaData; pub use builder::{ArrayReaderBuilder, CacheOptions, CacheOptionsBuilder}; pub use byte_array::make_byte_array_reader; +// Re-exported (beyond the `experimental` feature) so `file::metadata::dictionary` +// can PLAIN-decode a raw dictionary page without duplicating this logic. +pub(crate) use byte_array::ByteArrayDecoderPlain; pub use byte_array_dictionary::make_byte_array_dictionary_reader; #[cfg_attr(not(feature = "experimental"), expect(unused_imports))] pub use byte_view_array::make_byte_view_array_reader; diff --git a/parquet/src/arrow/async_reader/mod.rs b/parquet/src/arrow/async_reader/mod.rs index 3eddd549353b..c593823ad90d 100644 --- a/parquet/src/arrow/async_reader/mod.rs +++ b/parquet/src/arrow/async_reader/mod.rs @@ -33,7 +33,7 @@ use futures::future::{BoxFuture, FutureExt}; use futures::stream::Stream; use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt}; -use arrow_array::RecordBatch; +use arrow_array::{ArrayRef, RecordBatch}; use arrow_schema::{Schema, SchemaRef}; use crate::arrow::arrow_reader::{ @@ -670,6 +670,30 @@ impl ParquetRecordBatchStreamBuilder { self } + /// Read and decode the dictionary page for a column in a row group, if any. + /// + /// Returns `Ok(None)` if the column chunk has no dictionary page, or if + /// its physical type is not `BYTE_ARRAY` (the only physical type + /// currently supported). + /// + /// Note this does not verify that the *entire* column chunk is + /// dictionary-encoded -- callers that need that guarantee (e.g. to treat + /// the dictionary as an exhaustive set of the column's values) should + /// check the column chunk's page encoding statistics themselves. + pub async fn get_row_group_column_dictionary( + &mut self, + row_group_idx: usize, + column_idx: usize, + ) -> Result> { + ParquetMetaDataReader::read_column_dictionary_async( + &mut self.input.0, + &self.metadata, + row_group_idx, + column_idx, + ) + .await + } + /// Build a new [`ParquetRecordBatchStream`] /// /// See examples on [`ParquetRecordBatchStreamBuilder::new`] @@ -1120,6 +1144,44 @@ mod tests { ); } + #[tokio::test] + async fn test_get_row_group_column_dictionary() { + let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)])); + let values: Vec<&str> = ["alpha", "beta", "gamma"] + .iter() + .copied() + .cycle() + .take(30) + .collect(); + let array: ArrayRef = Arc::new(StringArray::from(values)); + let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(true) + .build(); + let mut buf = Vec::new(); + { + let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } + let data = Bytes::from(buf); + + let async_reader = TestReader::new(data); + let mut builder = ParquetRecordBatchStreamBuilder::new(async_reader) + .await + .unwrap(); + + let dictionary = builder + .get_row_group_column_dictionary(0, 0) + .await + .unwrap() + .unwrap(); + let dictionary = dictionary.as_any().downcast_ref::().unwrap(); + let dictionary_values: Vec<&str> = dictionary.iter().map(|v| v.unwrap()).collect(); + assert_eq!(dictionary_values, vec!["alpha", "beta", "gamma"]); + } + #[tokio::test] async fn test_async_reader_with_next_row_group() { let testdata = arrow::util::test_util::parquet_test_data(); diff --git a/parquet/src/arrow/mod.rs b/parquet/src/arrow/mod.rs index ff9924ffef40..2c225d5f0667 100644 --- a/parquet/src/arrow/mod.rs +++ b/parquet/src/arrow/mod.rs @@ -180,9 +180,13 @@ //! ``` experimental!(mod array_reader); +// Re-exported (beyond the `experimental` feature) so `file::metadata::dictionary` +// can PLAIN-decode a raw dictionary page without duplicating this logic. +pub(crate) use array_reader::ByteArrayDecoderPlain; pub mod arrow_reader; pub mod arrow_writer; mod buffer; +pub(crate) use buffer::offset_buffer::OffsetBuffer; mod decoder; #[cfg(feature = "async")] diff --git a/parquet/src/file/metadata/dictionary.rs b/parquet/src/file/metadata/dictionary.rs new file mode 100644 index 000000000000..b27c733ae3fc --- /dev/null +++ b/parquet/src/file/metadata/dictionary.rs @@ -0,0 +1,383 @@ +// 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. + +//! Decoding a column chunk's dictionary page directly into an Arrow array, +//! independent of the row-by-row [`ArrayReader`] machinery. +//! +//! This is useful for callers that want the *set* of distinct values stored +//! in a dictionary-encoded column chunk without reading any data pages, e.g. +//! to prune a row group when the query predicate's literals are known not to +//! be in the dictionary. +//! +//! [`ArrayReader`]: crate::arrow::array_reader::ArrayReader + +use crate::arrow::{ByteArrayDecoderPlain, OffsetBuffer}; +use crate::basic::{ConvertedType, LogicalType, PageType, Type as PhysicalType}; +use crate::column::page::Page; +use crate::compression::{CodecOptions, create_codec}; +#[cfg(feature = "encryption")] +use crate::encryption::decrypt::CryptoContext; +use crate::errors::{ParquetError, Result}; +#[cfg(feature = "encryption")] +use crate::file::metadata::ColumnChunkMetaData; +use crate::file::metadata::ParquetMetaData; +use crate::file::serialized_reader::{ + SerializedPageReaderContext, decode_page, read_page_header_len_from_bytes, verify_page_size, +}; +use crate::schema::types::ColumnDescriptor; +use arrow_array::ArrayRef; +use arrow_schema::DataType as ArrowType; +use bytes::Bytes; +#[cfg(feature = "encryption")] +use std::sync::Arc; + +/// Decodes the dictionary page of a column chunk into an [`ArrayRef`]. +/// +/// `buffer` must contain the entire dictionary page, byte-for-byte, i.e. the +/// range `[dictionary_page_offset, data_page_offset)` of the column chunk. +/// +/// Only `BYTE_ARRAY` columns are currently supported; other physical types +/// return an error. The returned array never contains nulls: dictionary +/// pages only store the distinct non-null values, with nulls represented via +/// definition levels in the data pages. +/// +/// Note this only decodes whatever dictionary page is present -- it does +/// **not** verify that the entire column chunk is dictionary-encoded (i.e. +/// that every value in the chunk is drawn from this dictionary). Callers +/// that need that guarantee (for example, to use the dictionary as an exact +/// membership index) must check that themselves, e.g. via +/// [`crate::file::metadata::ColumnChunkMetaData::page_encoding_stats_mask`]. +pub(crate) fn decode_dictionary_page( + buffer: Bytes, + parquet_meta_data: &ParquetMetaData, + row_group_idx: usize, + column_idx: usize, +) -> Result { + let column_metadata = parquet_meta_data + .row_group(row_group_idx) + .column(column_idx); + let column_descriptor = column_metadata.column_descr(); + + if column_descriptor.physical_type() != PhysicalType::BYTE_ARRAY { + return Err(ParquetError::General(format!( + "decode_dictionary_page only supports BYTE_ARRAY columns, got {}", + column_descriptor.physical_type() + ))); + } + + // Dictionary pages are subject to the same modular encryption as data + // pages: both the page header and the page body may be ciphertext, so + // we must route through the same crypto-aware header/data path that + // `SerializedPageReader` uses rather than parsing the header directly. + let page_context = SerializedPageReaderContext { + read_stats: true, + #[cfg(feature = "encryption")] + crypto_context: dictionary_page_crypto_context( + parquet_meta_data, + column_metadata, + row_group_idx, + column_idx, + )?, + }; + + let (consumed, header) = + read_page_header_len_from_bytes(&page_context, buffer.as_ref(), 0, true)?; + if header.r#type != PageType::DICTIONARY_PAGE { + return Err(ParquetError::General(format!( + "Expected a dictionary page, found {:?}", + header.r#type + ))); + } + + // `compressed_page_size` comes from the (possibly maliciously crafted) + // file header; `verify_page_size` bounds-checks it against what we + // actually fetched before we slice, instead of trusting it blindly. + let remaining = (buffer.len() - consumed) as u64; + verify_page_size( + header.compressed_page_size, + header.uncompressed_page_size, + remaining, + )?; + let compressed_size = header.compressed_page_size as usize; + let page_buf = buffer.slice(consumed..consumed + compressed_size); + let page_buf = page_context.decrypt_page_data(page_buf, 0, true)?; + + let mut decompressor = create_codec(column_metadata.compression(), &CodecOptions::default())?; + let page = decode_page( + header, + page_buf, + column_descriptor.physical_type(), + decompressor.as_mut(), + )?; + let Page::DictionaryPage { + buf, num_values, .. + } = page + else { + return Err(ParquetError::General( + "Expected a dictionary page".to_string(), + )); + }; + let num_values = num_values as usize; + + // The dictionary page is always PLAIN-encoded, regardless of what the + // data pages' encoding is (RLE_DICTIONARY/PLAIN_DICTIONARY only describe + // how *data* pages reference the dictionary by index). + let is_utf8 = is_utf8(column_descriptor); + let mut decoder = ByteArrayDecoderPlain::new(buf, num_values, Some(num_values), is_utf8); + let mut offsets = OffsetBuffer::::with_capacity(num_values); + decoder.read(&mut offsets, usize::MAX)?; + + let arrow_type = if is_utf8 { + ArrowType::Utf8 + } else { + ArrowType::Binary + }; + Ok(offsets.into_array(None, arrow_type)) +} + +/// Builds the crypto context needed to decrypt the dictionary page of +/// `column_metadata`, or `None` if the file (or this column) isn't encrypted. +#[cfg(feature = "encryption")] +fn dictionary_page_crypto_context( + parquet_meta_data: &ParquetMetaData, + column_metadata: &ColumnChunkMetaData, + row_group_idx: usize, + column_idx: usize, +) -> Result>> { + let Some(file_decryptor) = parquet_meta_data.file_decryptor() else { + return Ok(None); + }; + let Some(crypto_metadata) = column_metadata.crypto_metadata() else { + return Ok(None); + }; + let crypto_context = + CryptoContext::for_column(file_decryptor, crypto_metadata, row_group_idx, column_idx)? + .for_dictionary_page(); + Ok(Some(Arc::new(crypto_context))) +} + +/// Whether `column_descriptor` should be decoded as UTF-8 text (`Utf8`) +/// rather than raw `Binary`. +fn is_utf8(column_descriptor: &ColumnDescriptor) -> bool { + matches!( + column_descriptor.logical_type_ref(), + Some(LogicalType::String) | Some(LogicalType::Json) + ) || matches!( + column_descriptor.converted_type(), + ConvertedType::UTF8 | ConvertedType::JSON + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::ArrowWriter; + use crate::file::properties::WriterProperties; + use crate::file::reader::{ChunkReader, FileReader, SerializedFileReader}; + use arrow_array::{Array, RecordBatch, StringArray}; + use arrow_schema::{Field, Schema}; + use std::sync::Arc; + + fn write_dictionary_encoded_strings(values: &[&str]) -> Bytes { + let schema = Arc::new(Schema::new(vec![Field::new("s", ArrowType::Utf8, false)])); + let array = Arc::new(StringArray::from_iter_values(values.iter().copied())); + let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(true) + .build(); + let mut buf = Vec::new(); + { + let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } + Bytes::from(buf) + } + + #[test] + fn decode_dictionary_page_round_trips_strings() { + let distinct_values = ["alpha", "beta", "gamma"]; + // Repeat so the column is worth dictionary-encoding but the + // dictionary itself only contains the distinct values. + let values: Vec<&str> = distinct_values.iter().copied().cycle().take(30).collect(); + let data = write_dictionary_encoded_strings(&values); + + let reader = SerializedFileReader::new(data.clone()).unwrap(); + let metadata = reader.metadata(); + let column_metadata = metadata.row_group(0).column(0); + + assert!( + column_metadata.dictionary_page_offset().is_some(), + "expected the column chunk to be dictionary-encoded" + ); + + let start = column_metadata.dictionary_page_offset().unwrap() as u64; + let end = column_metadata.data_page_offset() as u64; + let buffer = data.get_bytes(start, (end - start) as usize).unwrap(); + + let array = decode_dictionary_page(buffer, metadata, 0, 0).unwrap(); + let array = array.as_any().downcast_ref::().unwrap(); + let decoded: Vec<&str> = array.iter().map(|v| v.unwrap()).collect(); + assert_eq!(decoded, distinct_values); + } + + #[test] + fn decode_dictionary_page_errors_on_truncated_buffer() { + let distinct_values = ["alpha", "beta", "gamma"]; + let values: Vec<&str> = distinct_values.iter().copied().cycle().take(30).collect(); + let data = write_dictionary_encoded_strings(&values); + + let reader = SerializedFileReader::new(data.clone()).unwrap(); + let metadata = reader.metadata(); + let column_metadata = metadata.row_group(0).column(0); + + let start = column_metadata.dictionary_page_offset().unwrap() as u64; + let end = column_metadata.data_page_offset() as u64; + let buffer = data.get_bytes(start, (end - start) as usize).unwrap(); + + // Simulate a truncated/malformed file: the page header's declared + // `compressed_page_size` no longer fits in what was actually + // fetched. This must return an error rather than panic while + // slicing (`Bytes::slice` panics on out-of-bounds ranges). + let truncated = buffer.slice(..buffer.len() - 1); + let err = decode_dictionary_page(truncated, metadata, 0, 0).unwrap_err(); + assert!( + matches!(err, ParquetError::EOF(_)), + "unexpected error: {err}" + ); + } + + #[test] + fn decode_dictionary_page_rejects_non_byte_array() { + let schema = Arc::new(Schema::new(vec![Field::new("i", ArrowType::Int32, false)])); + let array = Arc::new(arrow_array::Int32Array::from(vec![1, 2, 3])); + let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap(); + let mut buf = Vec::new(); + { + let mut writer = ArrowWriter::try_new(&mut buf, schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } + let data = Bytes::from(buf); + + let reader = SerializedFileReader::new(data).unwrap(); + let metadata = reader.metadata(); + + let err = decode_dictionary_page(Bytes::new(), metadata, 0, 0).unwrap_err(); + assert!(err.to_string().contains("BYTE_ARRAY")); + } + + #[test] + fn read_column_dictionary_round_trips_via_metadata_reader() { + use crate::file::metadata::ParquetMetaDataReader; + + let distinct_values = ["alpha", "beta", "gamma"]; + let values: Vec<&str> = distinct_values.iter().copied().cycle().take(30).collect(); + let data = write_dictionary_encoded_strings(&values); + + let reader = SerializedFileReader::new(data.clone()).unwrap(); + let metadata = reader.metadata(); + + let array = ParquetMetaDataReader::read_column_dictionary(&data, metadata, 0, 0) + .unwrap() + .unwrap(); + let array = array.as_any().downcast_ref::().unwrap(); + let decoded: Vec<&str> = array.iter().map(|v| v.unwrap()).collect(); + assert_eq!(decoded, distinct_values); + } + + #[cfg(feature = "encryption")] + #[test] + fn read_column_dictionary_round_trips_with_encryption() { + use crate::encryption::decrypt::FileDecryptionProperties; + use crate::encryption::encrypt::FileEncryptionProperties; + use crate::file::metadata::ParquetMetaDataReader; + + const FOOTER_KEY: &[u8] = b"0123456789012345"; + + let distinct_values = ["alpha", "beta", "gamma"]; + let values: Vec<&str> = distinct_values.iter().copied().cycle().take(30).collect(); + + let schema = Arc::new(Schema::new(vec![Field::new("s", ArrowType::Utf8, false)])); + let array = Arc::new(StringArray::from_iter_values(values.iter().copied())); + let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap(); + + let encryption_properties = FileEncryptionProperties::builder(FOOTER_KEY.to_vec()) + .build() + .unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(true) + .with_file_encryption_properties(encryption_properties) + .build(); + let mut buf = Vec::new(); + { + let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } + let data = Bytes::from(buf); + + let decryption_properties = FileDecryptionProperties::builder(FOOTER_KEY.to_vec()) + .build() + .unwrap(); + let metadata = ParquetMetaDataReader::new() + .with_decryption_properties(Some(decryption_properties)) + .parse_and_finish(&data) + .unwrap(); + + let array = ParquetMetaDataReader::read_column_dictionary(&data, &metadata, 0, 0) + .unwrap() + .unwrap(); + let array = array.as_any().downcast_ref::().unwrap(); + let decoded: Vec<&str> = array.iter().map(|v| v.unwrap()).collect(); + assert_eq!(decoded, distinct_values); + } + + #[test] + fn read_column_dictionary_returns_none_without_dictionary_page() { + use crate::file::metadata::ParquetMetaDataReader; + + let schema = Arc::new(Schema::new(vec![Field::new("s", ArrowType::Utf8, false)])); + let array = Arc::new(StringArray::from_iter_values(["a", "b", "c"])); + let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .build(); + let mut buf = Vec::new(); + { + let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } + let data = Bytes::from(buf); + + let reader = SerializedFileReader::new(data.clone()).unwrap(); + let metadata = reader.metadata(); + assert!( + metadata + .row_group(0) + .column(0) + .dictionary_page_offset() + .is_none() + ); + + let result = ParquetMetaDataReader::read_column_dictionary(&data, metadata, 0, 0).unwrap(); + assert!(result.is_none()); + } +} diff --git a/parquet/src/file/metadata/mod.rs b/parquet/src/file/metadata/mod.rs index f81ae4fb3576..fba55cbd9ef2 100644 --- a/parquet/src/file/metadata/mod.rs +++ b/parquet/src/file/metadata/mod.rs @@ -49,6 +49,47 @@ //! Please see [`external_metadata.rs`] //! //! [`external_metadata.rs`]: https://github.com/apache/arrow-rs/tree/master/parquet/examples/external_metadata.rs +//! +//! # Metadata Encodings and Structures +//! +//! There are three different encodings of Parquet Metadata in this crate: +//! +//! 1. `bytes`:encoded with the Thrift `TCompactProtocol` as defined in +//! [parquet.thrift] +//! +//! 2. [`format`]: Rust structures automatically generated by the thrift compiler +//! from [parquet.thrift]. These structures are low level and mirror +//! the thrift definitions. +//! +//! 3. [`file::metadata`] (this module): Easier to use Rust structures +//! with a more idiomatic API. Note that, confusingly, some but not all +//! of these structures have the same name as the [`format`] structures. +//! +//! [`file::metadata`]: crate::file::metadata +//! [parquet.thrift]: https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift +//! +//! Graphically, this is how the different structures relate to each other: +//! +//! ```text +//! ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ +//! ┌──────────────┐ │ ┌───────────────────────┐ │ +//! │ │ ColumnIndex │ ││ ParquetMetaData │ +//! └──────────────┘ │ └───────────────────────┘ │ +//! ┌──────────────┐ │ ┌────────────────┐ │┌───────────────────────┐ +//! │ ..0x24.. │ ◀────▶ │ OffsetIndex │ │ ◀────▶ │ ParquetMetaData │ │ +//! └──────────────┘ │ └────────────────┘ │└───────────────────────┘ +//! ... │ ... │ +//! │ ┌──────────────────┐ │ ┌──────────────────┐ +//! bytes │ FileMetaData* │ │ │ FileMetaData* │ │ +//! (thrift encoded) │ └──────────────────┘ │ └──────────────────┘ +//! ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘ +//! +//! format::meta structures file::metadata structures +//! +//! * Same name, different struct +//! ``` +#[cfg(feature = "arrow")] +mod dictionary; mod footer_tail; mod memory; mod options; diff --git a/parquet/src/file/metadata/reader.rs b/parquet/src/file/metadata/reader.rs index 018cf440cd3f..039c6e1e4945 100644 --- a/parquet/src/file/metadata/reader.rs +++ b/parquet/src/file/metadata/reader.rs @@ -19,6 +19,8 @@ use crate::encryption::decrypt::FileDecryptionProperties; use crate::errors::{ParquetError, Result}; use crate::file::FOOTER_SIZE; +#[cfg(feature = "arrow")] +use crate::file::metadata::dictionary::decode_dictionary_page; use crate::file::metadata::parser::decode_metadata; use crate::file::metadata::thrift::parquet_schema_from_bytes; use crate::file::metadata::{ @@ -26,6 +28,8 @@ use crate::file::metadata::{ }; use crate::file::reader::ChunkReader; use crate::schema::types::SchemaDescriptor; +#[cfg(feature = "arrow")] +use arrow_array::ArrayRef; use bytes::Bytes; use std::sync::Arc; use std::{io::Read, ops::Range}; @@ -469,6 +473,48 @@ impl ParquetMetaDataReader { self.load_page_index_with_remainder(fetch, None).await } + /// Reads and decodes the dictionary page of a column chunk into an Arrow array. + /// + /// Returns `Ok(None)` if the column chunk has no dictionary page, or if + /// its physical type is not `BYTE_ARRAY` (the only physical type + /// currently supported). + /// + /// Note this does not verify that the *entire* column chunk is + /// dictionary-encoded (i.e. that the dictionary contains every value in + /// the chunk) -- callers that need that guarantee should check + /// [`crate::file::metadata::ColumnChunkMetaData::page_encoding_stats_mask`] themselves. + #[cfg(feature = "arrow")] + pub fn read_column_dictionary( + reader: &R, + metadata: &ParquetMetaData, + row_group_idx: usize, + column_idx: usize, + ) -> Result> { + let Some((start, end)) = dictionary_page_byte_range(metadata, row_group_idx, column_idx)? + else { + return Ok(None); + }; + let length = usize::try_from(end - start)?; + let buffer = reader.get_bytes(start, length)?; + decode_dictionary_page(buffer, metadata, row_group_idx, column_idx).map(Some) + } + + /// Asynchronous version of [`Self::read_column_dictionary`]. + #[cfg(all(feature = "async", feature = "arrow"))] + pub async fn read_column_dictionary_async( + mut fetch: F, + metadata: &ParquetMetaData, + row_group_idx: usize, + column_idx: usize, + ) -> Result> { + let Some((start, end)) = dictionary_page_byte_range(metadata, row_group_idx, column_idx)? + else { + return Ok(None); + }; + let buffer = fetch.fetch(start..end).await?; + decode_dictionary_page(buffer, metadata, row_group_idx, column_idx).map(Some) + } + #[cfg(all(feature = "async", feature = "arrow"))] async fn load_page_index_with_remainder( &mut self, @@ -838,6 +884,41 @@ fn parse_index_data(push_decoder: &mut ParquetMetaDataPushDecoder) -> Result Result> { + let column_metadata = metadata.row_group(row_group_idx).column(column_idx); + let column_descriptor = column_metadata.column_descr(); + + if column_descriptor.physical_type() != crate::basic::Type::BYTE_ARRAY { + return Ok(None); + } + let Some(start) = column_metadata.dictionary_page_offset() else { + return Ok(None); + }; + let start: u64 = start + .try_into() + .map_err(|_| ParquetError::General("Dictionary page offset is invalid".to_string()))?; + let end: u64 = column_metadata + .data_page_offset() + .try_into() + .map_err(|_| ParquetError::General("Data page offset is invalid".to_string()))?; + if end < start { + return Err(ParquetError::General( + "Data page offset precedes dictionary page offset".to_string(), + )); + } + + Ok(Some((start, end))) +} + #[cfg(test)] mod tests { use super::*; diff --git a/parquet/src/file/serialized_reader.rs b/parquet/src/file/serialized_reader.rs index bc525a78aba7..89e6fe449a3c 100644 --- a/parquet/src/file/serialized_reader.rs +++ b/parquet/src/file/serialized_reader.rs @@ -567,12 +567,12 @@ enum SerializedPageReaderState { } #[derive(Default)] -struct SerializedPageReaderContext { +pub(crate) struct SerializedPageReaderContext { /// Controls decoding of page-level statistics - read_stats: bool, + pub(crate) read_stats: bool, /// Crypto context carrying objects required for decryption #[cfg(feature = "encryption")] - crypto_context: Option>, + pub(crate) crypto_context: Option>, } /// A serialized implementation for Parquet [`PageReader`]. @@ -785,23 +785,30 @@ impl SerializedPageReader { let header = context.read_page_header(&mut tracked, page_index, dictionary_page)?; Ok((tracked.bytes_read, header)) } +} - fn read_page_header_len_from_bytes( - context: &SerializedPageReaderContext, - buffer: &[u8], - page_index: usize, - dictionary_page: bool, - ) -> Result<(usize, PageHeader)> { - let mut input = std::io::Cursor::new(buffer); - let header = context.read_page_header(&mut input, page_index, dictionary_page)?; - let header_len = input.position() as usize; - Ok((header_len, header)) - } +/// Reads (and decrypts, if `context` carries a crypto context) the page header stored +/// at the front of `buffer`, returning the header and the number of bytes of `buffer` +/// it occupies. +/// +/// This is exposed for callers that need to decode a single page directly from an +/// already-fetched byte range, outside of the normal [`SerializedPageReader`] iteration +/// -- e.g. decoding a dictionary page standalone. +pub(crate) fn read_page_header_len_from_bytes( + context: &SerializedPageReaderContext, + buffer: &[u8], + page_index: usize, + dictionary_page: bool, +) -> Result<(usize, PageHeader)> { + let mut input = std::io::Cursor::new(buffer); + let header = context.read_page_header(&mut input, page_index, dictionary_page)?; + let header_len = input.position() as usize; + Ok((header_len, header)) } #[cfg(not(feature = "encryption"))] impl SerializedPageReaderContext { - fn read_page_header( + pub(crate) fn read_page_header( &self, input: &mut T, _page_index: usize, @@ -815,7 +822,7 @@ impl SerializedPageReaderContext { } } - fn decrypt_page_data( + pub(crate) fn decrypt_page_data( &self, buffer: T, _page_index: usize, @@ -827,7 +834,7 @@ impl SerializedPageReaderContext { #[cfg(feature = "encryption")] impl SerializedPageReaderContext { - fn read_page_header( + pub(crate) fn read_page_header( &self, input: &mut T, page_index: usize, @@ -865,7 +872,12 @@ impl SerializedPageReaderContext { } } - fn decrypt_page_data(&self, buffer: T, page_index: usize, dictionary_page: bool) -> Result + pub(crate) fn decrypt_page_data( + &self, + buffer: T, + page_index: usize, + dictionary_page: bool, + ) -> Result where T: AsRef<[u8]>, T: From>, @@ -911,7 +923,7 @@ fn verify_page_header_len(header_len: usize, remaining_bytes: u64) -> Result<()> Ok(()) } -fn verify_page_size( +pub(crate) fn verify_page_size( compressed_size: i32, uncompressed_size: i32, remaining_bytes: u64, @@ -1005,7 +1017,7 @@ impl PageReader for SerializedPageReader { 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) = Self::read_page_header_len_from_bytes( + let (offset, header) = read_page_header_len_from_bytes( &self.context, buffer.as_ref(), *page_index, From d1f6bb24db2a9c63ad2a715b4b76fff7664c554b Mon Sep 17 00:00:00 2001 From: "Oleg V. Kozlyuk" Date: Sat, 15 Aug 2026 18:21:32 +0200 Subject: [PATCH 2/4] test(parquet): demonstrate dictionary row group pruning --- parquet/src/arrow/async_reader/mod.rs | 87 ++++++++++++++++++++++++++- parquet/src/file/metadata/reader.rs | 5 +- 2 files changed, 90 insertions(+), 2 deletions(-) diff --git a/parquet/src/arrow/async_reader/mod.rs b/parquet/src/arrow/async_reader/mod.rs index c593823ad90d..0f039bcd1d43 100644 --- a/parquet/src/arrow/async_reader/mod.rs +++ b/parquet/src/arrow/async_reader/mod.rs @@ -676,10 +676,15 @@ impl ParquetRecordBatchStreamBuilder { /// its physical type is not `BYTE_ARRAY` (the only physical type /// currently supported). /// + /// This can be used to inspect dictionary values when selecting or pruning + /// row groups before passing the selected indices to + /// [`ParquetRecordBatchStreamBuilder::with_row_groups`]. + /// /// Note this does not verify that the *entire* column chunk is /// dictionary-encoded -- callers that need that guarantee (e.g. to treat /// the dictionary as an exhaustive set of the column's values) should - /// check the column chunk's page encoding statistics themselves. + /// check + /// [`crate::file::metadata::ColumnChunkMetaData::page_encoding_stats_mask`]. pub async fn get_row_group_column_dictionary( &mut self, row_group_idx: usize, @@ -991,6 +996,7 @@ mod tests { use crate::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; use crate::arrow::schema::virtual_type::RowNumber; use crate::arrow::{ArrowWriter, AsyncArrowWriter, ProjectionMask}; + use crate::basic::Encoding; use crate::file::metadata::ParquetMetaDataReader; use crate::file::metadata::{PageIndex, PageIndexPolicy}; use crate::file::properties::WriterProperties; @@ -1182,6 +1188,85 @@ mod tests { assert_eq!(dictionary_values, vec!["alpha", "beta", "gamma"]); } + #[tokio::test] + async fn test_dictionary_selects_row_groups_without_reading_skipped_data() { + let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)])); + let row_group_values = ["skip", "target"]; + let props = WriterProperties::builder() + .set_dictionary_enabled(true) + .build(); + let mut buf = Vec::new(); + { + let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), Some(props)).unwrap(); + for value in row_group_values { + let array: ArrayRef = Arc::new(StringArray::from(vec![value; 30])); + let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap(); + writer.write(&batch).unwrap(); + writer.flush().unwrap(); + } + writer.close().unwrap(); + } + + let async_reader = TestReader::new(Bytes::from(buf)); + let requests = async_reader.requests.clone(); + let mut builder = ParquetRecordBatchStreamBuilder::new(async_reader) + .await + .unwrap(); + let metadata = builder.metadata().clone(); + assert_eq!(metadata.num_row_groups(), 2); + + let mut selected_row_groups = Vec::new(); + let mut dictionary_ranges = Vec::new(); + for row_group_idx in 0..metadata.num_row_groups() { + let column = metadata.row_group(row_group_idx).column(0); + let encoding_mask = column.page_encoding_stats_mask().unwrap(); + assert!( + encoding_mask.is_only(Encoding::PLAIN_DICTIONARY) + || encoding_mask.is_only(Encoding::RLE_DICTIONARY) + ); + + let dictionary_start = column.dictionary_page_offset().unwrap() as usize; + let data_start = column.data_page_offset() as usize; + dictionary_ranges.push(dictionary_start..data_start); + + let dictionary = builder + .get_row_group_column_dictionary(row_group_idx, 0) + .await + .unwrap() + .unwrap(); + let dictionary = dictionary.as_string::(); + if dictionary.iter().any(|value| value == Some("target")) { + selected_row_groups.push(row_group_idx); + } + } + assert_eq!(selected_row_groups, vec![1]); + + let batches: Vec<_> = builder + .with_row_groups(selected_row_groups) + .build() + .unwrap() + .try_collect() + .await + .unwrap(); + let values: Vec<_> = batches + .iter() + .flat_map(|batch| batch.column(0).as_string::().iter()) + .collect(); + assert_eq!(values, vec![Some("target"); 30]); + + let skipped_column = metadata.row_group(0).column(0); + let (skipped_start, skipped_len) = skipped_column.byte_range(); + let skipped_data_range = + skipped_column.data_page_offset() as usize..(skipped_start + skipped_len) as usize; + let requests = requests.lock().unwrap(); + for dictionary_range in dictionary_ranges { + assert!(requests.contains(&dictionary_range)); + } + assert!(requests.iter().all(|request| { + request.end <= skipped_data_range.start || request.start >= skipped_data_range.end + })); + } + #[tokio::test] async fn test_async_reader_with_next_row_group() { let testdata = arrow::util::test_util::parquet_test_data(); diff --git a/parquet/src/file/metadata/reader.rs b/parquet/src/file/metadata/reader.rs index 039c6e1e4945..f2803aee4f96 100644 --- a/parquet/src/file/metadata/reader.rs +++ b/parquet/src/file/metadata/reader.rs @@ -479,10 +479,13 @@ impl ParquetMetaDataReader { /// its physical type is not `BYTE_ARRAY` (the only physical type /// currently supported). /// + /// This can be used to inspect dictionary values when selecting or pruning + /// row groups before reading their data pages. + /// /// Note this does not verify that the *entire* column chunk is /// dictionary-encoded (i.e. that the dictionary contains every value in /// the chunk) -- callers that need that guarantee should check - /// [`crate::file::metadata::ColumnChunkMetaData::page_encoding_stats_mask`] themselves. + /// [`crate::file::metadata::ColumnChunkMetaData::page_encoding_stats_mask`]. #[cfg(feature = "arrow")] pub fn read_column_dictionary( reader: &R, From 90eb024d743f3313f143a3a74e62db2b6dd93ff0 Mon Sep 17 00:00:00 2001 From: "Oleg V. Kozlyuk" Date: Sat, 22 Aug 2026 22:46:33 +0200 Subject: [PATCH 3/4] fix(parquet): nest or-patterns in is_utf8 to satisfy clippy Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01Y1dzYmhkPzEC6H48UVneAJ --- parquet/src/file/metadata/dictionary.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/parquet/src/file/metadata/dictionary.rs b/parquet/src/file/metadata/dictionary.rs index b27c733ae3fc..049e94144e31 100644 --- a/parquet/src/file/metadata/dictionary.rs +++ b/parquet/src/file/metadata/dictionary.rs @@ -175,7 +175,7 @@ fn dictionary_page_crypto_context( fn is_utf8(column_descriptor: &ColumnDescriptor) -> bool { matches!( column_descriptor.logical_type_ref(), - Some(LogicalType::String) | Some(LogicalType::Json) + Some(LogicalType::String | LogicalType::Json) ) || matches!( column_descriptor.converted_type(), ConvertedType::UTF8 | ConvertedType::JSON From b00bc69e557d4af0e425d8bed01e5969524b2dbf Mon Sep 17 00:00:00 2001 From: "Oleg V. Kozlyuk" Date: Wed, 2 Sep 2026 07:54:00 +0200 Subject: [PATCH 4/4] Apply suggestion from @etseidl Co-authored-by: Ed Seidl --- parquet/src/file/metadata/mod.rs | 38 -------------------------------- 1 file changed, 38 deletions(-) diff --git a/parquet/src/file/metadata/mod.rs b/parquet/src/file/metadata/mod.rs index fba55cbd9ef2..311bc31c54ca 100644 --- a/parquet/src/file/metadata/mod.rs +++ b/parquet/src/file/metadata/mod.rs @@ -50,44 +50,6 @@ //! //! [`external_metadata.rs`]: https://github.com/apache/arrow-rs/tree/master/parquet/examples/external_metadata.rs //! -//! # Metadata Encodings and Structures -//! -//! There are three different encodings of Parquet Metadata in this crate: -//! -//! 1. `bytes`:encoded with the Thrift `TCompactProtocol` as defined in -//! [parquet.thrift] -//! -//! 2. [`format`]: Rust structures automatically generated by the thrift compiler -//! from [parquet.thrift]. These structures are low level and mirror -//! the thrift definitions. -//! -//! 3. [`file::metadata`] (this module): Easier to use Rust structures -//! with a more idiomatic API. Note that, confusingly, some but not all -//! of these structures have the same name as the [`format`] structures. -//! -//! [`file::metadata`]: crate::file::metadata -//! [parquet.thrift]: https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift -//! -//! Graphically, this is how the different structures relate to each other: -//! -//! ```text -//! ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ -//! ┌──────────────┐ │ ┌───────────────────────┐ │ -//! │ │ ColumnIndex │ ││ ParquetMetaData │ -//! └──────────────┘ │ └───────────────────────┘ │ -//! ┌──────────────┐ │ ┌────────────────┐ │┌───────────────────────┐ -//! │ ..0x24.. │ ◀────▶ │ OffsetIndex │ │ ◀────▶ │ ParquetMetaData │ │ -//! └──────────────┘ │ └────────────────┘ │└───────────────────────┘ -//! ... │ ... │ -//! │ ┌──────────────────┐ │ ┌──────────────────┐ -//! bytes │ FileMetaData* │ │ │ FileMetaData* │ │ -//! (thrift encoded) │ └──────────────────┘ │ └──────────────────┘ -//! ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘ -//! -//! format::meta structures file::metadata structures -//! -//! * Same name, different struct -//! ``` #[cfg(feature = "arrow")] mod dictionary; mod footer_tail;