alamb commented on code in PR #10842: URL: https://github.com/apache/arrow-rs/pull/10842#discussion_r3866791969
########## parquet/examples/custom_page_index.rs: ########## @@ -0,0 +1,259 @@ +// 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. + +//! Example of implementing a custom PageIndexProvider +//! +//! This example demonstrates how to create a custom page index provider that: +//! - Only stores page indexes for specified columns (selective storage) +//! - Uses nested HashMaps for efficient storage and lookup +//! - Implements all required PageIndexProvider trait methods +//! +//! This approach can significantly reduce memory usage when working with wide +//! tables where only a few columns need page-level statistics. + +use bytes::Bytes; +use parquet::DecodeResult; +use parquet::errors::{ParquetError, Result}; +use parquet::file::metadata::{PageIndexProvider, ParquetMetaData, ParquetMetaDataPushDecoder}; +use parquet::file::page_index::column_index::ColumnIndexMetaData; +use parquet::file::page_index::offset_index::OffsetIndexMetaData; +use std::collections::HashMap; +use std::fs::File; +use std::sync::Arc; +use tempfile::TempDir; + +/// A custom PageIndexProvider that only stores indexes for a subset of columns +/// +/// This provider uses nested hash maps to store only the necessary indexes +/// to satisfy a query +#[derive(Debug, Clone)] +struct SparsePageIndexProvider { + column_indexes: Option<HashMap<usize, HashMap<usize, ColumnIndexMetaData>>>, + offset_indexes: Option<HashMap<usize, HashMap<usize, OffsetIndexMetaData>>>, +} + +impl SparsePageIndexProvider { + fn new( + column_indexes: Option<HashMap<usize, HashMap<usize, ColumnIndexMetaData>>>, + offset_indexes: Option<HashMap<usize, HashMap<usize, OffsetIndexMetaData>>>, + ) -> Self { + Self { + column_indexes, + offset_indexes, + } + } +} + +impl PageIndexProvider for SparsePageIndexProvider { Review Comment: Is it worth annotating some comments here like this? Maybe it is too obvious / redundant ```rust /// To provide runtime page index information, you should /// implement the `PageIndexProvider` API impl PageIndexProvider for SparsePageIndexProvider { ``` ########## parquet/examples/custom_page_index.rs: ########## @@ -0,0 +1,259 @@ +// 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. + +//! Example of implementing a custom PageIndexProvider +//! +//! This example demonstrates how to create a custom page index provider that: +//! - Only stores page indexes for specified columns (selective storage) +//! - Uses nested HashMaps for efficient storage and lookup +//! - Implements all required PageIndexProvider trait methods +//! +//! This approach can significantly reduce memory usage when working with wide +//! tables where only a few columns need page-level statistics. + +use bytes::Bytes; +use parquet::DecodeResult; +use parquet::errors::{ParquetError, Result}; +use parquet::file::metadata::{PageIndexProvider, ParquetMetaData, ParquetMetaDataPushDecoder}; +use parquet::file::page_index::column_index::ColumnIndexMetaData; +use parquet::file::page_index::offset_index::OffsetIndexMetaData; +use std::collections::HashMap; +use std::fs::File; +use std::sync::Arc; +use tempfile::TempDir; + +/// A custom PageIndexProvider that only stores indexes for a subset of columns +/// +/// This provider uses nested hash maps to store only the necessary indexes +/// to satisfy a query +#[derive(Debug, Clone)] +struct SparsePageIndexProvider { + column_indexes: Option<HashMap<usize, HashMap<usize, ColumnIndexMetaData>>>, + offset_indexes: Option<HashMap<usize, HashMap<usize, OffsetIndexMetaData>>>, +} + +impl SparsePageIndexProvider { + fn new( + column_indexes: Option<HashMap<usize, HashMap<usize, ColumnIndexMetaData>>>, + offset_indexes: Option<HashMap<usize, HashMap<usize, OffsetIndexMetaData>>>, + ) -> Self { + Self { + column_indexes, + offset_indexes, + } + } +} + +impl PageIndexProvider for SparsePageIndexProvider { + fn has_offset_indexes(&self) -> bool { + self.offset_indexes.is_some() + } + + fn has_column_indexes(&self) -> bool { + self.column_indexes.is_some() + } + + fn column_index( + &self, + row_group_idx: usize, + column_idx: usize, + ) -> Option<&ColumnIndexMetaData> { + self.column_indexes + .as_ref()? + .get(&row_group_idx)? + .get(&column_idx) + } + + fn offset_index( + &self, + row_group_idx: usize, + column_idx: usize, + ) -> Option<&OffsetIndexMetaData> { + self.offset_indexes + .as_ref()? + .get(&row_group_idx)? + .get(&column_idx) + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +fn dump_page_index(metadata: &ParquetMetaData) -> Result<()> { + if let Some(page_index) = metadata.page_index() { + let num_columns = metadata.file_metadata().schema_descr().num_columns(); + + println!("Standard page index present"); + println!(" Has column indexes: {}", page_index.has_column_indexes()); + println!(" Has offset indexes: {}", page_index.has_offset_indexes()); + + println!("\nPage counts for row group 0:"); + for col_idx in 0..num_columns { + println!( + " Column {col_idx}: has offset idx {}, has column idx {}", + page_index.offset_index(0, col_idx).is_some(), + page_index.column_index(0, col_idx).is_some() + ); + } + println!(); + } else { + println!("No page index in metadata"); + println!("Note: This example requires a file with page indexes."); + println!("Try using alltypes_tiny_pages.parquet or another file with page indexes."); + return Err(ParquetError::General("no page index".to_string())); + } + Ok(()) +} + +fn main() -> Result<()> { + // Create a sample parquet file with page indexes for this example + use arrow::array::{Int32Array, StringArray}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use parquet::arrow::ArrowWriter; + use parquet::file::properties::{EnabledStatistics, WriterProperties}; + + let tempdir = TempDir::new().unwrap(); + let temp_path = tempdir.path().join("custom_page_index_example.parquet"); + println!("Creating sample file: {}", temp_path.display()); + + // Create a sample dataset with multiple columns + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + Field::new("name", DataType::Utf8, false), + Field::new("score", DataType::Int32, false), + Field::new("category", DataType::Utf8, false), + Field::new("amount", DataType::Int32, false), + ])); + + // Create multiple row groups with multiple pages + let file = File::create(&temp_path)?; + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .set_data_page_size_limit(100) // Small pages for demonstration + .set_write_batch_size(10) + .build(); + + let mut writer = ArrowWriter::try_new(file, schema.clone(), Some(props))?; + + // Write several row groups Review Comment: I think the setup and creation of the file somewhat obscures the core API -- I suggest moving the setup to a helper function (`setup_file()` or something) ########## parquet/examples/custom_page_index.rs: ########## @@ -0,0 +1,259 @@ +// 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. + +//! Example of implementing a custom PageIndexProvider +//! +//! This example demonstrates how to create a custom page index provider that: Review Comment: Nice ########## parquet/src/file/metadata/mod.rs: ########## @@ -97,11 +97,11 @@ use std::sync::Arc; pub use writer::ParquetMetaDataWriter; pub(crate) use writer::ThriftMetadataWriter; -/// Encapsulates the Parquet [Page Index] for efficient page-level data skipping +/// Trait for accessing Parquet [Page Index] data for efficient page-level skipping /// -/// The Page Index is optional metadata that enables query engines to skip irrelevant -/// data pages during scans, significantly improving I/O efficiency. It consists of two -/// complementary structures: +/// The Page Index enables query engines to skip irrelevant data pages during scans, Review Comment: Bonus points for linking to the page index parquet page: https://parquet.apache.org/docs/file-format/pageindex/ ########## parquet/src/file/metadata/mod.rs: ########## @@ -362,18 +283,460 @@ impl PageIndex { /// Returns: /// * `Some(&Vec<PageLocation>)` - Vector of page locations if offset index exists /// * `None` - Offset index not available - pub fn page_locations( + fn page_locations( &self, row_group_idx: usize, column_idx: usize, ) -> Option<&Vec<PageLocation>> { - if let Some(offset_indexes) = self.offset_indexes.as_ref() { - let rg = offset_indexes.get(row_group_idx)?; - let off_idx = rg.get(column_idx)?.as_ref()?; - Some(off_idx.page_locations()) + Some( + self.offset_index(row_group_idx, column_idx)? + .page_locations(), + ) + } + + /// Returns a reference to the trait object as `&dyn Any` for downcasting + /// + /// This allows downcasting to concrete types when needed (e.g., for serialization) + fn as_any(&self) -> &dyn std::any::Any; +} + +/// Provides convenient access to page index data for a specific row group +/// +/// This struct wraps a [`PageIndexProvider`] and automatically applies the row group +/// index, simplifying access to column and offset indexes for a single row group. +/// It is primarily used by readers to avoid repeatedly passing the row group index +/// when accessing page-level metadata. +/// +/// # Example +/// +/// ``` +/// use parquet::file::metadata::ParquetMetaData; +/// # use parquet::errors::Result; +/// +/// fn process_row_group_pages(metadata: &ParquetMetaData, row_group_idx: usize) -> Result<()> { +/// if let Some(page_index) = metadata.page_index() { +/// // Create a row-group-specific view +/// let rg_page_index = parquet::file::metadata::RowGroupPageIndex::new( +/// row_group_idx, +/// metadata.page_index().cloned(), +/// ); +/// +/// // Now access column indexes without specifying row_group_idx each time +/// for col_idx in 0..metadata.file_metadata().schema_descr().num_columns() { +/// if let Some(col_idx_data) = rg_page_index.column_index(col_idx) { +/// println!("Column {} has {} pages", col_idx, col_idx_data.num_pages()); +/// } +/// } +/// } +/// Ok(()) +/// } +/// ``` +#[derive(Debug)] +pub struct RowGroupPageIndex { Review Comment: this is a neat idea and will likely be useful downtream ########## parquet/src/file/metadata/mod.rs: ########## @@ -362,18 +283,460 @@ impl PageIndex { /// Returns: /// * `Some(&Vec<PageLocation>)` - Vector of page locations if offset index exists /// * `None` - Offset index not available - pub fn page_locations( + fn page_locations( &self, row_group_idx: usize, column_idx: usize, ) -> Option<&Vec<PageLocation>> { - if let Some(offset_indexes) = self.offset_indexes.as_ref() { - let rg = offset_indexes.get(row_group_idx)?; - let off_idx = rg.get(column_idx)?.as_ref()?; - Some(off_idx.page_locations()) + Some( + self.offset_index(row_group_idx, column_idx)? + .page_locations(), + ) + } + + /// Returns a reference to the trait object as `&dyn Any` for downcasting + /// + /// This allows downcasting to concrete types when needed (e.g., for serialization) + fn as_any(&self) -> &dyn std::any::Any; +} + +/// Provides convenient access to page index data for a specific row group +/// +/// This struct wraps a [`PageIndexProvider`] and automatically applies the row group +/// index, simplifying access to column and offset indexes for a single row group. +/// It is primarily used by readers to avoid repeatedly passing the row group index +/// when accessing page-level metadata. +/// +/// # Example +/// +/// ``` +/// use parquet::file::metadata::ParquetMetaData; +/// # use parquet::errors::Result; +/// +/// fn process_row_group_pages(metadata: &ParquetMetaData, row_group_idx: usize) -> Result<()> { +/// if let Some(page_index) = metadata.page_index() { +/// // Create a row-group-specific view +/// let rg_page_index = parquet::file::metadata::RowGroupPageIndex::new( +/// row_group_idx, +/// metadata.page_index().cloned(), +/// ); +/// +/// // Now access column indexes without specifying row_group_idx each time +/// for col_idx in 0..metadata.file_metadata().schema_descr().num_columns() { +/// if let Some(col_idx_data) = rg_page_index.column_index(col_idx) { +/// println!("Column {} has {} pages", col_idx, col_idx_data.num_pages()); +/// } +/// } +/// } +/// Ok(()) +/// } +/// ``` +#[derive(Debug)] +pub struct RowGroupPageIndex { + row_group_idx: usize, + page_index: Option<Arc<dyn PageIndexProvider>>, +} + +impl RowGroupPageIndex { + /// Creates a new [`RowGroupPageIndex`] for the specified row group + /// + /// # Arguments + /// + /// * `row_group_idx` - The index of the row group within the file + /// * `page_index` - Optional page index provider containing the index data + pub fn new(row_group_idx: usize, page_index: Option<Arc<dyn PageIndexProvider>>) -> Self { + Self { + row_group_idx, + page_index, + } + } + + /// Returns the column index for a specific column in this row group + /// + /// This is a convenience method that wraps [`PageIndexProvider::column_index`], + /// automatically applying the row group index stored in this struct. + /// + /// # Returns + /// + /// * `Some(&ColumnIndexMetaData)` - Column index is available with page-level statistics + /// * `None` - Index unavailable (no page index, column out of bounds, or no statistics) + /// + /// # See Also + /// + /// * [`PageIndexProvider::column_index`] for more details on column indexes + pub fn column_index(&self, column_idx: usize) -> Option<&ColumnIndexMetaData> { + self.page_index + .as_ref()? + .column_index(self.row_group_idx, column_idx) + } + + /// Returns the offset index for a specific column in this row group + /// + /// This is a convenience method that wraps [`PageIndexProvider::offset_index`], + /// automatically applying the row group index stored in this struct. + /// + /// # Returns + /// + /// * `Some(&OffsetIndexMetaData)` - Offset index is available with page locations + /// * `None` - Index unavailable (no page index, column out of bounds) + /// + /// # See Also + /// + /// * [`PageIndexProvider::offset_index`] for more details on offset indexes + pub fn offset_index(&self, column_idx: usize) -> Option<&OffsetIndexMetaData> { + self.page_index + .as_ref()? + .offset_index(self.row_group_idx, column_idx) + } + + /// Returns the physical locations of all data pages for a specific column in this row group + /// + /// This is a convenience method that wraps [`PageIndexProvider::page_locations`], + /// automatically applying the row group index stored in this struct. + /// + /// This enables direct I/O to specific pages without reading the entire column chunk. + /// + /// # Returns + /// + /// * `Some(&Vec<PageLocation>)` - Vector of page locations if offset index exists + /// * `None` - Offset index not available for this column + /// + /// # See Also + /// + /// * [`PageIndexProvider::page_locations`] for more details on page locations + pub fn page_locations(&self, column_idx: usize) -> Option<&Vec<PageLocation>> { + Some(self.offset_index(column_idx)?.page_locations()) + } + + /// Returns the expected number of data pages for a specific column in this row group + /// + /// This count includes only data pages, not dictionary pages or other metadata pages. + /// + /// This is a convenience method that wraps [`PageIndexProvider::num_data_pages`], + /// automatically applying the row group index stored in this struct. + /// + /// # Returns + /// + /// * `Some(usize)` - Number of data pages if any index is available + /// * `None` - No index information available for this column + /// + /// # See Also + /// + /// * [`PageIndexProvider::num_data_pages`] for more details + pub fn num_data_pages(&self, column_idx: usize) -> Option<usize> { + self.page_index + .as_ref()? + .num_data_pages(self.row_group_idx, column_idx) + } +} + +/// Struct to encapsulate the Parquet [Page Index] +/// +/// This struct provides a dense representation of the Page Index. It is +/// used internally by this crate when assembling and writing the Page +/// Index. It is also the default implmentation of the [`PageIndexProvider`] +/// contained in the [`ParquetMetaData`]. +/// +/// # Example: Constructing a synthetic `PageIndex` +/// +/// This example builds a [`ParquetMetaData`] for a file with a single row +/// group containing a single `BYTE_ARRAY` column with one data page, and +/// attaches a matching `PageIndex`, as might be done in tests that +/// exercise page-level statistics handling. +/// +/// ``` +/// # use std::sync::Arc; +/// # use parquet::basic::{BoundaryOrder, Type as PhysicalType}; +/// # use parquet::file::metadata::{ +/// # ColumnChunkMetaData, ColumnIndexBuilder, FileMetaData, OffsetIndexBuilder, +/// # PageIndexBuilder, ParquetMetaData, RowGroupMetaData, +/// # }; +/// # use parquet::schema::types::{SchemaDescriptor, Type}; +/// // Create metadata for a file with a single row group containing a +/// // single BYTE_ARRAY column "s" with three values +/// let schema = Arc::new(SchemaDescriptor::new(Arc::new( +/// Type::group_type_builder("schema") +/// .with_fields(vec![Arc::new( +/// Type::primitive_type_builder("s", PhysicalType::BYTE_ARRAY) +/// .build() +/// .unwrap(), +/// )]) +/// .build() +/// .unwrap(), +/// ))); +/// let column = ColumnChunkMetaData::builder(schema.column(0)) +/// .set_num_values(3) +/// .build() +/// .unwrap(); +/// let row_group = RowGroupMetaData::builder(Arc::clone(&schema)) +/// .set_num_rows(3) +/// .set_column_metadata(vec![column]) +/// .build() +/// .unwrap(); +/// let file_metadata = FileMetaData::new(1, 3, None, None, schema, None); +/// let metadata = ParquetMetaData::new(file_metadata, vec![row_group]); +/// +/// // Build a column index with min/max statistics for the single page +/// let mut column_index = ColumnIndexBuilder::new(PhysicalType::BYTE_ARRAY); +/// column_index.append(false, b"az".to_vec(), b"b".to_vec(), 0, None); +/// column_index.set_boundary_order(BoundaryOrder::ASCENDING); +/// let column_index = column_index.build().unwrap(); +/// +/// // Build an offset index recording the location of the single page +/// let mut offset_index = OffsetIndexBuilder::new(); +/// offset_index.append_row_count(3); +/// offset_index.append_offset_and_size(4, 100); +/// let offset_index = offset_index.build(); +/// +/// // Assemble the PageIndex (one entry per row group, each with one +/// // entry per column) and attach it to the metadata +/// let mut page_index = PageIndexBuilder::new(1, 1); +/// page_index.put_column_index(column_index, 0, 0); +/// page_index.put_offset_index(offset_index, 0, 0); +/// let page_index = page_index.build(); +/// let metadata = metadata +/// .into_builder() +/// .set_page_index(Some(Arc::new(page_index))) +/// .build(); +/// assert!(metadata.page_index().unwrap().is_complete()); +/// ``` +/// +/// [Page Index]: https://github.com/apache/parquet-format/blob/master/PageIndex.md +/// [`ColumnIndex`]: crate::file::page_index::column_index::ColumnIndexMetaData +/// [`OffsetIndex`]: crate::file::page_index::offset_index::OffsetIndexMetaData +#[derive(Debug, Clone, PartialEq)] +pub struct PageIndex { + column_indexes: Option<Vec<Vec<Option<ColumnIndexMetaData>>>>, + offset_indexes: Option<Vec<Vec<Option<OffsetIndexMetaData>>>>, +} + +impl PageIndex { + pub(crate) fn new( + column_indexes: Option<Vec<Vec<Option<ColumnIndexMetaData>>>>, + offset_indexes: Option<Vec<Vec<Option<OffsetIndexMetaData>>>>, + ) -> Self { + Self { + column_indexes, + offset_indexes, + } + } + + /// Convert this `PageIndex` into a [`PageIndexBuilder`] + pub fn into_builder(self) -> PageIndexBuilder { + self.into() + } + + /// Returns a reference to the raw column indexes structure + /// + /// This method provides access to the underlying column index data for serialization + /// and other low-level operations. + pub(crate) fn column_indexes_raw(&self) -> Option<&Vec<Vec<Option<ColumnIndexMetaData>>>> { + self.column_indexes.as_ref() + } + + /// Returns a reference to the raw offset indexes structure + /// + /// This method provides access to the underlying offset index data for serialization + /// and other low-level operations. + pub(crate) fn offset_indexes_raw(&self) -> Option<&Vec<Vec<Option<OffsetIndexMetaData>>>> { + self.offset_indexes.as_ref() + } +} + +impl PageIndexProvider for PageIndex { + fn has_offset_indexes(&self) -> bool { + self.offset_indexes.is_some() + } + + fn has_column_indexes(&self) -> bool { + self.column_indexes.is_some() + } + + fn column_index( + &self, + row_group_idx: usize, + column_idx: usize, + ) -> Option<&ColumnIndexMetaData> { + let rg = self.column_indexes.as_ref()?.get(row_group_idx)?; + rg.get(column_idx)?.as_ref() + } + + fn offset_index( + &self, + row_group_idx: usize, + column_idx: usize, + ) -> Option<&OffsetIndexMetaData> { + let rg = self.offset_indexes.as_ref()?.get(row_group_idx)?; + rg.get(column_idx)?.as_ref() + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +/// Builder for constructing [`PageIndex`] structures +/// +/// It supports: +/// - Allocating space for indexes based on [`PageIndexPolicy`] +/// - Populating column indexes for predicate columns (for page filtering) +/// - Populating offset indexes for projected columns (for direct I/O) +/// - Automatic conversion of empty structures to `None` to save memory +pub struct PageIndexBuilder { + column_indexes: Option<Vec<Vec<Option<ColumnIndexMetaData>>>>, + offset_indexes: Option<Vec<Vec<Option<OffsetIndexMetaData>>>>, +} + +impl PageIndexBuilder { + /// Creates an empty index structure with space for the specified number of row groups and columns Review Comment: I recommend we also give guidance on what the `T` is (and how to find the corresponding T for each parquet type) ########## parquet/src/file/metadata/mod.rs: ########## @@ -362,18 +283,460 @@ impl PageIndex { /// Returns: /// * `Some(&Vec<PageLocation>)` - Vector of page locations if offset index exists /// * `None` - Offset index not available - pub fn page_locations( + fn page_locations( &self, row_group_idx: usize, column_idx: usize, ) -> Option<&Vec<PageLocation>> { - if let Some(offset_indexes) = self.offset_indexes.as_ref() { - let rg = offset_indexes.get(row_group_idx)?; - let off_idx = rg.get(column_idx)?.as_ref()?; - Some(off_idx.page_locations()) + Some( + self.offset_index(row_group_idx, column_idx)? + .page_locations(), + ) + } + + /// Returns a reference to the trait object as `&dyn Any` for downcasting + /// + /// This allows downcasting to concrete types when needed (e.g., for serialization) + fn as_any(&self) -> &dyn std::any::Any; +} + +/// Provides convenient access to page index data for a specific row group +/// +/// This struct wraps a [`PageIndexProvider`] and automatically applies the row group +/// index, simplifying access to column and offset indexes for a single row group. +/// It is primarily used by readers to avoid repeatedly passing the row group index +/// when accessing page-level metadata. +/// +/// # Example +/// +/// ``` +/// use parquet::file::metadata::ParquetMetaData; +/// # use parquet::errors::Result; +/// +/// fn process_row_group_pages(metadata: &ParquetMetaData, row_group_idx: usize) -> Result<()> { +/// if let Some(page_index) = metadata.page_index() { +/// // Create a row-group-specific view +/// let rg_page_index = parquet::file::metadata::RowGroupPageIndex::new( +/// row_group_idx, +/// metadata.page_index().cloned(), +/// ); +/// +/// // Now access column indexes without specifying row_group_idx each time +/// for col_idx in 0..metadata.file_metadata().schema_descr().num_columns() { +/// if let Some(col_idx_data) = rg_page_index.column_index(col_idx) { +/// println!("Column {} has {} pages", col_idx, col_idx_data.num_pages()); +/// } +/// } +/// } +/// Ok(()) +/// } +/// ``` +#[derive(Debug)] +pub struct RowGroupPageIndex { + row_group_idx: usize, + page_index: Option<Arc<dyn PageIndexProvider>>, +} + +impl RowGroupPageIndex { + /// Creates a new [`RowGroupPageIndex`] for the specified row group + /// + /// # Arguments + /// + /// * `row_group_idx` - The index of the row group within the file + /// * `page_index` - Optional page index provider containing the index data + pub fn new(row_group_idx: usize, page_index: Option<Arc<dyn PageIndexProvider>>) -> Self { + Self { + row_group_idx, + page_index, + } + } + + /// Returns the column index for a specific column in this row group + /// + /// This is a convenience method that wraps [`PageIndexProvider::column_index`], + /// automatically applying the row group index stored in this struct. + /// + /// # Returns + /// + /// * `Some(&ColumnIndexMetaData)` - Column index is available with page-level statistics + /// * `None` - Index unavailable (no page index, column out of bounds, or no statistics) + /// + /// # See Also + /// + /// * [`PageIndexProvider::column_index`] for more details on column indexes + pub fn column_index(&self, column_idx: usize) -> Option<&ColumnIndexMetaData> { + self.page_index + .as_ref()? + .column_index(self.row_group_idx, column_idx) + } + + /// Returns the offset index for a specific column in this row group + /// + /// This is a convenience method that wraps [`PageIndexProvider::offset_index`], + /// automatically applying the row group index stored in this struct. + /// + /// # Returns + /// + /// * `Some(&OffsetIndexMetaData)` - Offset index is available with page locations + /// * `None` - Index unavailable (no page index, column out of bounds) + /// + /// # See Also + /// + /// * [`PageIndexProvider::offset_index`] for more details on offset indexes + pub fn offset_index(&self, column_idx: usize) -> Option<&OffsetIndexMetaData> { + self.page_index + .as_ref()? + .offset_index(self.row_group_idx, column_idx) + } + + /// Returns the physical locations of all data pages for a specific column in this row group + /// + /// This is a convenience method that wraps [`PageIndexProvider::page_locations`], + /// automatically applying the row group index stored in this struct. + /// + /// This enables direct I/O to specific pages without reading the entire column chunk. + /// + /// # Returns + /// + /// * `Some(&Vec<PageLocation>)` - Vector of page locations if offset index exists + /// * `None` - Offset index not available for this column + /// + /// # See Also + /// + /// * [`PageIndexProvider::page_locations`] for more details on page locations + pub fn page_locations(&self, column_idx: usize) -> Option<&Vec<PageLocation>> { + Some(self.offset_index(column_idx)?.page_locations()) + } + + /// Returns the expected number of data pages for a specific column in this row group + /// + /// This count includes only data pages, not dictionary pages or other metadata pages. + /// + /// This is a convenience method that wraps [`PageIndexProvider::num_data_pages`], + /// automatically applying the row group index stored in this struct. + /// + /// # Returns + /// + /// * `Some(usize)` - Number of data pages if any index is available + /// * `None` - No index information available for this column + /// + /// # See Also + /// + /// * [`PageIndexProvider::num_data_pages`] for more details + pub fn num_data_pages(&self, column_idx: usize) -> Option<usize> { + self.page_index + .as_ref()? + .num_data_pages(self.row_group_idx, column_idx) + } +} + +/// Struct to encapsulate the Parquet [Page Index] +/// +/// This struct provides a dense representation of the Page Index. It is +/// used internally by this crate when assembling and writing the Page +/// Index. It is also the default implmentation of the [`PageIndexProvider`] +/// contained in the [`ParquetMetaData`]. +/// +/// # Example: Constructing a synthetic `PageIndex` +/// +/// This example builds a [`ParquetMetaData`] for a file with a single row +/// group containing a single `BYTE_ARRAY` column with one data page, and +/// attaches a matching `PageIndex`, as might be done in tests that +/// exercise page-level statistics handling. +/// +/// ``` +/// # use std::sync::Arc; +/// # use parquet::basic::{BoundaryOrder, Type as PhysicalType}; +/// # use parquet::file::metadata::{ +/// # ColumnChunkMetaData, ColumnIndexBuilder, FileMetaData, OffsetIndexBuilder, +/// # PageIndexBuilder, ParquetMetaData, RowGroupMetaData, +/// # }; +/// # use parquet::schema::types::{SchemaDescriptor, Type}; +/// // Create metadata for a file with a single row group containing a +/// // single BYTE_ARRAY column "s" with three values +/// let schema = Arc::new(SchemaDescriptor::new(Arc::new( +/// Type::group_type_builder("schema") +/// .with_fields(vec![Arc::new( +/// Type::primitive_type_builder("s", PhysicalType::BYTE_ARRAY) +/// .build() +/// .unwrap(), +/// )]) +/// .build() +/// .unwrap(), +/// ))); +/// let column = ColumnChunkMetaData::builder(schema.column(0)) +/// .set_num_values(3) +/// .build() +/// .unwrap(); +/// let row_group = RowGroupMetaData::builder(Arc::clone(&schema)) +/// .set_num_rows(3) +/// .set_column_metadata(vec![column]) +/// .build() +/// .unwrap(); +/// let file_metadata = FileMetaData::new(1, 3, None, None, schema, None); +/// let metadata = ParquetMetaData::new(file_metadata, vec![row_group]); +/// +/// // Build a column index with min/max statistics for the single page +/// let mut column_index = ColumnIndexBuilder::new(PhysicalType::BYTE_ARRAY); +/// column_index.append(false, b"az".to_vec(), b"b".to_vec(), 0, None); +/// column_index.set_boundary_order(BoundaryOrder::ASCENDING); +/// let column_index = column_index.build().unwrap(); +/// +/// // Build an offset index recording the location of the single page +/// let mut offset_index = OffsetIndexBuilder::new(); +/// offset_index.append_row_count(3); +/// offset_index.append_offset_and_size(4, 100); +/// let offset_index = offset_index.build(); +/// +/// // Assemble the PageIndex (one entry per row group, each with one +/// // entry per column) and attach it to the metadata +/// let mut page_index = PageIndexBuilder::new(1, 1); +/// page_index.put_column_index(column_index, 0, 0); +/// page_index.put_offset_index(offset_index, 0, 0); +/// let page_index = page_index.build(); +/// let metadata = metadata +/// .into_builder() +/// .set_page_index(Some(Arc::new(page_index))) +/// .build(); +/// assert!(metadata.page_index().unwrap().is_complete()); +/// ``` +/// +/// [Page Index]: https://github.com/apache/parquet-format/blob/master/PageIndex.md +/// [`ColumnIndex`]: crate::file::page_index::column_index::ColumnIndexMetaData +/// [`OffsetIndex`]: crate::file::page_index::offset_index::OffsetIndexMetaData +#[derive(Debug, Clone, PartialEq)] +pub struct PageIndex { + column_indexes: Option<Vec<Vec<Option<ColumnIndexMetaData>>>>, + offset_indexes: Option<Vec<Vec<Option<OffsetIndexMetaData>>>>, +} + +impl PageIndex { + pub(crate) fn new( + column_indexes: Option<Vec<Vec<Option<ColumnIndexMetaData>>>>, + offset_indexes: Option<Vec<Vec<Option<OffsetIndexMetaData>>>>, + ) -> Self { + Self { + column_indexes, + offset_indexes, + } + } + + /// Convert this `PageIndex` into a [`PageIndexBuilder`] + pub fn into_builder(self) -> PageIndexBuilder { + self.into() + } + + /// Returns a reference to the raw column indexes structure + /// + /// This method provides access to the underlying column index data for serialization + /// and other low-level operations. + pub(crate) fn column_indexes_raw(&self) -> Option<&Vec<Vec<Option<ColumnIndexMetaData>>>> { + self.column_indexes.as_ref() + } + + /// Returns a reference to the raw offset indexes structure + /// + /// This method provides access to the underlying offset index data for serialization + /// and other low-level operations. + pub(crate) fn offset_indexes_raw(&self) -> Option<&Vec<Vec<Option<OffsetIndexMetaData>>>> { + self.offset_indexes.as_ref() + } +} + +impl PageIndexProvider for PageIndex { + fn has_offset_indexes(&self) -> bool { + self.offset_indexes.is_some() + } + + fn has_column_indexes(&self) -> bool { + self.column_indexes.is_some() + } + + fn column_index( + &self, + row_group_idx: usize, + column_idx: usize, + ) -> Option<&ColumnIndexMetaData> { + let rg = self.column_indexes.as_ref()?.get(row_group_idx)?; + rg.get(column_idx)?.as_ref() + } + + fn offset_index( + &self, + row_group_idx: usize, + column_idx: usize, + ) -> Option<&OffsetIndexMetaData> { + let rg = self.offset_indexes.as_ref()?.get(row_group_idx)?; + rg.get(column_idx)?.as_ref() + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +/// Builder for constructing [`PageIndex`] structures +/// +/// It supports: +/// - Allocating space for indexes based on [`PageIndexPolicy`] +/// - Populating column indexes for predicate columns (for page filtering) +/// - Populating offset indexes for projected columns (for direct I/O) +/// - Automatic conversion of empty structures to `None` to save memory +pub struct PageIndexBuilder { + column_indexes: Option<Vec<Vec<Option<ColumnIndexMetaData>>>>, + offset_indexes: Option<Vec<Vec<Option<OffsetIndexMetaData>>>>, +} + +impl PageIndexBuilder { + /// Creates an empty index structure with space for the specified number of row groups and columns + /// + /// Returns `Some` containing a nested vector structure where all entries are initialized to `None`. + /// The outer vector has one entry per row group, and each inner vector has one entry per column. + fn empty_index<T>(num_row_groups: usize, num_columns: usize) -> Option<Vec<Vec<Option<T>>>> { + Some( + (0..num_row_groups) + .map(|_| { + let mut idx = Vec::with_capacity(num_columns); + idx.resize_with(num_columns, || None); + idx + }) + .collect(), + ) + } + + /// Creates a new [`PageIndexBuilder`] with space allocated for both column and offset indexes + /// + /// This allocates empty index structures for the specified number of row groups and columns. + /// All index entries are initialized to `None` and can be populated using + /// [`put_column_index`](Self::put_column_index) and [`put_offset_index`](Self::put_offset_index). + pub fn new(num_row_groups: usize, num_columns: usize) -> Self { + Self { + column_indexes: Self::empty_index(num_row_groups, num_columns), + offset_indexes: Self::empty_index(num_row_groups, num_columns), + } + } + + /// Creates a new [`PageIndexBuilder`] with selective allocation based on policies + /// + /// This allows fine-grained control over which indexes are allocated: + /// - [`PageIndexPolicy::Skip`]: No allocation, the index structure is set to `None` Review Comment: does it also mean that all subsequent operations on the builder are NoOps? This feels somewhat niche to me (like what else is going to use this API other than the ParquetMetadataDecoder?) Maybe we can put the `Option` in the metadata decoder itself and leave the builder a straigh builder ########## parquet/examples/custom_page_index.rs: ########## @@ -0,0 +1,259 @@ +// 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. + +//! Example of implementing a custom PageIndexProvider +//! +//! This example demonstrates how to create a custom page index provider that: +//! - Only stores page indexes for specified columns (selective storage) +//! - Uses nested HashMaps for efficient storage and lookup +//! - Implements all required PageIndexProvider trait methods +//! +//! This approach can significantly reduce memory usage when working with wide +//! tables where only a few columns need page-level statistics. + +use bytes::Bytes; +use parquet::DecodeResult; +use parquet::errors::{ParquetError, Result}; +use parquet::file::metadata::{PageIndexProvider, ParquetMetaData, ParquetMetaDataPushDecoder}; +use parquet::file::page_index::column_index::ColumnIndexMetaData; +use parquet::file::page_index::offset_index::OffsetIndexMetaData; +use std::collections::HashMap; +use std::fs::File; +use std::sync::Arc; +use tempfile::TempDir; + +/// A custom PageIndexProvider that only stores indexes for a subset of columns +/// +/// This provider uses nested hash maps to store only the necessary indexes +/// to satisfy a query +#[derive(Debug, Clone)] +struct SparsePageIndexProvider { + column_indexes: Option<HashMap<usize, HashMap<usize, ColumnIndexMetaData>>>, + offset_indexes: Option<HashMap<usize, HashMap<usize, OffsetIndexMetaData>>>, +} + +impl SparsePageIndexProvider { + fn new( + column_indexes: Option<HashMap<usize, HashMap<usize, ColumnIndexMetaData>>>, + offset_indexes: Option<HashMap<usize, HashMap<usize, OffsetIndexMetaData>>>, + ) -> Self { + Self { + column_indexes, + offset_indexes, + } + } +} + +impl PageIndexProvider for SparsePageIndexProvider { + fn has_offset_indexes(&self) -> bool { + self.offset_indexes.is_some() + } + + fn has_column_indexes(&self) -> bool { + self.column_indexes.is_some() + } + + fn column_index( + &self, + row_group_idx: usize, + column_idx: usize, + ) -> Option<&ColumnIndexMetaData> { + self.column_indexes + .as_ref()? + .get(&row_group_idx)? + .get(&column_idx) + } + + fn offset_index( + &self, + row_group_idx: usize, + column_idx: usize, + ) -> Option<&OffsetIndexMetaData> { + self.offset_indexes + .as_ref()? + .get(&row_group_idx)? + .get(&column_idx) + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +fn dump_page_index(metadata: &ParquetMetaData) -> Result<()> { + if let Some(page_index) = metadata.page_index() { + let num_columns = metadata.file_metadata().schema_descr().num_columns(); + + println!("Standard page index present"); + println!(" Has column indexes: {}", page_index.has_column_indexes()); + println!(" Has offset indexes: {}", page_index.has_offset_indexes()); + + println!("\nPage counts for row group 0:"); + for col_idx in 0..num_columns { + println!( + " Column {col_idx}: has offset idx {}, has column idx {}", + page_index.offset_index(0, col_idx).is_some(), + page_index.column_index(0, col_idx).is_some() + ); + } + println!(); + } else { + println!("No page index in metadata"); + println!("Note: This example requires a file with page indexes."); + println!("Try using alltypes_tiny_pages.parquet or another file with page indexes."); + return Err(ParquetError::General("no page index".to_string())); + } + Ok(()) +} + +fn main() -> Result<()> { + // Create a sample parquet file with page indexes for this example + use arrow::array::{Int32Array, StringArray}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use parquet::arrow::ArrowWriter; + use parquet::file::properties::{EnabledStatistics, WriterProperties}; + + let tempdir = TempDir::new().unwrap(); + let temp_path = tempdir.path().join("custom_page_index_example.parquet"); + println!("Creating sample file: {}", temp_path.display()); + + // Create a sample dataset with multiple columns + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + Field::new("name", DataType::Utf8, false), + Field::new("score", DataType::Int32, false), + Field::new("category", DataType::Utf8, false), + Field::new("amount", DataType::Int32, false), + ])); + + // Create multiple row groups with multiple pages + let file = File::create(&temp_path)?; + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .set_data_page_size_limit(100) // Small pages for demonstration + .set_write_batch_size(10) + .build(); + + let mut writer = ArrowWriter::try_new(file, schema.clone(), Some(props))?; + + // Write several row groups + for row_group in 0..3 { + for batch_num in 0..5 { + let offset = (row_group * 50) + (batch_num * 10); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from( + (offset..offset + 10).collect::<Vec<i32>>(), + )), + Arc::new(Int32Array::from( + (offset..offset + 10).map(|x| x * 2).collect::<Vec<i32>>(), + )), + Arc::new(StringArray::from( + (offset..offset + 10) + .map(|x| format!("name{x}")) + .collect::<Vec<String>>(), + )), + Arc::new(Int32Array::from( + (offset..offset + 10).map(|x| x * 3).collect::<Vec<i32>>(), + )), + Arc::new(StringArray::from( + (offset..offset + 10) + .map(|x| if x % 2 == 0 { "even" } else { "odd" }) + .collect::<Vec<&str>>(), + )), + Arc::new(Int32Array::from( + (offset..offset + 10).map(|x| x * 4).collect::<Vec<i32>>(), + )), + ], + )?; + writer.write(&batch)?; + } + writer.flush()?; + } + + writer.close()?; + println!("Sample file created with page indexes\n"); + + // Now read it back with page indexes + use parquet::file::metadata::PageIndexPolicy; + + let file_bytes = Bytes::from(std::fs::read(temp_path)?); + let file_len = file_bytes.len() as u64; + let mut decoder = ParquetMetaDataPushDecoder::try_new(file_len)? + .with_page_index_policy(PageIndexPolicy::Required); + #[expect(clippy::single_range_in_vec_init)] + decoder.push_ranges(vec![0..file_len], vec![file_bytes.clone()])?; + let metadata = match decoder.try_decode() { + Ok(DecodeResult::Data(metadata)) => metadata, // decode successful + other => { + panic!("expected DecodeResult::Data, got: {other:?}") + } + }; + + let num_columns = metadata.file_metadata().schema_descr().num_columns(); + println!("Number of row groups: {}", metadata.num_row_groups()); + println!("Number of columns: {num_columns}"); + println!(); + + // Example 1: Use the standard PageIndex provider (all columns accessible) + println!("=== Example 1: Standard PageIndex (all columns) ==="); + dump_page_index(&metadata)?; + + // Save original index to populate the custom one. A real application could + // cache the indexes externally. + let page_index = metadata.page_index().cloned().unwrap(); + + // Example 2: Read metadata and then add custom PageIndexProvider + decoder = ParquetMetaDataPushDecoder::try_new(file_len)? + .with_page_index_policy(PageIndexPolicy::Skip); + #[expect(clippy::single_range_in_vec_init)] + decoder.push_ranges(vec![0..file_len], vec![file_bytes])?; + let metadata = match decoder.try_decode() { + Ok(DecodeResult::Data(metadata)) => metadata, // decode successful + other => { + panic!("expected DecodeResult::Data, got: {other:?}") + } + }; + println!("=== Example 2: Selective PageIndex (columns 0, 1, 4 only) ==="); + let mut builder = metadata.into_builder(); + + // create partial indexes. column index for column 0 only (predicate column), Review Comment: This is kind of strange to decode the entire metadata (and all page indexes) but then only serve a subset -- it kind of defeats the use case of not loading the entire thing I think a more compelling example (and one that would more likely be used) would be -- a "on demand" page index loader -- basically something that only loaded the page indexes on demand (maybe caching them). Is that possible with this API? I think we could do it if you loaded the metadata *without* the page index and then implemented your own read from the underling file for the relvant structures 🤔 ########## parquet/src/file/metadata/mod.rs: ########## @@ -209,70 +208,31 @@ pub(crate) use writer::ThriftMetadataWriter; /// [Page Index]: https://github.com/apache/parquet-format/blob/master/PageIndex.md /// [`ColumnIndex`]: crate::file::page_index::column_index::ColumnIndexMetaData /// [`OffsetIndex`]: crate::file::page_index::offset_index::OffsetIndexMetaData -#[derive(Debug, Clone, PartialEq)] -pub struct PageIndex { - column_indexes: Option<Vec<Vec<Option<ColumnIndexMetaData>>>>, - offset_indexes: Option<Vec<Vec<Option<OffsetIndexMetaData>>>>, -} - -impl PageIndex { - pub(crate) fn new( - column_indexes: Option<Vec<Vec<Option<ColumnIndexMetaData>>>>, - offset_indexes: Option<Vec<Vec<Option<OffsetIndexMetaData>>>>, - ) -> Self { - Self { - column_indexes, - offset_indexes, - } - } - +pub trait PageIndexProvider: Send + Sync + std::fmt::Debug { /// Returns `true` if offset index structures are present /// /// This indicates whether [`OffsetIndexMetaData`] structures were loaded or created. /// Returns `true` even if some individual columns lack offset indexes. /// /// To check if a specific column has an offset index, use [`Self::offset_index`]. - pub fn has_offset_indexes(&self) -> bool { - self.offset_indexes.is_some() - } + fn has_offset_indexes(&self) -> bool; /// Returns `true` if column index structures are present /// /// This indicates whether [`ColumnIndexMetaData`] structures were loaded or created. /// Returns `true` even if some individual columns lack column indexes. /// /// To check if a specific column has a column index, use [`Self::column_index`]. - pub fn has_column_indexes(&self) -> bool { - self.column_indexes.is_some() - } + fn has_column_indexes(&self) -> bool; /// Returns `true` if both the offset and column index structures are present /// /// This is equivalent to both [`Self::has_offset_indexes`] and [`Self::has_column_indexes`] /// returning `true`. - pub fn is_complete(&self) -> bool { + fn is_complete(&self) -> bool { self.has_column_indexes() && self.has_offset_indexes() } - /// Returns column indexes for all columns in the specified row group - /// - /// Returns `None` if: - /// - Column indexes were not loaded or are not available - /// - The row group index is out of bounds - /// - /// Returns `Some(&[Option<ColumnIndexMetaData>])` where: - /// - The slice length equals the number of columns in the row group - /// - Each element is `Some` if that column has statistics, `None` otherwise - pub fn column_indexes_for_rowgroup( Review Comment: Does this mean we lose the ability to access all column indexes without a method cal for each one? Maybe that is ok given that processing the index is likely to be much more expensive than the method call to retrieve it 🤔 ########## parquet/src/file/metadata/mod.rs: ########## @@ -209,70 +208,31 @@ pub(crate) use writer::ThriftMetadataWriter; /// [Page Index]: https://github.com/apache/parquet-format/blob/master/PageIndex.md /// [`ColumnIndex`]: crate::file::page_index::column_index::ColumnIndexMetaData /// [`OffsetIndex`]: crate::file::page_index::offset_index::OffsetIndexMetaData -#[derive(Debug, Clone, PartialEq)] -pub struct PageIndex { - column_indexes: Option<Vec<Vec<Option<ColumnIndexMetaData>>>>, - offset_indexes: Option<Vec<Vec<Option<OffsetIndexMetaData>>>>, -} - -impl PageIndex { - pub(crate) fn new( - column_indexes: Option<Vec<Vec<Option<ColumnIndexMetaData>>>>, - offset_indexes: Option<Vec<Vec<Option<OffsetIndexMetaData>>>>, - ) -> Self { - Self { - column_indexes, - offset_indexes, - } - } - +pub trait PageIndexProvider: Send + Sync + std::fmt::Debug { /// Returns `true` if offset index structures are present /// /// This indicates whether [`OffsetIndexMetaData`] structures were loaded or created. /// Returns `true` even if some individual columns lack offset indexes. /// /// To check if a specific column has an offset index, use [`Self::offset_index`]. - pub fn has_offset_indexes(&self) -> bool { - self.offset_indexes.is_some() - } + fn has_offset_indexes(&self) -> bool; Review Comment: This is a really nice thing ########## parquet/src/file/metadata/mod.rs: ########## Review Comment: Given the new code / structures, maybe putting them in their own module like `parquet/src/file/metadata/page_index.rs` would keep things better organized ########## parquet/src/file/metadata/mod.rs: ########## @@ -362,18 +283,460 @@ impl PageIndex { /// Returns: /// * `Some(&Vec<PageLocation>)` - Vector of page locations if offset index exists /// * `None` - Offset index not available - pub fn page_locations( + fn page_locations( &self, row_group_idx: usize, column_idx: usize, ) -> Option<&Vec<PageLocation>> { - if let Some(offset_indexes) = self.offset_indexes.as_ref() { - let rg = offset_indexes.get(row_group_idx)?; - let off_idx = rg.get(column_idx)?.as_ref()?; - Some(off_idx.page_locations()) + Some( + self.offset_index(row_group_idx, column_idx)? + .page_locations(), + ) + } + + /// Returns a reference to the trait object as `&dyn Any` for downcasting + /// + /// This allows downcasting to concrete types when needed (e.g., for serialization) + fn as_any(&self) -> &dyn std::any::Any; +} + +/// Provides convenient access to page index data for a specific row group +/// +/// This struct wraps a [`PageIndexProvider`] and automatically applies the row group +/// index, simplifying access to column and offset indexes for a single row group. +/// It is primarily used by readers to avoid repeatedly passing the row group index +/// when accessing page-level metadata. +/// +/// # Example +/// +/// ``` +/// use parquet::file::metadata::ParquetMetaData; +/// # use parquet::errors::Result; +/// +/// fn process_row_group_pages(metadata: &ParquetMetaData, row_group_idx: usize) -> Result<()> { +/// if let Some(page_index) = metadata.page_index() { +/// // Create a row-group-specific view +/// let rg_page_index = parquet::file::metadata::RowGroupPageIndex::new( +/// row_group_idx, +/// metadata.page_index().cloned(), +/// ); +/// +/// // Now access column indexes without specifying row_group_idx each time +/// for col_idx in 0..metadata.file_metadata().schema_descr().num_columns() { +/// if let Some(col_idx_data) = rg_page_index.column_index(col_idx) { +/// println!("Column {} has {} pages", col_idx, col_idx_data.num_pages()); +/// } +/// } +/// } +/// Ok(()) +/// } +/// ``` +#[derive(Debug)] +pub struct RowGroupPageIndex { + row_group_idx: usize, + page_index: Option<Arc<dyn PageIndexProvider>>, +} + +impl RowGroupPageIndex { + /// Creates a new [`RowGroupPageIndex`] for the specified row group + /// + /// # Arguments + /// + /// * `row_group_idx` - The index of the row group within the file + /// * `page_index` - Optional page index provider containing the index data + pub fn new(row_group_idx: usize, page_index: Option<Arc<dyn PageIndexProvider>>) -> Self { + Self { + row_group_idx, + page_index, + } + } + + /// Returns the column index for a specific column in this row group + /// + /// This is a convenience method that wraps [`PageIndexProvider::column_index`], + /// automatically applying the row group index stored in this struct. + /// + /// # Returns + /// + /// * `Some(&ColumnIndexMetaData)` - Column index is available with page-level statistics + /// * `None` - Index unavailable (no page index, column out of bounds, or no statistics) + /// + /// # See Also + /// + /// * [`PageIndexProvider::column_index`] for more details on column indexes + pub fn column_index(&self, column_idx: usize) -> Option<&ColumnIndexMetaData> { + self.page_index + .as_ref()? + .column_index(self.row_group_idx, column_idx) + } + + /// Returns the offset index for a specific column in this row group + /// + /// This is a convenience method that wraps [`PageIndexProvider::offset_index`], + /// automatically applying the row group index stored in this struct. + /// + /// # Returns + /// + /// * `Some(&OffsetIndexMetaData)` - Offset index is available with page locations + /// * `None` - Index unavailable (no page index, column out of bounds) + /// + /// # See Also + /// + /// * [`PageIndexProvider::offset_index`] for more details on offset indexes + pub fn offset_index(&self, column_idx: usize) -> Option<&OffsetIndexMetaData> { + self.page_index + .as_ref()? + .offset_index(self.row_group_idx, column_idx) + } + + /// Returns the physical locations of all data pages for a specific column in this row group + /// + /// This is a convenience method that wraps [`PageIndexProvider::page_locations`], + /// automatically applying the row group index stored in this struct. + /// + /// This enables direct I/O to specific pages without reading the entire column chunk. + /// + /// # Returns + /// + /// * `Some(&Vec<PageLocation>)` - Vector of page locations if offset index exists + /// * `None` - Offset index not available for this column + /// + /// # See Also + /// + /// * [`PageIndexProvider::page_locations`] for more details on page locations + pub fn page_locations(&self, column_idx: usize) -> Option<&Vec<PageLocation>> { + Some(self.offset_index(column_idx)?.page_locations()) + } + + /// Returns the expected number of data pages for a specific column in this row group + /// + /// This count includes only data pages, not dictionary pages or other metadata pages. + /// + /// This is a convenience method that wraps [`PageIndexProvider::num_data_pages`], + /// automatically applying the row group index stored in this struct. + /// + /// # Returns + /// + /// * `Some(usize)` - Number of data pages if any index is available + /// * `None` - No index information available for this column + /// + /// # See Also + /// + /// * [`PageIndexProvider::num_data_pages`] for more details + pub fn num_data_pages(&self, column_idx: usize) -> Option<usize> { + self.page_index + .as_ref()? + .num_data_pages(self.row_group_idx, column_idx) + } +} + +/// Struct to encapsulate the Parquet [Page Index] +/// +/// This struct provides a dense representation of the Page Index. It is +/// used internally by this crate when assembling and writing the Page +/// Index. It is also the default implmentation of the [`PageIndexProvider`] +/// contained in the [`ParquetMetaData`]. +/// +/// # Example: Constructing a synthetic `PageIndex` +/// +/// This example builds a [`ParquetMetaData`] for a file with a single row +/// group containing a single `BYTE_ARRAY` column with one data page, and +/// attaches a matching `PageIndex`, as might be done in tests that +/// exercise page-level statistics handling. +/// +/// ``` +/// # use std::sync::Arc; +/// # use parquet::basic::{BoundaryOrder, Type as PhysicalType}; +/// # use parquet::file::metadata::{ +/// # ColumnChunkMetaData, ColumnIndexBuilder, FileMetaData, OffsetIndexBuilder, +/// # PageIndexBuilder, ParquetMetaData, RowGroupMetaData, +/// # }; +/// # use parquet::schema::types::{SchemaDescriptor, Type}; +/// // Create metadata for a file with a single row group containing a Review Comment: For this example, I recomemnd hiding via the `#` the setup with schema, columns, etc ########## parquet/src/file/metadata/mod.rs: ########## @@ -323,18 +254,8 @@ impl PageIndex { /// Returns: /// * `Some(&OffsetIndexMetaData)` - Offset index is available /// * `None` - Index unavailable (not loaded, row group/column out of bounds) - pub fn offset_index( - &self, - row_group_idx: usize, - column_idx: usize, - ) -> Option<&OffsetIndexMetaData> { - if let Some(offset_indexes) = self.offset_indexes.as_ref() { - let rg = offset_indexes.get(row_group_idx)?; - rg.get(column_idx)?.as_ref() - } else { - None - } - } + fn offset_index(&self, row_group_idx: usize, column_idx: usize) Review Comment: maybe worth pointing here to the new `RowGroupPageIndex` structure in docs to help readers discover it ########## parquet/src/file/metadata/mod.rs: ########## @@ -2293,19 +2700,26 @@ mod tests { .set_row_groups(row_group_meta_with_stats) .build(); + // Base size without page index #[cfg(not(feature = "encryption"))] - let base_expected_size = 2798; + let base_expected_size = 2766; Review Comment: 🎉 it get ssmaller -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
