You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by gi...@apache.org on 2026/10/07 04:24:49 UTC
(arrow-rs) 01/01: feat: Allow for partial population of the Parquet Page Index (#11157)
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 <et...@users.noreply.github.com>
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 <17...@users.noreply.github.com>
---
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}"));