etseidl commented on code in PR #10784:
URL: https://github.com/apache/arrow-rs/pull/10784#discussion_r3847665240
##########
parquet/src/file/metadata/mod.rs:
##########
@@ -412,6 +412,176 @@ impl PageIndex {
None
}
}
+
+ /// Convert this `PageIndex` into a [`PageIndexBuilder`]
+ pub fn into_builder(self) -> PageIndexBuilder {
+ self.into()
+ }
+}
+
+/// 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`
+ /// - [`PageIndexPolicy::Optional`] or [`PageIndexPolicy::Required`]:
Allocates empty index structure
+ ///
+ /// This is more memory-efficient than [`new`](Self::new) when only one
type of index is needed.
+ ///
+ /// # Arguments
+ /// * `num_row_groups` - Number of row groups in the file
+ /// * `num_columns` - Number of columns in the schema
+ /// * `column_index_policy` - Policy for column index allocation
+ /// * `offset_index_policy` - Policy for offset index allocation
+ pub fn new_with_policy(
+ num_row_groups: usize,
+ num_columns: usize,
+ column_index_policy: &PageIndexPolicy,
+ offset_index_policy: &PageIndexPolicy,
+ ) -> Self {
+ use reader::PageIndexPolicy;
+
+ let column_indexes = match column_index_policy {
+ PageIndexPolicy::Skip => None,
+ _ => Self::empty_index(num_row_groups, num_columns),
+ };
+
+ let offset_indexes = match offset_index_policy {
+ PageIndexPolicy::Skip => None,
+ _ => Self::empty_index(num_row_groups, num_columns),
+ };
+
+ Self {
+ column_indexes,
+ offset_indexes,
+ }
+ }
+
+ /// Creates a new [`PageIndexBuilder`] from an existing [`PageIndex`]
+ ///
+ /// This takes ownership of the index structures from the provided
[`PageIndex`],
+ /// allowing them to be modified and rebuilt. Useful for updating existing
page indexes.
+ pub(crate) fn new_from(page_index: PageIndex) -> Self {
+ Self {
+ column_indexes: page_index.column_indexes,
+ offset_indexes: page_index.offset_indexes,
+ }
+ }
+
+ /// Sets the column index for a specific row group and column
Review Comment:
I thought of `add_column_index`, but that sounds more like appending than
doing random access. I don't think `copy` is appropriate since this takes
ownership of the passed index.
--
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]
