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/6] 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/6] 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/6] 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/6] 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; From f3377f353d5ff9c6acdba68c5d35803656547b85 Mon Sep 17 00:00:00 2001 From: "Oleg V. Kozlyuk" Date: Fri, 18 Sep 2026 20:35:02 +0200 Subject: [PATCH 5/6] Address review notes --- parquet/src/arrow/async_reader/mod.rs | 20 ++++++++++++++++---- parquet/src/file/metadata/reader.rs | 16 +++++++--------- 2 files changed, 23 insertions(+), 13 deletions(-) diff --git a/parquet/src/arrow/async_reader/mod.rs b/parquet/src/arrow/async_reader/mod.rs index c91404a72545..54e890245777 100644 --- a/parquet/src/arrow/async_reader/mod.rs +++ b/parquet/src/arrow/async_reader/mod.rs @@ -690,7 +690,7 @@ impl ParquetRecordBatchStreamBuilder { /// the dictionary as an exhaustive set of the column's values) should /// check /// [`crate::file::metadata::ColumnChunkMetaData::page_encoding_stats_mask`]. - pub async fn get_row_group_column_dictionary( + pub async fn get_column_chunk_dictionary( &mut self, row_group_idx: usize, column_idx: usize, @@ -1157,7 +1157,7 @@ mod tests { } #[tokio::test] - async fn test_get_row_group_column_dictionary() { + async fn test_get_column_chunk_dictionary() { let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)])); let values: Vec<&str> = ["alpha", "beta", "gamma"] .iter() @@ -1185,7 +1185,7 @@ mod tests { .unwrap(); let dictionary = builder - .get_row_group_column_dictionary(0, 0) + .get_column_chunk_dictionary(0, 0) .await .unwrap() .unwrap(); @@ -1194,8 +1194,13 @@ mod tests { assert_eq!(dictionary_values, vec!["alpha", "beta", "gamma"]); } + // This test demonstrates row group pruning using dictionary pages, + // verifying that data pages of skipped row groups are not read #[tokio::test] async fn test_dictionary_selects_row_groups_without_reading_skipped_data() { + // Write two row groups, each dictionary-encoded and containing a single + // distinct string repeated 30 times: row group 0 is all "skip", row + // group 1 is all "target". let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)])); let row_group_values = ["skip", "target"]; let props = WriterProperties::builder() @@ -1214,6 +1219,8 @@ mod tests { } let async_reader = TestReader::new(Bytes::from(buf)); + // `requests` records the byte ranges fetched from the underlying reader, + // which we later use to check that we read only the needed page let requests = async_reader.requests.clone(); let mut builder = ParquetRecordBatchStreamBuilder::new(async_reader) .await @@ -1221,6 +1228,8 @@ mod tests { let metadata = builder.metadata().clone(); assert_eq!(metadata.num_row_groups(), 2); + // For each row group, fetch and decode just the dictionary page + // to decide whether to read the whole row group let mut selected_row_groups = Vec::new(); let mut dictionary_ranges = Vec::new(); for row_group_idx in 0..metadata.num_row_groups() { @@ -1236,7 +1245,7 @@ mod tests { dictionary_ranges.push(dictionary_start..data_start); let dictionary = builder - .get_row_group_column_dictionary(row_group_idx, 0) + .get_column_chunk_dictionary(row_group_idx, 0) .await .unwrap() .unwrap(); @@ -1247,6 +1256,7 @@ mod tests { } assert_eq!(selected_row_groups, vec![1]); + // Read the filtered row group and verify the values let batches: Vec<_> = builder .with_row_groups(selected_row_groups) .build() @@ -1260,6 +1270,8 @@ mod tests { .collect(); assert_eq!(values, vec![Some("target"); 30]); + // Finally, verify none of the requests overlapped the data pages + // of the skipped row group let skipped_column = metadata.row_group(0).column(0); let (skipped_start, skipped_len) = skipped_column.byte_range(); let skipped_data_range = diff --git a/parquet/src/file/metadata/reader.rs b/parquet/src/file/metadata/reader.rs index c8f226bf7ab6..37e9fe900c89 100644 --- a/parquet/src/file/metadata/reader.rs +++ b/parquet/src/file/metadata/reader.rs @@ -496,12 +496,11 @@ impl ParquetMetaDataReader { row_group_idx: usize, column_idx: usize, ) -> Result> { - let Some((start, end)) = dictionary_page_byte_range(metadata, row_group_idx, column_idx)? - else { + let Some(range) = 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)?; + let length = usize::try_from(range.end - range.start)?; + let buffer = reader.get_bytes(range.start, length)?; decode_dictionary_page(buffer, metadata, row_group_idx, column_idx).map(Some) } @@ -513,11 +512,10 @@ impl ParquetMetaDataReader { row_group_idx: usize, column_idx: usize, ) -> Result> { - let Some((start, end)) = dictionary_page_byte_range(metadata, row_group_idx, column_idx)? - else { + let Some(range) = dictionary_page_byte_range(metadata, row_group_idx, column_idx)? else { return Ok(None); }; - let buffer = fetch.fetch(start..end).await?; + let buffer = fetch.fetch(range).await?; decode_dictionary_page(buffer, metadata, row_group_idx, column_idx).map(Some) } @@ -899,7 +897,7 @@ fn dictionary_page_byte_range( metadata: &ParquetMetaData, row_group_idx: usize, column_idx: usize, -) -> Result> { +) -> Result>> { let column_metadata = metadata.row_group(row_group_idx).column(column_idx); let column_descriptor = column_metadata.column_descr(); @@ -922,7 +920,7 @@ fn dictionary_page_byte_range( )); } - Ok(Some((start, end))) + Ok(Some(start..end)) } #[cfg(test)] From fecbc60704f2857b6d4f4a863467b6e9de731654 Mon Sep 17 00:00:00 2001 From: "Oleg V. Kozlyuk" Date: Fri, 25 Sep 2026 20:30:49 +0200 Subject: [PATCH 6/6] Fix dictionary page decoding and encrypted row group lookup --- parquet/src/arrow/async_reader/mod.rs | 40 ++++- parquet/src/file/metadata/dictionary.rs | 191 +++++++++++++++++++----- parquet/src/file/metadata/reader.rs | 4 + 3 files changed, 190 insertions(+), 45 deletions(-) diff --git a/parquet/src/arrow/async_reader/mod.rs b/parquet/src/arrow/async_reader/mod.rs index 54e890245777..7d5a07f36985 100644 --- a/parquet/src/arrow/async_reader/mod.rs +++ b/parquet/src/arrow/async_reader/mod.rs @@ -681,6 +681,10 @@ impl ParquetRecordBatchStreamBuilder { /// its physical type is not `BYTE_ARRAY` (the only physical type /// currently supported). /// + /// The returned array contains raw `Binary` values, even for columns + /// annotated as strings. Callers can compare byte slices directly or + /// convert values to UTF-8 explicitly. + /// /// This can be used to inspect dictionary values when selecting or pruning /// row groups before passing the selected indices to /// [`ParquetRecordBatchStreamBuilder::with_row_groups`]. @@ -1012,8 +1016,8 @@ mod tests { use arrow_array::cast::AsArray; use arrow_array::types::Int32Type; use arrow_array::{ - Array, ArrayRef, BooleanArray, Int32Array, RecordBatchReader, Scalar, StringArray, - StructArray, UInt64Array, + Array, ArrayRef, BinaryArray, BooleanArray, Int32Array, RecordBatchReader, Scalar, + StringArray, StructArray, UInt64Array, }; use arrow_schema::{DataType, Field, Schema}; use futures::{StreamExt, TryStreamExt}; @@ -1179,6 +1183,22 @@ mod tests { } let data = Bytes::from(buf); + let direct_metadata = ParquetMetaDataReader::new() + .parse_and_finish(&data) + .unwrap(); + let mut direct_reader = TestReader::new(data.clone()); + let direct = ParquetMetaDataReader::read_column_dictionary_async( + &mut direct_reader, + &direct_metadata, + 0, + 0, + ) + .await + .unwrap() + .unwrap(); + let direct = direct.as_any().downcast_ref::().unwrap(); + assert_eq!(direct.value(0), b"alpha"); + let async_reader = TestReader::new(data); let mut builder = ParquetRecordBatchStreamBuilder::new(async_reader) .await @@ -1189,9 +1209,12 @@ mod tests { .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"]); + let dictionary = dictionary.as_any().downcast_ref::().unwrap(); + let dictionary_values: Vec<&[u8]> = dictionary.iter().map(|v| v.unwrap()).collect(); + assert_eq!( + dictionary_values, + vec![b"alpha".as_slice(), b"beta", b"gamma"] + ); } // This test demonstrates row group pruning using dictionary pages, @@ -1249,8 +1272,11 @@ mod tests { .await .unwrap() .unwrap(); - let dictionary = dictionary.as_string::(); - if dictionary.iter().any(|value| value == Some("target")) { + let dictionary = dictionary.as_binary::(); + if dictionary + .iter() + .any(|value| value == Some(b"target".as_slice())) + { selected_row_groups.push(row_group_idx); } } diff --git a/parquet/src/file/metadata/dictionary.rs b/parquet/src/file/metadata/dictionary.rs index 049e94144e31..809fbff355e7 100644 --- a/parquet/src/file/metadata/dictionary.rs +++ b/parquet/src/file/metadata/dictionary.rs @@ -26,7 +26,7 @@ //! [`ArrayReader`]: crate::arrow::array_reader::ArrayReader use crate::arrow::{ByteArrayDecoderPlain, OffsetBuffer}; -use crate::basic::{ConvertedType, LogicalType, PageType, Type as PhysicalType}; +use crate::basic::{Encoding, PageType, Type as PhysicalType}; use crate::column::page::Page; use crate::compression::{CodecOptions, create_codec}; #[cfg(feature = "encryption")] @@ -38,7 +38,6 @@ 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; @@ -54,6 +53,7 @@ use std::sync::Arc; /// 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. +/// The returned array has `Binary` values, regardless of the column's logical type. /// /// Note this only decodes whatever dictionary page is present -- it does /// **not** verify that the entire column chunk is dictionary-encoded (i.e. @@ -124,7 +124,10 @@ pub(crate) fn decode_dictionary_page( decompressor.as_mut(), )?; let Page::DictionaryPage { - buf, num_values, .. + buf, + num_values, + encoding, + .. } = page else { return Err(ParquetError::General( @@ -136,17 +139,22 @@ pub(crate) fn decode_dictionary_page( // 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); + if encoding != Encoding::PLAIN { + return Err(ParquetError::General(format!( + "Dictionary page encoding must be PLAIN, got {encoding:?}" + ))); + } + let mut decoder = ByteArrayDecoderPlain::new(buf, num_values, Some(num_values), false); let mut offsets = OffsetBuffer::::with_capacity(num_values); decoder.read(&mut offsets, usize::MAX)?; + if offsets.len() != num_values { + return Err(ParquetError::General(format!( + "Expected {num_values} dictionary values, decoded {}", + offsets.len() + ))); + } - let arrow_type = if is_utf8 { - ArrowType::Utf8 - } else { - ArrowType::Binary - }; - Ok(offsets.into_array(None, arrow_type)) + Ok(offsets.into_array(None, ArrowType::Binary)) } /// Builds the crypto context needed to decrypt the dictionary page of @@ -164,31 +172,31 @@ fn dictionary_page_crypto_context( let Some(crypto_metadata) = column_metadata.crypto_metadata() else { return Ok(None); }; + let ordinal = parquet_meta_data + .row_group(row_group_idx) + .ordinal() + .ok_or_else(|| { + ParquetError::General("Encrypted row group is missing its file ordinal".to_string()) + })?; + let ordinal = usize::try_from(ordinal).map_err(|_| { + ParquetError::General("Encrypted row group has an invalid file ordinal".to_string()) + })?; let crypto_context = - CryptoContext::for_column(file_decryptor, crypto_metadata, row_group_idx, column_idx)? + CryptoContext::for_column(file_decryptor, crypto_metadata, ordinal, 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 | LogicalType::Json) - ) || matches!( - column_descriptor.converted_type(), - ConvertedType::UTF8 | ConvertedType::JSON - ) -} - #[cfg(test)] mod tests { use super::*; use crate::arrow::ArrowWriter; + use crate::basic::Encoding; + use crate::file::metadata::ParquetMetaDataReader; use crate::file::properties::WriterProperties; use crate::file::reader::{ChunkReader, FileReader, SerializedFileReader}; - use arrow_array::{Array, RecordBatch, StringArray}; + use crate::parquet_thrift::{ThriftCompactOutputProtocol, WriteThrift}; + use arrow_array::{Array, BinaryArray, RecordBatch, StringArray}; use arrow_schema::{Field, Schema}; use std::sync::Arc; @@ -231,9 +239,9 @@ mod tests { 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); + let array = array.as_any().downcast_ref::().unwrap(); + let decoded: Vec<&[u8]> = array.iter().map(|v| v.unwrap()).collect(); + assert_eq!(decoded, distinct_values.map(str::as_bytes)); } #[test] @@ -262,6 +270,52 @@ mod tests { ); } + fn dictionary_page_with_header_change( + data: &Bytes, + change: impl FnOnce(&mut crate::file::metadata::thrift::PageHeader), + ) -> (Bytes, ParquetMetaData) { + let reader = SerializedFileReader::new(data.clone()).unwrap(); + let metadata = reader.metadata().clone(); + let column = metadata.row_group(0).column(0); + let start = column.dictionary_page_offset().unwrap() as u64; + let end = column.data_page_offset() as u64; + let buffer = data.get_bytes(start, (end - start) as usize).unwrap(); + let context = SerializedPageReaderContext { + read_stats: true, + #[cfg(feature = "encryption")] + crypto_context: None, + }; + let (header_len, mut header) = + read_page_header_len_from_bytes(&context, &buffer, 0, true).unwrap(); + change(&mut header); + let mut changed = Vec::new(); + header + .write_thrift(&mut ThriftCompactOutputProtocol::new(&mut changed)) + .unwrap(); + changed.extend_from_slice(&buffer[header_len..]); + (Bytes::from(changed), metadata) + } + + #[test] + fn decode_dictionary_page_rejects_missing_values() { + let data = write_dictionary_encoded_strings(&["alpha", "beta", "alpha"]); + let (buffer, metadata) = dictionary_page_with_header_change(&data, |header| { + header.dictionary_page_header.as_mut().unwrap().num_values += 1; + }); + let err = decode_dictionary_page(buffer, &metadata, 0, 0).unwrap_err(); + assert!(err.to_string().contains("dictionary values"), "{err}"); + } + + #[test] + fn decode_dictionary_page_rejects_non_plain_encoding() { + let data = write_dictionary_encoded_strings(&["alpha", "beta", "alpha"]); + let (buffer, metadata) = dictionary_page_with_header_change(&data, |header| { + header.dictionary_page_header.as_mut().unwrap().encoding = Encoding::RLE_DICTIONARY; + }); + let err = decode_dictionary_page(buffer, &metadata, 0, 0).unwrap_err(); + assert!(err.to_string().contains("PLAIN"), "{err}"); + } + #[test] fn decode_dictionary_page_rejects_non_byte_array() { let schema = Arc::new(Schema::new(vec![Field::new("i", ArrowType::Int32, false)])); @@ -284,8 +338,6 @@ mod tests { #[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); @@ -296,9 +348,9 @@ mod tests { 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); + let array = array.as_any().downcast_ref::().unwrap(); + let decoded: Vec<&[u8]> = array.iter().map(|v| v.unwrap()).collect(); + assert_eq!(decoded, distinct_values.map(str::as_bytes)); } #[cfg(feature = "encryption")] @@ -306,8 +358,6 @@ mod tests { 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"]; @@ -343,9 +393,74 @@ mod tests { 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); + let array = array.as_any().downcast_ref::().unwrap(); + let decoded: Vec<&[u8]> = array.iter().map(|v| v.unwrap()).collect(); + assert_eq!(decoded, distinct_values.map(str::as_bytes)); + } + + #[cfg(feature = "encryption")] + #[test] + fn read_column_dictionary_uses_file_ordinal_after_filtering() { + use crate::encryption::decrypt::FileDecryptionProperties; + use crate::encryption::encrypt::FileEncryptionProperties; + use crate::file::metadata::ParquetMetaDataBuilder; + + const FOOTER_KEY: &[u8] = b"0123456789012345"; + let schema = Arc::new(Schema::new(vec![Field::new("s", ArrowType::Utf8, false)])); + let props = WriterProperties::builder() + .set_dictionary_enabled(true) + .with_file_encryption_properties( + FileEncryptionProperties::builder(FOOTER_KEY.to_vec()) + .build() + .unwrap(), + ) + .build(); + let mut buf = Vec::new(); + { + let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), Some(props)).unwrap(); + for value in ["first", "second"] { + let array = Arc::new(StringArray::from_iter_values(std::iter::repeat_n( + value, 30, + ))); + let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap(); + writer.write(&batch).unwrap(); + writer.flush().unwrap(); + } + writer.close().unwrap(); + } + let data = Bytes::from(buf); + let metadata = ParquetMetaDataReader::new() + .with_decryption_properties(Some( + FileDecryptionProperties::builder(FOOTER_KEY.to_vec()) + .build() + .unwrap(), + )) + .parse_and_finish(&data) + .unwrap(); + assert_eq!(metadata.row_group(1).ordinal(), Some(1)); + let second = metadata.row_group(1).clone(); + let filtered = ParquetMetaDataBuilder::new_from_metadata(metadata) + .set_row_groups(vec![second]) + .build(); + let array = ParquetMetaDataReader::read_column_dictionary(&data, &filtered, 0, 0) + .unwrap() + .unwrap(); + let array = array.as_any().downcast_ref::().unwrap(); + assert_eq!(array.value(0), b"second"); + + let invalid = filtered + .row_group(0) + .clone() + .into_builder() + .set_ordinal(-1) + .build() + .unwrap(); + let invalid_metadata = ParquetMetaDataBuilder::new_from_metadata(filtered) + .set_row_groups(vec![invalid]) + .build(); + let err = ParquetMetaDataReader::read_column_dictionary(&data, &invalid_metadata, 0, 0) + .unwrap_err(); + assert!(err.to_string().contains("invalid file ordinal"), "{err}"); } #[test] diff --git a/parquet/src/file/metadata/reader.rs b/parquet/src/file/metadata/reader.rs index 37e9fe900c89..e2bd138e50e4 100644 --- a/parquet/src/file/metadata/reader.rs +++ b/parquet/src/file/metadata/reader.rs @@ -482,6 +482,10 @@ impl ParquetMetaDataReader { /// its physical type is not `BYTE_ARRAY` (the only physical type /// currently supported). /// + /// The returned array contains raw `Binary` values, even for columns + /// annotated as strings. Callers can compare byte slices directly or + /// convert values to UTF-8 explicitly. + /// /// This can be used to inspect dictionary values when selecting or pruning /// row groups before reading their data pages. ///