This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11337-e97b1aa7fa992b544ef7a5a0c98cbc3791a89776 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
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)]);
