etseidl commented on code in PR #11157:
URL: https://github.com/apache/arrow-rs/pull/11157#discussion_r4088324552
##########
parquet/src/file/metadata/reader.rs:
##########
@@ -104,6 +107,114 @@ impl From<bool> for PageIndexPolicy {
}
}
+/// Struct to specify column chunks for which metadata is required.
+///
+/// Column chunks are identified by row group index and column index. 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, Default)]
+pub struct ColumnChunkMask {
+ // using i32 because that's how thrift vectors are sized
+ row_groups: Option<Arc<BTreeSet<i32>>>,
+ columns: Option<Arc<BTreeSet<i32>>>,
+}
+
+impl ColumnChunkMask {
+ /// Select all row groups and columns.
+ pub fn all() -> Self {
+ Self::default()
+ }
+
+ /// Select only the listed columns.
+ ///
+ /// Any indices in `columns` that are less than zero will be ignored.
Passing an empty
+ /// set is treated the same as selecting all columns.
+ pub fn columns(columns: impl IntoIterator<Item = i32>) -> Self {
+ Self {
+ row_groups: None,
+ columns: Self::iter_to_set(columns),
+ }
+ }
+
+ /// Select only the listed row groups.
+ ///
+ /// Any indices in `row_groups` that are less than zero will be ignored.
Passing an empty
+ /// set is treated the same as selecting all row groups.
+ pub fn row_groups(row_groups: impl IntoIterator<Item = i32>) -> Self {
+ Self {
+ row_groups: Self::iter_to_set(row_groups),
+ columns: None,
+ }
+ }
+
+ /// Select only the listed row groups and columns.
+ ///
+ /// Any indices in `row_groups` or `columns` that are less than zero will
be ignored.
+ /// Passing an empty set for `row_groups` is treated as selecting all row
groups, and
+ /// an empty set for `columns` as selectiong all columns.
+ pub fn row_groups_and_columns(
+ row_groups: impl IntoIterator<Item = i32>,
+ columns: impl IntoIterator<Item = i32>,
+ ) -> 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.
+ ///
+ /// Returns `false` if `idx > `[`i32::MAX`].
+ pub fn includes_row_group(&self, idx: usize) -> bool {
+ let Ok(idx) = i32::try_from(idx) else {
+ return false;
+ };
+ self.row_groups
+ .as_ref()
+ .is_none_or(|keep| keep.contains(&idx))
+ }
+
+ /// Test if `idx` is in the column set.
+ ///
+ /// Returns `false` if `idx > `[`i32::MAX`].
+ pub fn includes_column(&self, idx: usize) -> bool {
+ let Ok(idx) = i32::try_from(idx) else {
+ return false;
+ };
+ self.columns.as_ref().is_none_or(|keep| keep.contains(&idx))
+ }
+
+ fn iter_to_set(indices: impl IntoIterator<Item = i32>) ->
Option<Arc<BTreeSet<i32>>> {
+ let set: BTreeSet<i32> = indices.into_iter().filter(|&i| i >=
0).collect();
+ (!set.is_empty()).then_some(Arc::new(set))
+ }
Review Comment:
I'm fine with revisiting this. I think I was just trying to avoid users
inadvertently specifying empty sets, but I agree that empty (as opposed to
`None`) should mean just that.
--
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]
