etseidl commented on code in PR #11157:
URL: https://github.com/apache/arrow-rs/pull/11157#discussion_r4139448659
##########
parquet/src/file/metadata/reader.rs:
##########
@@ -108,6 +115,166 @@ 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]>>,
+ columns: Option<Arc<[u32]>>,
+}
+
+impl ColumnChunkMask {
+ /// Select all row groups and columns.
+ pub fn all() -> Self {
+ Self::default()
+ }
+
+ /// Select no row groups or columns.
+ pub fn none() -> Self {
+ Self {
+ row_groups: Some(Arc::from([])),
+ columns: Some(Arc::from([])),
+ }
+ }
+
+ /// Select only the listed columns.
+ ///
+ /// Passing an empty iterator selects no columns.
+ pub fn columns(columns: impl IntoIterator<Item = usize>) -> Self {
+ Self {
+ row_groups: None,
+ columns: Self::iter_to_set(columns),
+ }
+ }
+
+ /// Select only the listed row groups.
+ ///
+ /// Passing an empty iterator selects no row groups.
+ pub fn row_groups(row_groups: impl IntoIterator<Item = usize>) -> Self {
+ Self {
+ row_groups: Self::iter_to_set(row_groups),
+ columns: None,
+ }
+ }
+
+ /// Select only the listed row groups and columns.
+ ///
+ /// An empty iterator for either dimension selects no column chunks.
+ pub fn row_groups_and_columns(
+ row_groups: impl IntoIterator<Item = usize>,
+ columns: impl IntoIterator<Item = usize>,
+ ) -> Self {
+ Self {
+ row_groups: Self::iter_to_set(row_groups),
+ columns: Self::iter_to_set(columns),
+ }
+ }
+
+ /// Test if `idx` is in the row group set.
+ pub fn includes_row_group(&self, idx: usize) -> bool {
+ Self::includes_index(self.row_groups.as_ref(), idx)
+ }
+
+ /// Test if `idx` is in the column set.
+ pub fn includes_column(&self, idx: usize) -> bool {
+ Self::includes_index(self.columns.as_ref(), idx)
+ }
+
+ fn includes_index(keep: Option<&Arc<[u32]>>, idx: usize) -> bool {
+ // return false for out-of-bounds index
+ let Ok(idx) = u32::try_from(idx) else {
+ return false;
+ };
+ keep.is_none_or(|keep| keep.binary_search(&idx).is_ok())
+ }
+
+ /// Returns `true` when this mask selects every column chunk.
+ pub fn is_all(&self) -> bool {
+ self.row_groups.is_none() && self.columns.is_none()
+ }
+
+ /// Returns selected row groups, or `None` when all row groups are
selected.
+ pub fn selected_row_groups(&self) -> Option<&[u32]> {
+ self.row_groups.as_deref()
+ }
+
+ /// Returns selected leaf columns, or `None` when all columns are selected.
+ pub fn selected_columns(&self) -> Option<&[u32]> {
+ self.columns.as_deref()
+ }
+
+ /// Creates a mask selecting the leaf columns in an Arrow projection.
+ #[cfg(feature = "arrow")]
+ pub fn from_projection(projection: &ProjectionMask, schema:
&SchemaDescriptor) -> Self {
+ Self::columns((0..schema.num_columns()).filter(|&i|
projection.leaf_included(i)))
+ }
+
+ /// Returns an iterator over the row group indices selected by this mask
+ pub fn row_group_indices(&self, num_row_groups: usize) -> Box<dyn
Iterator<Item = usize> + '_> {
+ Self::axis_indices(self.row_groups.as_deref(), num_row_groups)
+ }
+
+ /// Returns an iterator over the column indices selected by this mask
+ pub fn column_indices(&self, num_columns: usize) -> Box<dyn Iterator<Item
= usize> + '_> {
+ Self::axis_indices(self.columns.as_deref(), num_columns)
+ }
+
+ fn axis_indices(axis: Option<&[u32]>, len: usize) -> Box<dyn Iterator<Item
= usize> + '_> {
Review Comment:
I had Codex whip up some tests to show the difference. Added a concrete
iterator in
https://github.com/apache/arrow-rs/pull/11157/commits/ad865184537abeab2463330e0cc26bd39905072d
--
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]
