alamb commented on code in PR #10784:
URL: https://github.com/apache/arrow-rs/pull/10784#discussion_r3847489114
##########
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:
maybe `copy_column_index` would be a more specific name 🤔
##########
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>>>>,
Review Comment:
Having some outlet for different representations (so we don't have to hard
code one choice into the base crate that downstream projects have to accept)
sounds like a good idea to me
That way we could implement some fancy Grid structure for DataFusion for
example, but leave a basic implementation in arrow-rs that is simple
##########
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>>>>,
Review Comment:
Yeah the more I think about it, the more I like the idea of separating out
the API for accessing the page index from the concrete implementation.
The builder needs to be part of the concrete implementation, but it would be
really neat if we could switch out PageIndex with `Arc<dyn PageIndexSource>` or
something 🤔 that would let downstream users provide their own imlementation
Something like this perhaps
```rust
enum PageIndexSource {
// concrete pageIndex (our old `Option<Vec<Option<Vec<..>>>>` friend
PageIndex(PageIndex),
// user supplied call back
Callback(Arc<dyn PageIndexSource>),
}
// and then wire it into the ParquetMetadata 🤔
struct ParquetMetadata { '
...
page_index: PageIndexSource,
}
```
--
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]
