From 8145bb8ef9dcdddcbe22f90285c9218a5bec9b25 Mon Sep 17 00:00:00 2001 From: "Oleg V. Kozlyuk" Date: Fri, 24 Jul 2026 00:16:15 +0200 Subject: [PATCH] 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 | 2 + parquet/src/file/metadata/reader.rs | 81 +++++ parquet/src/file/serialized_reader.rs | 52 ++-- 7 files changed, 568 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 347058329f38..cd0e5cb0e700 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; #[allow(unused_imports)] // Only used for benchmarks 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 5a0083b7164d..f886b328c7b2 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::{ @@ -570,6 +570,30 @@ impl ParquetRecordBatchStreamBuilder { Ok(Some(Sbbf::new(&bitset))) } + /// 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`] @@ -984,6 +1008,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 b89db361eda1..6dffb53b471b 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 646438d2e979..1910466b78ca 100644 --- a/parquet/src/file/metadata/mod.rs +++ b/parquet/src/file/metadata/mod.rs @@ -86,6 +86,8 @@ //! //! * 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 844ae747c7c9..6cab133e5172 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}; @@ -470,6 +474,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, @@ -840,6 +886,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 113e5203c94b..bf5b32df7096 100644 --- a/parquet/src/file/serialized_reader.rs +++ b/parquet/src/file/serialized_reader.rs @@ -553,12 +553,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`]. @@ -771,23 +771,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, @@ -801,7 +808,7 @@ impl SerializedPageReaderContext { } } - fn decrypt_page_data( + pub(crate) fn decrypt_page_data( &self, buffer: T, _page_index: usize, @@ -813,7 +820,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, @@ -851,7 +858,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>, @@ -897,7 +909,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, @@ -991,7 +1003,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,