This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11157-29ec1c09fa6c8454848b596f95eafcf36662fc87 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit efb652c27242870f35e009df26d7d84a0c1ee06a Author: Ed Seidl <[email protected]> AuthorDate: Wed Oct 7 04:24:24 2026 +0000 feat: Allow for partial population of the Parquet Page Index (#11157) # Which issue does this PR close? - Part of #7582. # Rationale for this change The Parquet Page Index can be quite large, but in many cases (column projection, row group pruning) only a small subset of the structure is required. # What changes are included in this PR? This PR introduces a new struct `ColumnChunkMask` which can be used to identify the subset of the column and offset indexes that need to be retrieved and parsed. This also modifies the `ParquetMetaDataPushDecoder` to request ranges based on the offset and column index masks to reduce unnecessary I/O. Clarifies and standardizes repeated page index reads: each successful call now replaces the existing page index state rather than implicitly preserving previously loaded indexes. # Are these changes tested? Yes # Are there any user-facing changes? There is a behavior change to repeated page index reads as mentioned above, but the prior behavior was not documented. Adds some new public APIs, but no breaking changes to existing public APIs # AI assistance Claude Code and Codex were used in the development of this PR, but I take responsibility for the end result. --------- Co-authored-by: Adrian Garcia Badaracco <[email protected]> --- parquet/benches/metadata.rs | 28 +- parquet/src/arrow/arrow_reader/mod.rs | 83 +++- parquet/src/arrow/arrow_reader/statistics.rs | 9 + parquet/src/arrow/async_reader/store.rs | 10 + parquet/src/file/metadata/mod.rs | 2 +- parquet/src/file/metadata/parser.rs | 96 ++--- parquet/src/file/metadata/push_decoder.rs | 109 +++++- parquet/src/file/metadata/reader.rs | 292 +++++++++++++- parquet/src/file/page_index/index_reader.rs | 11 - .../arrow_reader/custom_page_index_provider.rs | 2 +- parquet/tests/arrow_reader/io/mod.rs | 4 +- parquet/tests/arrow_reader/mod.rs | 1 + parquet/tests/arrow_reader/page_index.rs | 434 +++++++++++++++++++++ parquet/tests/arrow_reader/partial_offset_index.rs | 67 +++- 14 files changed, 1058 insertions(+), 90 deletions(-) diff --git a/parquet/benches/metadata.rs b/parquet/benches/metadata.rs index 0ec1b3cd17..b5ee07c495 100644 --- a/parquet/benches/metadata.rs +++ b/parquet/benches/metadata.rs @@ -21,9 +21,9 @@ use std::sync::Arc; use parquet::basic::{Encoding, PageType, Type as PhysicalType}; use parquet::file::metadata::{ - ColumnChunkMetaData, FileMetaData, LevelHistogram, PageEncodingStats, ParquetMetaData, - ParquetMetaDataOptions, ParquetMetaDataReader, ParquetMetaDataWriter, ParquetStatisticsPolicy, - RowGroupMetaData, + ColumnChunkMask, ColumnChunkMetaData, FileMetaData, LevelHistogram, PageEncodingStats, + PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions, ParquetMetaDataReader, + ParquetMetaDataWriter, ParquetStatisticsPolicy, RowGroupMetaData, }; use parquet::file::statistics::Statistics; use parquet::file::writer::TrackedWrite; @@ -191,6 +191,28 @@ fn criterion_benchmark(c: &mut Criterion) { }) }); + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Skip) + .parse_and_finish(&data) + .unwrap(); + let mut reader = ParquetMetaDataReader::new_with_metadata(metadata.clone()) + .with_page_index_policy(PageIndexPolicy::Required); + c.bench_function("read page index", |b| { + b.iter(|| { + reader.read_page_indexes(&data).unwrap(); + }) + }); + + let mut reader = ParquetMetaDataReader::new_with_metadata(metadata) + .with_page_index_policy(PageIndexPolicy::Required) + .with_column_index_mask(ColumnChunkMask::columns([0])) + .with_offset_index_mask(ColumnChunkMask::columns([0, 1, 4])); + c.bench_function("read page index reduced columns", |b| { + b.iter(|| { + reader.read_page_indexes(&data).unwrap(); + }) + }); + let meta_data = get_footer_bytes(data.clone()); let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false); c.bench_function("decode parquet metadata", |b| { diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs index 73bcb2048a..132236c7cd 100644 --- a/parquet/src/arrow/arrow_reader/mod.rs +++ b/parquet/src/arrow/arrow_reader/mod.rs @@ -45,8 +45,8 @@ use crate::column::page::{PageIterator, PageReader}; use crate::encryption::decrypt::FileDecryptionProperties; use crate::errors::{ParquetError, Result}; use crate::file::metadata::{ - PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions, ParquetMetaDataReader, - ParquetStatisticsPolicy, RowGroupMetaData, + ColumnChunkMask, PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions, + ParquetMetaDataReader, ParquetStatisticsPolicy, RowGroupMetaData, }; use crate::file::reader::{ChunkReader, SerializedPageReader}; use crate::schema::types::SchemaDescriptor; @@ -593,8 +593,10 @@ pub struct ArrowReaderOptions { /// [ARROW_SCHEMA_META_KEY]: crate::arrow::ARROW_SCHEMA_META_KEY supplied_schema: Option<SchemaRef>, - pub(crate) column_index: PageIndexPolicy, - pub(crate) offset_index: PageIndexPolicy, + column_index: PageIndexPolicy, + offset_index: PageIndexPolicy, + column_index_mask: ColumnChunkMask, + offset_index_mask: ColumnChunkMask, /// Options to control reading of Parquet metadata metadata_options: ParquetMetaDataOptions, @@ -774,6 +776,49 @@ impl ArrowReaderOptions { self } + // TODO: deprecate the following as part of https://github.com/apache/arrow-rs/issues/11395 + + /// Sets the same [`ColumnChunkMask`] for both page-index structures. + pub fn with_page_index_mask(self, mask: ColumnChunkMask) -> Self { + self.with_column_index_mask(mask.clone()) + .with_offset_index_mask(mask) + } + + /// Sets the [`ColumnChunkMask`] for the Parquet [ColumnIndex] structure. + /// + /// The column index can be costly to decode and store, especially when it is needed + /// only for a subset of row groups or columns (such as when filtering by a predicate + /// on a single column). Providing a [`ColumnChunkMask`] can greatly decrease + /// the time needed to decode this metadata. + /// + /// The mask applies only if the column-index policy is not [`PageIndexPolicy::Skip`] + /// (the default), or an underlying reader is configured to preload the index. It is + /// honored by loading APIs such as [`ArrowReaderMetadata::load`]; + /// [`ArrowReaderMetadata::try_new`] does not load or filter page indexes. + /// + /// [ColumnIndex]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + pub fn with_column_index_mask(mut self, mask: ColumnChunkMask) -> Self { + self.column_index_mask = mask; + self + } + + /// Sets the [`ColumnChunkMask`] for the Parquet [OffsetIndex] structure. + /// + /// The offset index can be costly to decode and store, especially when it is needed + /// only for a subset of row groups or columns (such as when projecting a small subset + /// of columns). Providing a [`ColumnChunkMask`] can greatly decrease + /// the time needed to decode this metadata. + /// + /// Page pruning also needs the offset index for predicate columns, so callers should + /// include those columns in addition to projected columns. The same loading and policy + /// qualifications as [`Self::with_column_index_mask`] apply. + /// + /// [OffsetIndex]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + pub fn with_offset_index_mask(mut self, mask: ColumnChunkMask) -> Self { + self.offset_index_mask = mask; + self + } + /// Provide a Parquet schema to use when decoding the metadata. The schema in the Parquet /// footer will be skipped. /// @@ -924,6 +969,22 @@ impl ArrowReaderOptions { self.column_index } + // TODO: deprecate the following as part of https://github.com/apache/arrow-rs/issues/11395 + + /// Retrieve the currently set [`ColumnChunkMask`] for the offset index. + /// + /// This can be set via [`with_offset_index_mask`][Self::with_offset_index_mask]. + pub fn offset_index_mask(&self) -> &ColumnChunkMask { + &self.offset_index_mask + } + + /// Retrieve the currently set [`ColumnChunkMask`] for the column index. + /// + /// This can be set via [`with_column_index_mask`][Self::with_column_index_mask]. + pub fn column_index_mask(&self) -> &ColumnChunkMask { + &self.column_index_mask + } + /// Retrieve the currently set metadata decoding options. pub fn metadata_options(&self) -> &ParquetMetaDataOptions { &self.metadata_options @@ -973,6 +1034,18 @@ impl ParquetMetaDataReader { .with_offset_index_policy(options.offset_index_policy()); } + // Preload settings on the underlying reader may enable an index even when the + // corresponding options policy is `Skip`, so apply non-default masks independently. + // + // TODO: https://github.com/apache/arrow-rs/issues/11395 may deprecate this API. Decide + // at that time if masks should be added to ParquetObjectReader. + if !options.column_index_mask().is_all() { + self = self.with_column_index_mask(options.column_index_mask().clone()); + } + if !options.offset_index_mask().is_all() { + self = self.with_offset_index_mask(options.offset_index_mask().clone()); + } + self } } @@ -1019,6 +1092,8 @@ impl ArrowReaderMetadata { let metadata = ParquetMetaDataReader::new() .with_column_index_policy(options.column_index_policy()) .with_offset_index_policy(options.offset_index_policy()) + .with_column_index_mask(options.column_index_mask().clone()) + .with_offset_index_mask(options.offset_index_mask().clone()) .with_metadata_options(Some(options.metadata_options.clone())); #[cfg(feature = "encryption")] let metadata = metadata.with_decryption_properties( diff --git a/parquet/src/arrow/arrow_reader/statistics.rs b/parquet/src/arrow/arrow_reader/statistics.rs index 486e937a79..32684b2b12 100644 --- a/parquet/src/arrow/arrow_reader/statistics.rs +++ b/parquet/src/arrow/arrow_reader/statistics.rs @@ -2042,6 +2042,15 @@ impl<'a> StatisticsConverter<'a> { let mut nulls = NullBufferBuilder::new(0); for rg_idx in row_group_indices { let Some(offset_index) = page_index.offset_index(*rg_idx, parquet_index) else { + // A partial page index may contain the ColumnIndex but not the OffsetIndex. + // Preserve page-wise alignment with data_page_{mins,maxes,null_counts} by + // emitting one null row count for every page known from the ColumnIndex. + let num_pages = page_index + .column_index(*rg_idx, parquet_index) + .map(|index| index.num_pages() as usize) + .unwrap_or(0); + row_counts.resize(row_counts.len() + num_pages, 0); + nulls.append_n_nulls(num_pages); continue; }; let page_locations = offset_index.page_locations(); diff --git a/parquet/src/arrow/async_reader/store.rs b/parquet/src/arrow/async_reader/store.rs index 4bd65cfead..a38d0c3558 100644 --- a/parquet/src/arrow/async_reader/store.rs +++ b/parquet/src/arrow/async_reader/store.rs @@ -244,6 +244,16 @@ impl AsyncFileReader for ParquetObjectReader { .with_column_index_policy(options.column_index_policy()) .with_offset_index_policy(options.offset_index_policy()); } + // TODO: https://github.com/apache/arrow-rs/issues/11395 may deprecate this API. Decide + // at that time if masks should be added to ParquetObjectReader. + if let Some(options) = options { + if !options.column_index_mask().is_all() { + metadata = metadata.with_column_index_mask(options.column_index_mask().clone()); + } + if !options.offset_index_mask().is_all() { + metadata = metadata.with_offset_index_mask(options.offset_index_mask().clone()); + } + } let metadata = if let Some(file_size) = self.file_size { metadata.load_and_finish(self, file_size).await? diff --git a/parquet/src/file/metadata/mod.rs b/parquet/src/file/metadata/mod.rs index a12fbdd1de..f5ee3b20a2 100644 --- a/parquet/src/file/metadata/mod.rs +++ b/parquet/src/file/metadata/mod.rs @@ -95,7 +95,7 @@ use crate::{ pub use footer_tail::FooterTail; pub use options::{ParquetMetaDataOptions, ParquetStatisticsPolicy}; pub use push_decoder::ParquetMetaDataPushDecoder; -pub use reader::{PageIndexPolicy, ParquetMetaDataReader}; +pub use reader::{ColumnChunkMask, PageIndexPolicy, ParquetMetaDataReader}; use std::io::Write; use std::ops::Range; use std::sync::Arc; diff --git a/parquet/src/file/metadata/parser.rs b/parquet/src/file/metadata/parser.rs index 861bb2e53b..22f0564a9c 100644 --- a/parquet/src/file/metadata/parser.rs +++ b/parquet/src/file/metadata/parser.rs @@ -20,19 +20,20 @@ //! These functions parse thrift-encoded metadata from a byte slice //! into the corresponding Rust structures -use std::sync::Arc; +use std::{ops::Range, sync::Arc}; -use crate::errors::ParquetError; +use bytes::Bytes; + +use crate::errors::{ParquetError, Result}; use crate::file::metadata::page_index::{PageIndexBuilder, PageIndexProvider}; use crate::file::metadata::thrift::parquet_metadata_from_bytes; use crate::file::metadata::{ - ColumnChunkMetaData, PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions, + ColumnChunkMask, ColumnChunkMetaData, PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions, }; use crate::file::page_index::column_index::ColumnIndexMetaData; use crate::file::page_index::index_reader::{decode_column_index, decode_offset_index}; use crate::file::page_index::offset_index::OffsetIndexMetaData; -use bytes::Bytes; /// Helper struct for metadata parsing /// @@ -49,7 +50,6 @@ mod inner { use super::*; use crate::encryption::decrypt::FileDecryptionProperties; - use crate::errors::Result; /// API for decoding metadata that may be encrypted #[derive(Debug, Default)] @@ -107,7 +107,7 @@ mod inner { column: &ColumnChunkMetaData, row_group_index: usize, col_index: usize, - ) -> crate::errors::Result<ColumnIndexMetaData> { + ) -> Result<ColumnIndexMetaData> { use crate::encryption::decrypt::CryptoContext; match &column.column_crypto_metadata { Some(crypto_metadata) => { @@ -135,7 +135,7 @@ mod inner { column: &ColumnChunkMetaData, row_group_index: usize, col_index: usize, - ) -> crate::errors::Result<OffsetIndexMetaData> { + ) -> Result<OffsetIndexMetaData> { use crate::encryption::decrypt::CryptoContext; match &column.column_crypto_metadata { Some(crypto_metadata) => { @@ -161,7 +161,6 @@ mod inner { #[cfg(not(feature = "encryption"))] mod inner { use super::*; - use crate::errors::Result; use std::sync::Arc; /// parallel implementation when encryption feature is not enabled /// @@ -207,7 +206,7 @@ mod inner { column: &ColumnChunkMetaData, _row_group_index: usize, _col_index: usize, - ) -> crate::errors::Result<ColumnIndexMetaData> { + ) -> Result<ColumnIndexMetaData> { decode_column_index(bytes, column.column_type()) } @@ -217,7 +216,7 @@ mod inner { _column: &ColumnChunkMetaData, _row_group_index: usize, _col_index: usize, - ) -> crate::errors::Result<OffsetIndexMetaData> { + ) -> Result<OffsetIndexMetaData> { decode_offset_index(bytes) } } @@ -232,39 +231,43 @@ mod inner { pub(crate) fn decode_metadata( buf: &[u8], options: Option<&ParquetMetaDataOptions>, -) -> crate::errors::Result<ParquetMetaData> { +) -> Result<ParquetMetaData> { parquet_metadata_from_bytes(buf, options) } -/// Parses page index from the provided bytes and adds it to the metadata. +/// Parses page indexes from the provided bytes and replaces those in the metadata. /// /// Arguments -/// * `metadata` - The ParquetMetaData to which the parsed column index will be added. +/// * `metadata` - The ParquetMetaData whose page index will be replaced. /// * `column_index_policy` - The policy for handling column index parsing (e.g., /// Required, Optional, Skip). /// * `offset_index_policy` - The policy for handling offset index parsing (e.g., /// Required, Optional, Skip). +/// * `column_index_mask` - The row groups and leaf columns whose column indexes are parsed. +/// Ignored when `column_index_policy` is [`PageIndexPolicy::Skip`]. +/// * `offset_index_mask` - The row groups and leaf columns whose offset indexes are parsed. +/// Ignored when `offset_index_policy` is [`PageIndexPolicy::Skip`]. /// * `bytes` - The byte slice containing the page index data. /// * `start_offset` - The offset where `bytes` begin in the file. pub(crate) fn parse_page_index( metadata: &mut ParquetMetaData, column_index_policy: PageIndexPolicy, offset_index_policy: PageIndexPolicy, + column_index_mask: &ColumnChunkMask, + offset_index_mask: &ColumnChunkMask, bytes: &Bytes, start_offset: u64, -) -> crate::errors::Result<()> { - if column_index_policy == PageIndexPolicy::Skip && offset_index_policy == PageIndexPolicy::Skip - { - return Ok(()); - } +) -> Result<()> { let num_row_groups = metadata.num_row_groups(); let num_columns = metadata.file_metadata().schema_descr().num_columns(); let mut builder = PageIndexBuilder::default(); + if column_index_policy != PageIndexPolicy::Skip { builder.allocate_column_indexes(num_row_groups, num_columns); parse_column_index( metadata, column_index_policy, + column_index_mask, &mut builder, bytes, start_offset, @@ -275,45 +278,45 @@ pub(crate) fn parse_page_index( parse_offset_index( metadata, offset_index_policy, + offset_index_mask, &mut builder, bytes, start_offset, )?; } + // Always replace the page index, even if both policies are Skip. + // This ensures that repeated calls to read_page_indexes replace the existing index + // rather than preserving it. let page_index = builder.build(); - // if both indexes are missing from the file, return without modifying `metadata` - if !page_index.has_column_indexes() && !page_index.has_offset_indexes() { - return Ok(()); + if page_index.has_column_indexes() || page_index.has_offset_indexes() { + metadata.set_page_index(Some(Arc::new(page_index))); + } else { + // If no indexes were read, clear the page index + metadata.set_page_index(None); } - metadata.set_page_index(Some(Arc::new(page_index))); Ok(()) } fn parse_column_index( metadata: &ParquetMetaData, column_index_policy: PageIndexPolicy, + mask: &ColumnChunkMask, page_index_builder: &mut PageIndexBuilder, bytes: &Bytes, start_offset: u64, -) -> crate::errors::Result<()> { +) -> Result<()> { if column_index_policy == PageIndexPolicy::Skip { return Ok(()); } - for rg_idx in 0..metadata.num_row_groups() { + for rg_idx in mask.row_group_indices(metadata.num_row_groups()) { let rg = metadata.row_group(rg_idx); - for col_idx in 0..rg.num_columns() { + for col_idx in mask.column_indices(rg.num_columns()) { let col = rg.column(col_idx); if let Some(r) = col.column_index_range() { - let r_start = usize::try_from(r.start - start_offset)?; - let r_end = usize::try_from(r.end - start_offset)?; - let idx = inner::parse_single_column_index( - &bytes[r_start..r_end], - metadata, - col, - rg_idx, - col_idx, - )?; + let idx_bytes = get_index_bytes(bytes, start_offset, r)?; + let idx = + inner::parse_single_column_index(idx_bytes, metadata, col, rg_idx, col_idx)?; page_index_builder.put_column_index(idx, rg_idx, col_idx); } } @@ -325,27 +328,22 @@ fn parse_column_index( fn parse_offset_index( metadata: &ParquetMetaData, offset_index_policy: PageIndexPolicy, + mask: &ColumnChunkMask, page_index_builder: &mut PageIndexBuilder, bytes: &Bytes, start_offset: u64, -) -> crate::errors::Result<()> { +) -> Result<()> { if offset_index_policy == PageIndexPolicy::Skip { return Ok(()); } - for rg_idx in 0..metadata.num_row_groups() { + for rg_idx in mask.row_group_indices(metadata.num_row_groups()) { let rg = metadata.row_group(rg_idx); - for col_idx in 0..rg.num_columns() { + for col_idx in mask.column_indices(rg.num_columns()) { let col = rg.column(col_idx); if let Some(r) = col.offset_index_range() { - let r_start = usize::try_from(r.start - start_offset)?; - let r_end = usize::try_from(r.end - start_offset)?; - let idx = inner::parse_single_offset_index( - &bytes[r_start..r_end], - metadata, - col, - rg_idx, - col_idx, - )?; + let idx_bytes = get_index_bytes(bytes, start_offset, r)?; + let idx = + inner::parse_single_offset_index(idx_bytes, metadata, col, rg_idx, col_idx)?; page_index_builder.put_offset_index(idx, rg_idx, col_idx); } else if offset_index_policy == PageIndexPolicy::Required { return Err(general_err!("missing offset index")); @@ -355,3 +353,9 @@ fn parse_offset_index( Ok(()) } + +fn get_index_bytes(bytes: &[u8], start_offset: u64, range: Range<u64>) -> Result<&[u8]> { + let start = usize::try_from(range.start - start_offset)?; + let end = usize::try_from(range.end - start_offset)?; + Ok(&bytes[start..end]) +} diff --git a/parquet/src/file/metadata/push_decoder.rs b/parquet/src/file/metadata/push_decoder.rs index c5e77cc224..90fd3ee7fa 100644 --- a/parquet/src/file/metadata/push_decoder.rs +++ b/parquet/src/file/metadata/push_decoder.rs @@ -21,8 +21,9 @@ use crate::encryption::decrypt::FileDecryptionProperties; use crate::errors::{ParquetError, Result}; use crate::file::FOOTER_SIZE; use crate::file::metadata::parser::{MetadataParser, parse_page_index}; -use crate::file::metadata::{FooterTail, PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions}; -use crate::file::page_index::index_reader::acc_range; +use crate::file::metadata::{ + ColumnChunkMask, FooterTail, PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions, +}; use crate::file::reader::ChunkReader; use bytes::Bytes; use std::ops::Range; @@ -226,6 +227,10 @@ pub struct ParquetMetaDataPushDecoder { column_index_policy: PageIndexPolicy, /// policy for loading OffsetIndex (part of the PageIndex) offset_index_policy: PageIndexPolicy, + /// which rows and columns of the column index should be read + column_index_mask: ColumnChunkMask, + /// which rows and columns of the offset index should be read + offset_index_mask: ColumnChunkMask, /// Underlying buffers buffers: crate::util::push_buffers::PushBuffers, /// Encryption API @@ -250,6 +255,8 @@ impl ParquetMetaDataPushDecoder { state: DecodeState::ReadingFooter, column_index_policy: PageIndexPolicy::Optional, offset_index_policy: PageIndexPolicy::Optional, + column_index_mask: ColumnChunkMask::all(), + offset_index_mask: ColumnChunkMask::all(), buffers: crate::util::push_buffers::PushBuffers::new(file_len), metadata_parser: MetadataParser::new(), }) @@ -266,6 +273,11 @@ impl ParquetMetaDataPushDecoder { /// /// This can be used to parse and populate the page index structures /// after the metadata has already been decoded. + /// + /// When decoding page indexes, any existing page indexes in `metadata` will be + /// replaced based on the configured policies. See + /// [`crate::file::metadata::ParquetMetaDataReader::read_page_indexes`] for details + /// on the replacement behavior. pub fn try_new_with_metadata(file_len: u64, metadata: ParquetMetaData) -> Result<Self> { let mut new_self = Self::try_new(file_len)?; new_self.state = DecodeState::ReadingPageIndex(Box::new(metadata)); @@ -300,6 +312,27 @@ impl ParquetMetaDataPushDecoder { self } + // TODO: deprecate the following as part of https://github.com/apache/arrow-rs/issues/11395 + + /// Select the row groups and columns for which both page index structures are read. + pub fn with_page_index_mask(mut self, mask: ColumnChunkMask) -> Self { + self.column_index_mask = mask.clone(); + self.offset_index_mask = mask; + self + } + + /// Select the row groups and columns for which column indexes are read. + pub fn with_column_index_mask(mut self, mask: ColumnChunkMask) -> Self { + self.column_index_mask = mask; + self + } + + /// Select the row groups and columns for which offset indexes are read. + pub fn with_offset_index_mask(mut self, mask: ColumnChunkMask) -> Self { + self.offset_index_mask = mask; + self + } + /// Set the options to use when decoding the Parquet metadata. pub fn with_metadata_options(mut self, options: Option<Arc<ParquetMetaDataOptions>>) -> Self { self.metadata_parser = self.metadata_parser.with_metadata_options(options); @@ -412,9 +445,16 @@ impl ParquetMetaDataPushDecoder { &metadata, self.column_index_policy, self.offset_index_policy, + &self.column_index_mask, + &self.offset_index_mask, ); let Some(page_index_range) = range else { + // There is nothing to decode because both policies are Skip, the masks + // select nothing, or the file has no index ranges. Clear any previously + // loaded indexes while preserving the compatibility behavior that + // PageIndexPolicy::Required accepts files with no page indexes at all. + metadata.set_page_index(None); self.state = DecodeState::Finished; return Ok(DecodeResult::Data(*metadata)); }; @@ -430,6 +470,8 @@ impl ParquetMetaDataPushDecoder { &mut metadata, self.column_index_policy, self.offset_index_policy, + &self.column_index_mask, + &self.offset_index_mask, &buffer, offset, )?; @@ -482,21 +524,39 @@ enum DecodeState { } /// Returns the byte range needed to read the offset/page indexes, based on the -/// specified policies +/// specified policies and masks /// /// Returns None if no page indexes are needed -pub fn range_for_page_index( +pub(crate) fn range_for_page_index( metadata: &ParquetMetaData, column_index_policy: PageIndexPolicy, offset_index_policy: PageIndexPolicy, + column_index_mask: &ColumnChunkMask, + offset_index_mask: &ColumnChunkMask, ) -> Option<Range<u64>> { - let mut range = None; - for c in metadata.row_groups().iter().flat_map(|r| r.columns()) { - if column_index_policy != PageIndexPolicy::Skip { - range = acc_range(range, c.column_index_range()); + let mut range: Option<Range<u64>> = None; + for (policy, mask, column_index) in [ + (column_index_policy, column_index_mask, true), + (offset_index_policy, offset_index_mask, false), + ] { + if policy == PageIndexPolicy::Skip { + continue; } - if offset_index_policy != PageIndexPolicy::Skip { - range = acc_range(range, c.offset_index_range()); + for row_group in mask.row_group_indices(metadata.num_row_groups()) { + for column in mask.column_indices(metadata.row_group(row_group).num_columns()) { + let column = metadata.row_group(row_group).column(column); + let index = if column_index { + column.column_index_range() + } else { + column.offset_index_range() + }; + if let Some(index) = index { + range = Some(match range { + Some(range) => range.start.min(index.start)..range.end.max(index.end), + None => index, + }); + } + } } } range @@ -532,6 +592,33 @@ mod tests { assert!(metadata.page_index().is_some_and(|idx| idx.is_complete())); } + #[test] + fn test_range_for_page_index_respects_column_mask() { + let mut decoder = ParquetMetaDataPushDecoder::try_new(test_file_len()).unwrap(); + push_ranges_to_metadata_decoder(&mut decoder, vec![test_file_range()]); + let metadata = expect_data(decoder.try_decode()); + + let all = range_for_page_index( + &metadata, + PageIndexPolicy::Required, + PageIndexPolicy::Skip, + &ColumnChunkMask::all(), + &ColumnChunkMask::all(), + ) + .unwrap(); + let masked = range_for_page_index( + &metadata, + PageIndexPolicy::Required, + PageIndexPolicy::Skip, + &ColumnChunkMask::columns([0]), + &ColumnChunkMask::all(), + ) + .unwrap(); + + assert!(all.start <= masked.start && masked.end <= all.end); + assert!(masked.end - masked.start < all.end - all.start); + } + /// It is possible to feed some, but not all, of the footer into the metadata decoder /// before asked. This avoids multiple IO requests #[test] @@ -612,6 +699,8 @@ mod tests { // expect the second request to read the offset indexes let ranges = expect_needs_data(metadata_decoder.try_decode()); + // since we're reading all indexes, there should be a single range needed + assert_eq!(ranges.len(), 1); push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges); // expect the third request to read the actual data diff --git a/parquet/src/file/metadata/reader.rs b/parquet/src/file/metadata/reader.rs index 24bc5c982b..78710600bf 100644 --- a/parquet/src/file/metadata/reader.rs +++ b/parquet/src/file/metadata/reader.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +#[cfg(feature = "arrow")] +use crate::arrow::ProjectionMask; #[cfg(feature = "encryption")] use crate::encryption::decrypt::FileDecryptionProperties; use crate::errors::{ParquetError, Result}; @@ -31,6 +33,9 @@ use crate::schema::types::SchemaDescriptor; #[cfg(feature = "arrow")] use arrow_array::ArrayRef; use bytes::Bytes; +use std::collections::BTreeSet; +use std::iter::FusedIterator; +use std::slice::Iter; use std::sync::Arc; use std::{io::Read, ops::Range}; @@ -75,6 +80,8 @@ pub struct ParquetMetaDataReader { metadata: Option<ParquetMetaData>, column_index: PageIndexPolicy, offset_index: PageIndexPolicy, + column_index_mask: ColumnChunkMask, + offset_index_mask: ColumnChunkMask, prefetch_hint: Option<usize>, metadata_options: Option<Arc<ParquetMetaDataOptions>>, // Size of the serialized thrift metadata plus the 8 byte footer. Only set if @@ -84,6 +91,10 @@ pub struct ParquetMetaDataReader { file_decryption_properties: Option<Arc<FileDecryptionProperties>>, } +// TODO: https://github.com/apache/arrow-rs/issues/11395 is tracking unifying PageIndexPolicy +// with the ColumnChunkMask to reduce the dizzying number of APIs that exist to control reading +// the page indexes. + /// Describes the policy for reading page indexes #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum PageIndexPolicy { @@ -92,10 +103,12 @@ pub enum PageIndexPolicy { Skip, /// Read the page index if it exists, otherwise do not error. Optional, - /// Require the page index to exist, and error if it does not. + /// Require offset indexes for every selected column chunk when any selected + /// page index ranges exist. /// - /// Only enforced for the offset index; this is the same as [`Self::Optional`] - /// for the column index. + /// For compatibility, a file with no page index ranges is accepted. This is only + /// enforced for the offset index; this is the same as [`Self::Optional`] for the + /// column index. Required, } @@ -108,6 +121,203 @@ impl From<bool> for PageIndexPolicy { } } +/// A cheaply cloneable struct to specify column chunks for which metadata is required. +/// +/// Column chunks are identified by row group index and leaf column index (the index of the +/// column in [`SchemaDescriptor::columns`], not the index of a root or Arrow field). This struct +/// allows for specifying vertical slices of column chunk data (via [`Self::columns`]), +/// horizontal slices (via [`Self::row_groups`]), or the intersection of the two +/// (via [`Self::row_groups_and_columns`]). +/// +/// Because Parquet collection cannot contain more than [`i32::MAX`] values, user provided indices +/// outside that range will be discarded. +/// +/// At present this is only used to select elements of the [Page Index] for decoding. +/// +/// # Examples +/// +/// To select columns 0 and 1 from all row groups: +/// ```rust +/// # use parquet::file::metadata::ColumnChunkMask; +/// let mask = ColumnChunkMask::columns([0, 1]); +/// ``` +/// +/// To select all columns from row group 2: +/// ```rust +/// # use parquet::file::metadata::ColumnChunkMask; +/// let mask = ColumnChunkMask::row_groups([2]); +/// ``` +/// +/// To select columns 1 and 3 from row group 0: +/// ```rust +/// # use parquet::file::metadata::ColumnChunkMask; +/// let mask = ColumnChunkMask::row_groups_and_columns([0], [1, 3]); +/// ``` +/// +/// [Page Index]: https://parquet.apache.org/docs/file-format/pageindex/ +#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)] +pub struct ColumnChunkMask { + // `row_groups` and `columns` are optional `u32` slices, with elements + // sorted in ascending order (to permit binary searching for member indices). + // `None` means all indices are included, while an empty slice means none are. + // Indices are stored as `u32` because Parquet/Thrift collections cannot contain + // more than i32::MAX entries. + row_groups: Option<Arc<[u32]>>, + columns: Option<Arc<[u32]>>, +} + +impl ColumnChunkMask { + /// Select all row groups and columns. + pub fn all() -> Self { + Self::default() + } + + /// Select no row groups or columns. + pub fn none() -> Self { + Self { + row_groups: Some(Arc::from([])), + columns: Some(Arc::from([])), + } + } + + /// Select only the listed columns. + /// + /// Passing an empty iterator selects no columns. + pub fn columns(columns: impl IntoIterator<Item = usize>) -> Self { + Self { + row_groups: None, + columns: Self::iter_to_set(columns), + } + } + + /// Select only the listed row groups. + /// + /// Passing an empty iterator selects no row groups. + pub fn row_groups(row_groups: impl IntoIterator<Item = usize>) -> Self { + Self { + row_groups: Self::iter_to_set(row_groups), + columns: None, + } + } + + /// Select only the listed row groups and columns. + /// + /// An empty iterator for either dimension selects no column chunks. + pub fn row_groups_and_columns( + row_groups: impl IntoIterator<Item = usize>, + columns: impl IntoIterator<Item = usize>, + ) -> Self { + Self { + row_groups: Self::iter_to_set(row_groups), + columns: Self::iter_to_set(columns), + } + } + + /// Returns whether row group `idx` is selected and within the bounds `[0, num_row_groups)`. + pub fn includes_row_group(&self, idx: usize, num_row_groups: usize) -> bool { + idx < num_row_groups && Self::includes_index(self.row_groups.as_ref(), idx) + } + + /// Returns whether column `idx` is selected and within the bounds `[0, num_columns)`. + pub fn includes_column(&self, idx: usize, num_columns: usize) -> bool { + idx < num_columns && Self::includes_index(self.columns.as_ref(), idx) + } + + fn includes_index(keep: Option<&Arc<[u32]>>, idx: usize) -> bool { + // return false for out-of-bounds index + let Ok(idx) = i32::try_from(idx) else { + return false; + }; + let idx = idx as u32; + keep.is_none_or(|keep| keep.binary_search(&idx).is_ok()) + } + + /// Returns `true` when this mask selects every column chunk. + pub fn is_all(&self) -> bool { + self.row_groups.is_none() && self.columns.is_none() + } + + /// Returns `true` when either configured axis is empty (and thus would select no column chunks). + /// + /// This does not consider the dimensions of a particular Parquet file, so may return `false` + /// when [`Self::column_indices`] or [`Self::row_group_indices`] returns an empty iterator due + /// to all elements exceeding the provided upper bound. + pub fn is_none(&self) -> bool { + self.row_groups.as_ref().is_some_and(|i| i.is_empty()) + || self.columns.as_ref().is_some_and(|i| i.is_empty()) + } + + /// Creates a mask selecting the leaf columns in an Arrow projection. + #[cfg(feature = "arrow")] + pub fn from_projection(projection: &ProjectionMask, schema: &SchemaDescriptor) -> Self { + Self::columns((0..schema.num_columns()).filter(|&i| projection.leaf_included(i))) + } + + /// Returns an iterator over the row group indices selected by this mask. If the mask is `all` + /// then all indices in the range `[0, num_row_groups)` are returned. + pub fn row_group_indices(&self, num_row_groups: usize) -> impl Iterator<Item = usize> + '_ { + Self::axis_indices(self.row_groups.as_deref(), num_row_groups) + } + + /// Returns an iterator over the column indices selected by this mask. If the mask is `all` + /// then all indices in the range `[0, num_columns)` are returned. + pub fn column_indices(&self, num_columns: usize) -> impl Iterator<Item = usize> + '_ { + Self::axis_indices(self.columns.as_deref(), num_columns) + } + + fn axis_indices(axis: Option<&[u32]>, len: usize) -> AxisIndices<'_> { + match axis { + None => AxisIndices::All(0..len), + Some(indices) => { + let end = indices.partition_point(|&i| (i as usize) < len); + AxisIndices::Selected(indices[..end].iter()) + } + } + } + + // Used during initialization, this helper takes an iterator over `usize` elements, + // discards out-of-range values and then sorts. + fn iter_to_set(indices: impl IntoIterator<Item = usize>) -> Option<Arc<[u32]>> { + Some( + indices + .into_iter() + .filter_map(|i| i32::try_from(i).ok()) // remove out-of-bounds values + .map(|i| i as u32) // store as u32. safe because 0 <= i <= i32::MAX + .collect::<BTreeSet<_>>() + .into_iter() + .collect::<Vec<_>>() + .into(), + ) + } +} + +// helper for ColumnChunkMask selected index traversal +enum AxisIndices<'a> { + All(Range<usize>), + Selected(Iter<'a, u32>), +} + +impl Iterator for AxisIndices<'_> { + type Item = usize; + + fn next(&mut self) -> Option<Self::Item> { + match self { + Self::All(indices) => indices.next(), + Self::Selected(indices) => indices.next().map(|&i| i as usize), + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + match self { + Self::All(indices) => indices.size_hint(), + Self::Selected(indices) => indices.size_hint(), + } + } +} + +impl ExactSizeIterator for AxisIndices<'_> {} +impl FusedIterator for AxisIndices<'_> {} + impl ParquetMetaDataReader { /// Create a new [`ParquetMetaDataReader`] pub fn new() -> Self { @@ -141,6 +351,36 @@ impl ParquetMetaDataReader { self } + // TODO: deprecate the following as part of https://github.com/apache/arrow-rs/issues/11395 + + /// Selects the row groups and leaf columns for which both page index structures are read. + /// + /// The mask applies only to index structures whose policy is not [`PageIndexPolicy::Skip`]. + /// This reader fetches one byte range covering the selected indexes, so a sparse mask + /// primarily reduces decoding and memory rather than necessarily reducing I/O. + pub fn with_page_index_mask(self, mask: ColumnChunkMask) -> Self { + self.with_column_index_mask(mask.clone()) + .with_offset_index_mask(mask) + } + + /// Selects the row groups and leaf columns for which column indexes are read. + /// + /// The mask applies only when the column index policy is not [`PageIndexPolicy::Skip`]. + /// [`PageIndexPolicy::Required`] validates only selected column chunks. + pub fn with_column_index_mask(mut self, mask: ColumnChunkMask) -> Self { + self.column_index_mask = mask; + self + } + + /// Selects the row groups and leaf columns for which offset indexes are read. + /// + /// The mask applies only when the offset index policy is not [`PageIndexPolicy::Skip`]. + /// [`PageIndexPolicy::Required`] validates only selected column chunks. + pub fn with_offset_index_mask(mut self, mask: ColumnChunkMask) -> Self { + self.offset_index_mask = mask; + self + } + /// Sets the [`ParquetMetaDataOptions`] to use when decoding pub fn with_metadata_options(mut self, options: Option<ParquetMetaDataOptions>) -> Self { self.metadata_options = options.map(Arc::new); @@ -322,6 +562,13 @@ impl ParquetMetaDataReader { /// Read the page index structures when a [`ParquetMetaData`] has already been obtained. /// See [`Self::new_with_metadata()`] and [`Self::has_metadata()`]. + /// + /// # Behavior on multiple calls + /// + /// On success, this method replaces any existing page index in the metadata. Index types + /// configured as [`PageIndexPolicy::Skip`] are cleared; other index types contain only the + /// indexes selected by their configured masks and available according to their policies. If + /// neither index type is populated, the page index is cleared. pub fn read_page_indexes<R: ChunkReader>(&mut self, reader: &R) -> Result<()> { self.read_page_indexes_sized(reader, reader.len()) } @@ -331,6 +578,9 @@ impl ParquetMetaDataReader { /// a [`Bytes`] struct containing the tail of the file). /// See [`Self::new_with_metadata()`] and [`Self::has_metadata()`]. Like /// [`Self::try_parse_sized()`] this function may return [`ParquetError::NeedMoreData`]. + /// + /// See [`Self::read_page_indexes()`] for a description of the replacement behavior on + /// multiple calls. pub fn read_page_indexes_sized<R: ChunkReader>( &mut self, reader: &R, @@ -345,6 +595,8 @@ impl ParquetMetaDataReader { let push_decoder = ParquetMetaDataPushDecoder::try_new_with_metadata(file_size, metadata)? .with_offset_index_policy(self.offset_index) .with_column_index_policy(self.column_index) + .with_offset_index_mask(self.offset_index_mask.clone()) + .with_column_index_mask(self.column_index_mask.clone()) .with_metadata_options(self.metadata_options.clone()); let mut push_decoder = self.prepare_push_decoder(push_decoder); @@ -471,6 +723,9 @@ impl ParquetMetaDataReader { /// Asynchronously fetch the page index structures when a [`ParquetMetaData`] has already /// been obtained. See [`Self::new_with_metadata()`]. + /// + /// See [`Self::read_page_indexes()`] for a description of the replacement behavior on + /// multiple calls. #[cfg(all(feature = "async", feature = "arrow"))] pub async fn load_page_index<F: MetadataFetch>(&mut self, fetch: F) -> Result<()> { self.load_page_index_with_remainder(fetch, None).await @@ -539,6 +794,8 @@ impl ParquetMetaDataReader { let push_decoder = ParquetMetaDataPushDecoder::try_new_with_metadata(file_size, metadata)? .with_offset_index_policy(self.offset_index) .with_column_index_policy(self.column_index) + .with_offset_index_mask(self.offset_index_mask.clone()) + .with_column_index_mask(self.column_index_mask.clone()) .with_metadata_options(self.metadata_options.clone()); let mut push_decoder = self.prepare_push_decoder(push_decoder); @@ -1073,6 +1330,35 @@ mod tests { "EOF: Parquet file too small. Size is 1728 but need 1729" ); } + + #[test] + fn test_chunk_mask() { + let mask = ColumnChunkMask::row_groups_and_columns([0], [1]); + assert!(mask.includes_row_group(0, 1)); + assert!(!mask.includes_row_group(1, 1)); + assert!(!mask.includes_column(0, 2)); + assert!(mask.includes_column(1, 2)); + + let mask = ColumnChunkMask::columns([]); + assert!(!mask.includes_column(0, 3)); + assert_eq!(mask.column_indices(3).count(), 0); + + let all = ColumnChunkMask::all(); + assert!(all.is_all()); + assert!(all.includes_column(2, 3)); + assert!(!all.includes_column(3, 3)); + + assert_eq!(all.row_group_indices(3).collect::<Vec<_>>(), [0, 1, 2]); + assert_eq!(all.column_indices(2).collect::<Vec<_>>(), [0, 1]); + + let none = ColumnChunkMask::none(); + assert_eq!(none.row_group_indices(3).count(), 0); + assert_eq!(none.column_indices(3).count(), 0); + + let sparse = ColumnChunkMask::row_groups_and_columns([0, 2, 4], [1, 3, 5]); + assert_eq!(sparse.row_group_indices(3).collect::<Vec<_>>(), [0, 2]); + assert_eq!(sparse.column_indices(4).collect::<Vec<_>>(), [1, 3]); + } } #[cfg(all(feature = "async", feature = "arrow", test))] diff --git a/parquet/src/file/page_index/index_reader.rs b/parquet/src/file/page_index/index_reader.rs index 963ee0ea97..235f0b2418 100644 --- a/parquet/src/file/page_index/index_reader.rs +++ b/parquet/src/file/page_index/index_reader.rs @@ -30,17 +30,6 @@ use crate::parquet_thrift::{ }; use crate::thrift_struct; use std::io::Write; -use std::ops::Range; - -/// Computes the covering range of two optional ranges -/// -/// For example `acc_range(Some(7..9), Some(1..3)) = Some(1..9)` -pub(crate) fn acc_range(a: Option<Range<u64>>, b: Option<Range<u64>>) -> Option<Range<u64>> { - match (a, b) { - (Some(a), Some(b)) => Some(a.start.min(b.start)..a.end.max(b.end)), - (None, x) | (x, None) => x, - } -} /// Decode a Thrift [`OffsetIndex`] from the provided bytes. /// diff --git a/parquet/tests/arrow_reader/custom_page_index_provider.rs b/parquet/tests/arrow_reader/custom_page_index_provider.rs index 147c9a7336..0266377427 100644 --- a/parquet/tests/arrow_reader/custom_page_index_provider.rs +++ b/parquet/tests/arrow_reader/custom_page_index_provider.rs @@ -404,7 +404,7 @@ impl PageIndexProvider for SelectivePageIndexProvider { } /// Create a test parquet file with multiple row groups and multiple pages per column -fn create_test_file() -> Bytes { +pub(super) fn create_test_file() -> Bytes { // Create small pages and small row groups for testing let props = WriterProperties::builder() .set_statistics_enabled(EnabledStatistics::Page) diff --git a/parquet/tests/arrow_reader/io/mod.rs b/parquet/tests/arrow_reader/io/mod.rs index 3fb63ffdb8..bb9b5aeba4 100644 --- a/parquet/tests/arrow_reader/io/mod.rs +++ b/parquet/tests/arrow_reader/io/mod.rs @@ -132,7 +132,9 @@ impl AsyncFileReader for TestReader { if let Some(options) = options { metadata_reader = metadata_reader .with_column_index_policy(options.column_index_policy()) - .with_offset_index_policy(options.offset_index_policy()); + .with_offset_index_policy(options.offset_index_policy()) + .with_column_index_mask(options.column_index_mask().clone()) + .with_offset_index_mask(options.offset_index_mask().clone()); } self.metadata = Some(Arc::new( diff --git a/parquet/tests/arrow_reader/mod.rs b/parquet/tests/arrow_reader/mod.rs index 1dfcfb6901..680524f64e 100644 --- a/parquet/tests/arrow_reader/mod.rs +++ b/parquet/tests/arrow_reader/mod.rs @@ -47,6 +47,7 @@ mod int96_stats_roundtrip; mod invalid_utf8; mod io; mod large_string_overflow; +mod page_index; mod parquet_testing; mod partial_offset_index; #[cfg(feature = "async")] diff --git a/parquet/tests/arrow_reader/page_index.rs b/parquet/tests/arrow_reader/page_index.rs new file mode 100644 index 0000000000..e1a8300286 --- /dev/null +++ b/parquet/tests/arrow_reader/page_index.rs @@ -0,0 +1,434 @@ +// 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. + +//! Tests of selective page index population + +use parquet::file::metadata::{ColumnChunkMask, PageIndexPolicy, ParquetMetaDataReader}; + +use crate::custom_page_index_provider::create_test_file; + +fn assert_page_index_cells( + metadata: &parquet::file::metadata::ParquetMetaData, + expect_ci: impl Fn(usize, usize) -> bool, + expect_oi: impl Fn(usize, usize) -> bool, +) { + let page_index = metadata.page_index().expect("page index should be loaded"); + let num_cols = metadata.file_metadata().schema_descr().num_columns(); + for rg in 0..metadata.num_row_groups() { + for col in 0..num_cols { + assert_eq!( + page_index.column_index(rg, col).is_some(), + expect_ci(rg, col), + "column index rg={rg} col={col}" + ); + assert_eq!( + page_index.offset_index(rg, col).is_some(), + expect_oi(rg, col), + "offset index rg={rg} col={col}" + ); + } + } +} + +#[test] +fn test_arrow_reader_options_page_index_masks() { + use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; + + let file = create_test_file(); + let column_mask = ColumnChunkMask::row_groups_and_columns([1], [0]); + let offset_mask = ColumnChunkMask::columns([2]); + let options = ArrowReaderOptions::new() + .with_page_index_policy(PageIndexPolicy::Required) + .with_column_index_mask(column_mask.clone()) + .with_offset_index_mask(offset_mask.clone()); + assert_eq!(options.column_index_mask(), &column_mask); + assert_eq!(options.offset_index_mask(), &offset_mask); + + let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); + assert_page_index_cells( + metadata.metadata(), + |rg, col| rg == 1 && col == 0, + |_, col| col == 2, + ); + + let metadata = ParquetMetaDataReader::new() + .with_arrow_reader_options(Some(&options)) + .parse_and_finish(&file) + .unwrap(); + assert_page_index_cells(&metadata, |rg, col| rg == 1 && col == 0, |_, col| col == 2); +} + +#[test] +fn test_parse_with_page_index_mask() { + let file = create_test_file(); + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .with_page_index_mask(ColumnChunkMask::row_groups_and_columns([2], [1, 3])) + .parse_and_finish(&file) + .unwrap(); + assert_page_index_cells( + &metadata, + |rg, col| rg == 2 && (col == 1 || col == 3), + |rg, col| rg == 2 && (col == 1 || col == 3), + ); +} + +#[test] +fn test_repeated_page_index_reads_replace() { + // verify that repeated calls to read_page_indexes replace existing indexes + let file = create_test_file(); + // First read: load indexes for row group 0, column 0 + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .with_page_index_mask(ColumnChunkMask::row_groups_and_columns([0], [0])) + .parse_and_finish(&file) + .unwrap(); + + // Second read: load indexes for row group 1, column 1 + // This should REPLACE the previous indexes, not merge with them + let mut reader = ParquetMetaDataReader::new_with_metadata(metadata) + .with_page_index_policy(PageIndexPolicy::Required) + .with_page_index_mask(ColumnChunkMask::row_groups_and_columns([1], [1])); + reader.read_page_indexes(&file).unwrap(); + let metadata = reader.finish().unwrap(); + + // After the second read, only row group 1, column 1 should have indexes + // The indexes from the first read (row group 0, column 0) should be gone + assert_page_index_cells( + &metadata, + |rg, col| rg == 1 && col == 1, + |rg, col| rg == 1 && col == 1, + ); +} + +#[test] +fn test_page_index_skip_policy_clears_index() { + // verify that Skip policy clears the corresponding index type + let file = create_test_file(); + // First read: load both column and offset indexes + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .parse_and_finish(&file) + .unwrap(); + + // Verify both indexes are present + let page_index = metadata.page_index().expect("page index should be loaded"); + assert!( + page_index.has_column_indexes(), + "column indexes should be present" + ); + assert!( + page_index.has_offset_indexes(), + "offset indexes should be present" + ); + + // Second read: Skip column index but load offset index + let mut reader = ParquetMetaDataReader::new_with_metadata(metadata) + .with_column_index_policy(PageIndexPolicy::Skip) + .with_offset_index_policy(PageIndexPolicy::Required); + reader.read_page_indexes(&file).unwrap(); + let metadata = reader.finish().unwrap(); + + // After the second read, only offset indexes should be present + let page_index = metadata.page_index().expect("page index should be loaded"); + assert!( + !page_index.has_column_indexes(), + "column indexes should be cleared" + ); + assert!( + page_index.has_offset_indexes(), + "offset indexes should be present" + ); + + // Third read: Skip both indexes + let mut reader = ParquetMetaDataReader::new_with_metadata(metadata) + .with_column_index_policy(PageIndexPolicy::Skip) + .with_offset_index_policy(PageIndexPolicy::Skip); + reader.read_page_indexes(&file).unwrap(); + let metadata = reader.finish().unwrap(); + + // After the third read, no page index should be present + assert!( + metadata.page_index().is_none(), + "page index should be cleared when both policies are Skip" + ); +} + +#[test] +fn test_empty_page_index_mask_clears_index() { + let file = create_test_file(); + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .parse_and_finish(&file) + .unwrap(); + assert!(metadata.page_index().is_some()); + + let mut reader = ParquetMetaDataReader::new_with_metadata(metadata) + .with_page_index_policy(PageIndexPolicy::Required) + .with_page_index_mask(ColumnChunkMask::none()); + reader.read_page_indexes(&file).unwrap(); + + assert!(reader.finish().unwrap().page_index().is_none()); +} + +#[test] +fn test_partial_page_statistics_remain_aligned() { + use arrow_array::Array; + use parquet::arrow::arrow_reader::statistics::StatisticsConverter; + use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; + + let file = create_test_file(); + let options = ArrowReaderOptions::new() + .with_page_index_policy(PageIndexPolicy::Required) + .with_column_index_mask(ColumnChunkMask::columns([0])) + .with_offset_index_mask(ColumnChunkMask::columns([1])); + let metadata = ArrowReaderMetadata::load(&file, options).unwrap(); + let parquet_metadata = metadata.metadata(); + let page_index = parquet_metadata.page_index().unwrap().as_ref(); + let row_groups = parquet_metadata.row_groups(); + let row_group_indices = [0, 1, 2]; + let converter = StatisticsConverter::try_new( + "id", + metadata.schema(), + parquet_metadata.file_metadata().schema_descr(), + ) + .unwrap(); + + let mins = converter + .data_page_mins(page_index, row_group_indices.iter()) + .unwrap(); + let row_counts = converter + .data_page_row_counts(page_index, row_groups, row_group_indices.iter()) + .unwrap() + .unwrap(); + assert_eq!(mins.len(), row_counts.len()); + assert_eq!(row_counts.null_count(), row_counts.len()); +} + +#[test] +fn test_parse_selected_columns() { + // test populating PageIndex with a subset of columns + let temp_file = create_test_file(); + + // populate column 0 for column index and columns 0 & 2 for the offset index + let mut reader = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Optional) + .with_column_index_mask(ColumnChunkMask::columns([0])) + .with_offset_index_mask(ColumnChunkMask::columns([0, 2])); + + // parse metadata + reader.try_parse(&temp_file).unwrap(); + let metadata = reader.finish().unwrap(); + let num_rg = metadata.num_row_groups(); + + // test page indexes + for rg in 0..num_rg { + let idx = metadata.page_index_for_row_group(rg); + // column 0 has both indexes + assert!(idx.column_index(0).is_some()); + assert!(idx.offset_index(0).is_some()); + // column 1 has no indexes + assert!(idx.column_index(1).is_none()); + assert!(idx.offset_index(1).is_none()); + // column 2 has offset index + assert!(idx.column_index(2).is_none()); + assert!(idx.offset_index(2).is_some()); + // column 3 has no indexes + assert!(idx.column_index(3).is_none()); + assert!(idx.offset_index(3).is_none()); + } +} + +#[test] +fn test_parse_selected_columns_mixed() { + // test populating PageIndex with a subset of columns + let temp_file = create_test_file(); + + // populate column 0 for column index and all columns for offset index + let mut reader = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Optional) + .with_column_index_mask(ColumnChunkMask::columns([0])) + .with_offset_index_mask(ColumnChunkMask::all()); + + // parse metadata + reader.try_parse(&temp_file).unwrap(); + let metadata = reader.finish().unwrap(); + let num_rg = metadata.num_row_groups(); + + // test page indexes + for rg in 0..num_rg { + let idx = metadata.page_index_for_row_group(rg); + // column 0 has both indexes + assert!(idx.column_index(0).is_some()); + assert!(idx.offset_index(0).is_some()); + // column 1 has no indexes + assert!(idx.column_index(1).is_none()); + assert!(idx.offset_index(1).is_some()); + // column 2 has offset index + assert!(idx.column_index(2).is_none()); + assert!(idx.offset_index(2).is_some()); + // column 3 has no indexes + assert!(idx.column_index(3).is_none()); + assert!(idx.offset_index(3).is_some()); + } +} + +#[test] +fn test_parse_selected_row_groups() { + // test populating PageIndex with a subset of row groups + let temp_file = create_test_file(); + + // populate indexes for row groups 0 and 2 + let mut reader = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Optional) + .with_column_index_mask(ColumnChunkMask::row_groups([0, 2])) + .with_offset_index_mask(ColumnChunkMask::row_groups([0, 2])); + + // parse metadata + reader.try_parse(&temp_file).unwrap(); + let metadata = reader.finish().unwrap(); + let num_cols = metadata.file_metadata().schema_descr().num_columns(); + assert!(metadata.page_index().is_some()); + let page_index = metadata.page_index().unwrap(); + + for col in 0..num_cols { + // row group 0 has both indexes + assert!(page_index.column_index(0, col).is_some()); + assert!(page_index.offset_index(0, col).is_some()); + // row group 1 has no indexes + assert!(page_index.column_index(1, col).is_none()); + assert!(page_index.offset_index(1, col).is_none()); + // row group 2 has both indexes + assert!(page_index.column_index(2, col).is_some()); + assert!(page_index.offset_index(2, col).is_some()); + } +} + +#[test] +fn test_parse_selected_row_groups_and_columns() { + // test populating PageIndex by row group and column + let temp_file = create_test_file(); + + // populate only row group 1, column index gets column 0, offset index gets + // columns 0 and 2. + let mut reader = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Optional) + .with_column_index_mask(ColumnChunkMask::row_groups_and_columns([1], [0])) + .with_offset_index_mask(ColumnChunkMask::row_groups_and_columns([1], [0, 2])); + + // parse metadata + reader.try_parse(&temp_file).unwrap(); + let metadata = reader.finish().unwrap(); + assert!(metadata.page_index().is_some()); + let num_cols = metadata.file_metadata().schema_descr().num_columns(); + + // row groups 0 and 2 should have no indexes + [0, 2].into_iter().for_each(|i| { + let rg_idx = metadata.page_index_for_row_group(i); + for col in 0..num_cols { + assert!(rg_idx.column_index(col).is_none()); + assert!(rg_idx.offset_index(col).is_none()); + } + }); + + // row group 1 should have column index for column 0 and offset index for columns 0 & 2 + let rg_idx = metadata.page_index_for_row_group(1); + assert!(rg_idx.column_index(0).is_some()); + assert!(rg_idx.offset_index(0).is_some()); + assert!(rg_idx.column_index(1).is_none()); + assert!(rg_idx.offset_index(1).is_none()); + assert!(rg_idx.column_index(2).is_none()); + assert!(rg_idx.offset_index(2).is_some()); + assert!(rg_idx.column_index(3).is_none()); + assert!(rg_idx.offset_index(3).is_none()); +} + +#[test] +fn test_page_index_sizes() { + // test populating PageIndex by row group and column + let temp_file = create_test_file(); + + // no index + let mut reader = ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Skip); + + // parse metadata + reader.try_parse(&temp_file).unwrap(); + let metadata = reader.finish().unwrap(); + assert!(metadata.page_index().is_none()); + #[cfg(not(feature = "encryption"))] + assert_eq!(metadata.memory_size(), 7393); + #[cfg(feature = "encryption")] + assert_eq!(metadata.memory_size(), 7817); + + // full index + let mut reader = ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required); + + // parse metadata + reader.try_parse(&temp_file).unwrap(); + let metadata = reader.finish().unwrap(); + assert!(metadata.page_index().is_some()); + #[cfg(not(feature = "encryption"))] + assert_eq!(metadata.memory_size(), 13897); + #[cfg(feature = "encryption")] + assert_eq!(metadata.memory_size(), 14321); + + // populate column 0 for column index and all columns for offset index + let mut reader = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Optional) + .with_column_index_mask(ColumnChunkMask::columns([0])) + .with_offset_index_mask(ColumnChunkMask::all()); + + // parse metadata + reader.try_parse(&temp_file).unwrap(); + let metadata = reader.finish().unwrap(); + assert!(metadata.page_index().is_some()); + #[cfg(not(feature = "encryption"))] + assert_eq!(metadata.memory_size(), 12776); + #[cfg(feature = "encryption")] + assert_eq!(metadata.memory_size(), 13200); + + // populate column 0 for column index and columns 0 & 2 for the offset index + let mut reader = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Optional) + .with_column_index_mask(ColumnChunkMask::columns([0])) + .with_offset_index_mask(ColumnChunkMask::columns([0, 2])); + + // parse metadata + reader.try_parse(&temp_file).unwrap(); + let metadata = reader.finish().unwrap(); + assert!(metadata.page_index().is_some()); + #[cfg(not(feature = "encryption"))] + assert_eq!(metadata.memory_size(), 12056); + #[cfg(feature = "encryption")] + assert_eq!(metadata.memory_size(), 12480); + + // populate only row group 1, column index gets column 0, offset index gets + // columns 0 and 2. + let mut reader = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Optional) + .with_column_index_mask(ColumnChunkMask::row_groups_and_columns([1], [0])) + .with_offset_index_mask(ColumnChunkMask::row_groups_and_columns([1], [0, 2])); + + // parse metadata + reader.try_parse(&temp_file).unwrap(); + let metadata = reader.finish().unwrap(); + assert!(metadata.page_index().is_some()); + #[cfg(not(feature = "encryption"))] + assert_eq!(metadata.memory_size(), 11326); + #[cfg(feature = "encryption")] + assert_eq!(metadata.memory_size(), 11750); +} diff --git a/parquet/tests/arrow_reader/partial_offset_index.rs b/parquet/tests/arrow_reader/partial_offset_index.rs index 72489da5c6..28aa4bb185 100644 --- a/parquet/tests/arrow_reader/partial_offset_index.rs +++ b/parquet/tests/arrow_reader/partial_offset_index.rs @@ -30,7 +30,7 @@ use parquet::arrow::push_decoder::ParquetPushDecoderBuilder; use parquet::arrow::{ArrowWriter, ProjectionMask}; use parquet::column::writer::ColumnCloseResult; use parquet::errors::Result; -use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader}; +use parquet::file::metadata::{ColumnChunkMask, PageIndexPolicy, ParquetMetaDataReader}; use parquet::file::properties::WriterProperties; use parquet::file::writer::SerializedFileWriter; use std::sync::Arc; @@ -144,15 +144,23 @@ fn configure<T>(builder: ArrowReaderBuilder<T>, trigger: Trigger) -> ArrowReader } } -fn read_sync(file: &Bytes, trigger: Trigger) -> Result<Vec<RecordBatch>> { - let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file.clone(), options())?; +fn read_sync( + file: &Bytes, + trigger: Trigger, + options: ArrowReaderOptions, +) -> Result<Vec<RecordBatch>> { + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file.clone(), options)?; Ok(configure(builder, trigger) .build()? .collect::<Result<_, _>>()?) } #[cfg(feature = "async")] -fn read_async(file: &Bytes, trigger: Trigger) -> Result<Vec<RecordBatch>> { +fn read_async( + file: &Bytes, + trigger: Trigger, + options: ArrowReaderOptions, +) -> Result<Vec<RecordBatch>> { use futures::TryStreamExt; use parquet::arrow::ParquetRecordBatchStreamBuilder; @@ -161,13 +169,17 @@ fn read_async(file: &Bytes, trigger: Trigger) -> Result<Vec<RecordBatch>> { .unwrap(); runtime.block_on(async { let input = std::io::Cursor::new(file.clone()); - let builder = ParquetRecordBatchStreamBuilder::new_with_options(input, options()).await?; + let builder = ParquetRecordBatchStreamBuilder::new_with_options(input, options).await?; configure(builder, trigger).build()?.try_collect().await }) } -fn read_push_decoder(file: &Bytes, trigger: Trigger) -> Result<Vec<RecordBatch>> { - let metadata = ArrowReaderMetadata::load(file, options())?; +fn read_push_decoder( + file: &Bytes, + trigger: Trigger, + options: ArrowReaderOptions, +) -> Result<Vec<RecordBatch>> { + let metadata = ArrowReaderMetadata::load(file, options)?; let builder = ParquetPushDecoderBuilder::new_with_metadata(metadata); let mut decoder = configure(builder, trigger).build()?; let mut batches = vec![]; @@ -202,10 +214,45 @@ fn check(has_index: impl Fn(usize, usize) -> bool) { let expected = test_batch().slice(150, 100); for trigger in [Trigger::RowSelection, Trigger::RowFilter] { let results = [ - ("sync", read_sync(&file, trigger)), - ("push decoder", read_push_decoder(&file, trigger)), + ("sync", read_sync(&file, trigger, options())), + ("push decoder", read_push_decoder(&file, trigger, options())), #[cfg(feature = "async")] - ("async", read_async(&file, trigger)), + ("async", read_async(&file, trigger, options())), + ]; + for (reader, result) in results { + let batches = result.unwrap_or_else(|e| panic!("{reader} reader, {trigger:?}: {e}")); + let actual = concat_batches(&expected.schema(), &batches).unwrap(); + assert_eq!(actual, expected, "{reader} reader, {trigger:?}"); + } + } +} + +#[test] +fn test_read_with_offset_index_mask() { + let file = test_file(|_, _| true); + let masked_options = + || options().with_offset_index_mask(ColumnChunkMask::row_groups_and_columns([1], [0])); + let metadata = ArrowReaderMetadata::load(&file, masked_options()).unwrap(); + let page_index = metadata.metadata().page_index().unwrap(); + for row_group in 0..2 { + for column in 0..3 { + assert_eq!( + page_index.offset_index(row_group, column).is_some(), + row_group == 1 && column == 0 + ); + } + } + + let expected = test_batch().slice(150, 100); + for trigger in [Trigger::RowSelection, Trigger::RowFilter] { + let results = [ + ("sync", read_sync(&file, trigger, masked_options())), + ( + "push decoder", + read_push_decoder(&file, trigger, masked_options()), + ), + #[cfg(feature = "async")] + ("async", read_async(&file, trigger, masked_options())), ]; for (reader, result) in results { let batches = result.unwrap_or_else(|e| panic!("{reader} reader, {trigger:?}: {e}"));
