alamb commented on code in PR #11157:
URL: https://github.com/apache/arrow-rs/pull/11157#discussion_r4144630473
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -774,6 +776,47 @@ impl ArrowReaderOptions {
self
}
+ /// 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
Review Comment:
The interaction between PageIndexPolicy and the ColumnChunkMask is
complicated, especially when there is an existing PageIndex in ParquetMetadata
One potential API we could consider is a new variant in
`PageIndexPolicy::Mask` to avoid having to thread through another set of
variables.
```rust
pub enum PageIndexPolicy {
/// Do not read the page index.
#[default]
Skip,
...
/// Ensure the specified indexes exist, and error if not
Required(ColumnChunkMask),
/// Load the specified indexes if they exist
Optional(ColumnChunkMask),
}
```
However, that would be an API change as `PageIndexPolicy` isn't marked as
`non_exhastive` so we would have to either
1. Wait for the next major release
2. Go with this API, and then unify PageIndexPolicy, and deprecate
`with_column_index_mask` et al methods in next major version 🤔
##########
parquet/src/file/metadata/reader.rs:
##########
@@ -108,6 +117,191 @@ impl From<bool> for PageIndexPolicy {
}
}
+/// 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`]).
+///
+/// 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 {
+ // `None` means all, while `Some(empty)` means none. Store u32 because
+ // Parquet/Thrift collections cannot contain more than i32::MAX entries.
+ row_groups: Option<Arc<[u32]>>,
Review Comment:
I recommend documenting what the `[u32]` means and the invariants. I think
it is something like "sorted indexes"
I was somewhat confused on first read if it was using bitmaps; It is fine
that it isn't, but if we want to get fancy / need more space efficiency here in
the future we could switch to bitmaps
##########
parquet/src/file/metadata/reader.rs:
##########
@@ -108,6 +117,191 @@ impl From<bool> for PageIndexPolicy {
}
}
+/// 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`]).
+///
+/// At present this is only used to select elements of the [Page Index] for
decoding.
Review Comment:
It may also be worth mentioning this is cheap to clone (some Arcs)
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -774,6 +776,47 @@ impl ArrowReaderOptions {
self
}
+ /// 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;
Review Comment:
is it worth making this API return `Result` and checking that the column
index mask doesn't refer to out of bounds columns (e.g. some column index
greater than the number of columns?)
--
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]
