etseidl commented on code in PR #10842:
URL: https://github.com/apache/arrow-rs/pull/10842#discussion_r3867017016
##########
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:
The idea was to not bother allocating the column index if all you want to do
is build the offset index, but yeah, maybe this is overkill since I later
modified `build` to set empty indexes to `None`.
--
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]
