This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new b98f76beab feat(parquet): expose RowSelection::force_mask (#11337)
b98f76beab is described below
commit b98f76beab88301fecb00afe1e2cbc306befed79
Author: Huaijin <[email protected]>
AuthorDate: Sun Oct 4 14:15:40 2026 +0000
feat(parquet): expose RowSelection::force_mask (#11337)
# Which issue does this PR close?
Follow-up to the suggestion in
https://github.com/apache/datafusion/pull/25731#discussion_r4158690507.
# Rationale for this change
Engines composing row selections need to convert selector-backed
selections to masks to keep an intersection or union bitmap-backed.
Parquet already implements this conversion internally, but its public
API requires callers such as DataFusion to rebuild the bitmap
themselves.
# What changes are included in this PR?
Add `RowSelection::force_mask(self) -> Self`, reusing the existing
bitmap conversion and preserving the full row span, including trailing
skips. Selections that are already mask-backed are returned unchanged.
The internal mask materialization policy now calls this method.
# Are these changes tested?
Unit tests cover empty, fully selected, fully skipped, and mixed
selector sequences; reuse of an existing mask with a non-byte-aligned
offset; repeated conversion; and preserving bitmap backing and selected
rows through intersection. A runnable documentation example demonstrates
the public API.
# Are there any significant user-facing changes?
No breaking changes.
---
parquet/src/arrow/arrow_reader/selection/cursor.rs | 4 +-
parquet/src/arrow/arrow_reader/selection/mod.rs | 98 ++++++++++++++++++++++
2 files changed, 99 insertions(+), 3 deletions(-)
diff --git a/parquet/src/arrow/arrow_reader/selection/cursor.rs
b/parquet/src/arrow/arrow_reader/selection/cursor.rs
index 2be8b0d48f..f78553028e 100644
--- a/parquet/src/arrow/arrow_reader/selection/cursor.rs
+++ b/parquet/src/arrow/arrow_reader/selection/cursor.rs
@@ -72,9 +72,7 @@ impl RowSelectionPolicy {
(RowSelectionStrategy::Mask, true) |
(RowSelectionStrategy::Selectors, false) => {
selection
}
- (RowSelectionStrategy::Mask, false) => {
-
RowSelection::from_boolean_buffer(selection.into_boolean_buffer())
- }
+ (RowSelectionStrategy::Mask, false) => selection.force_mask(),
(RowSelectionStrategy::Selectors, true) => {
RowSelection::from(Vec::<RowSelector>::from(selection))
}
diff --git a/parquet/src/arrow/arrow_reader/selection/mod.rs
b/parquet/src/arrow/arrow_reader/selection/mod.rs
index df9c4cde08..6ce7d71f9f 100644
--- a/parquet/src/arrow/arrow_reader/selection/mod.rs
+++ b/parquet/src/arrow/arrow_reader/selection/mod.rs
@@ -234,6 +234,37 @@ impl RowSelection {
}
}
+ /// Consume this selection and return a mask-backed selection.
+ ///
+ /// If this selection is already mask-backed, it is returned unchanged.
+ /// Otherwise, its selectors are converted to a bitmap, preserving all
+ /// selected and skipped rows, including trailing skips.
+ ///
+ /// This can be used to keep the result of [`Self::intersection`] or
+ /// [`Self::union`] mask-backed when the other selection is mask-backed.
+ ///
+ /// # Example
+ /// ```
+ /// use parquet::arrow::arrow_reader::{RowSelection, RowSelector};
+ ///
+ /// let selection = RowSelection::from(vec![
+ /// RowSelector::skip(2),
+ /// RowSelector::select(3),
+ /// RowSelector::skip(1),
+ /// ]).force_mask();
+ ///
+ /// assert!(selection.as_mask().is_some());
+ /// assert_eq!(selection.total_row_count(), 6);
+ /// assert_eq!(selection.row_count(), 3);
+ /// ```
+ pub fn force_mask(self) -> Self {
+ if self.as_mask().is_some() {
+ self
+ } else {
+ Self::from_boolean_buffer(self.into_boolean_buffer())
+ }
+ }
+
/// Consume the selection and return its internal storage.
pub(crate) fn into_inner(self) -> RowSelectionInner {
self.inner
@@ -1131,6 +1162,73 @@ mod tests {
filters
}
+ #[test]
+ fn test_force_mask_from_selectors() {
+ let cases = [
+ (vec![], vec![]),
+ (vec![RowSelector::select(9)], vec![true; 9]),
+ (vec![RowSelector::skip(9)], vec![false; 9]),
+ (
+ vec![
+ RowSelector::skip(3),
+ RowSelector::select(7),
+ RowSelector::skip(2),
+ RowSelector::select(1),
+ RowSelector::skip(4),
+ ],
+ vec![
+ false, false, false, true, true, true, true, true, true,
true, false, false,
+ true, false, false, false, false,
+ ],
+ ),
+ ];
+
+ for (selectors, bits) in cases {
+ let original = RowSelection::from(selectors);
+ let selection = original.clone().force_mask();
+ assert_eq!(selection.as_mask().unwrap(),
&BooleanBuffer::from(bits));
+ assert_eq!(selection.total_row_count(),
original.total_row_count());
+ assert_eq!(selection.row_count(), original.row_count());
+ assert_eq!(selection.skipped_row_count(),
original.skipped_row_count());
+ assert_eq!(selection, original);
+ }
+ }
+
+ #[test]
+ fn test_force_mask_preserves_existing_mask() {
+ let mask = BooleanBuffer::from(vec![
+ false, true, false, true, true, false, true, false, false, true,
false,
+ ])
+ .slice(3, 7);
+ let selection =
RowSelection::from_boolean_buffer(mask.clone()).force_mask();
+ assert!(selection.as_mask().unwrap().ptr_eq(&mask));
+ assert_eq!(selection.row_count(), 4);
+
+ // Repeated conversion also preserves the buffer and its bit offset.
+ let selection = selection.force_mask();
+ assert!(selection.as_mask().unwrap().ptr_eq(&mask));
+ assert_eq!(selection.total_row_count(), 7);
+ }
+
+ #[test]
+ fn test_force_mask_preserves_bitmap_intersection() {
+ let existing =
RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![
+ true, false, true, true, false, true, true, false, true,
+ ]));
+ let selection = RowSelection::from(vec![
+ RowSelector::skip(2),
+ RowSelector::select(5),
+ RowSelector::skip(2),
+ ]);
+ let expected = existing.intersection(&selection);
+ let actual = existing.intersection(&selection.force_mask());
+
+ assert!(actual.as_mask().is_some());
+ assert_eq!(actual, expected);
+ assert_eq!(actual.total_row_count(), 9);
+ assert_eq!(actual.row_count(), 4);
+ }
+
#[test]
fn test_total_row_count() {
let selection = RowSelection::from(vec![RowSelector::skip(5),
RowSelector::select(3)]);
