feat(parquet): RowSelection can be backed by a BooleanBuffer

parquet/src/arrow/arrow_reader/read_plan.rs

@@ -19,8 +19,9 @@
 //! from a Parquet file
 
 use crate::arrow::array_reader::ArrayReader;
-use crate::arrow::arrow_reader::selection::RowSelectionPolicy;
-use crate::arrow::arrow_reader::selection::RowSelectionStrategy;
+use crate::arrow::arrow_reader::selection::{
+    RowSelectionInner, RowSelectionPolicy, RowSelectionStrategy, mask_to_selectors,
+};
 use crate::arrow::arrow_reader::{
     ArrowPredicate, ParquetRecordBatchReader, RowSelection, RowSelectionCursor, RowSelector,
 };
@@ -157,6 +158,10 @@ impl ReadPlanBuilder {
                     None => return RowSelectionStrategy::Selectors,
                 };
 
+                if selection.as_mask().is_some() {
+                    return RowSelectionStrategy::Mask;
+                }
+
                 // total_rows: total number of rows selected / skipped
                 // effective_count: number of non-empty selectors
                 let (total_rows, effective_count) =
@@ -302,19 +307,8 @@ impl ReadPlanBuilder {
             row_selection_policy: _,
         } = self;
 
-        let selection = selection.map(|s| s.trim());
-
         let row_selection_cursor = selection
-            .map(|s| {
-                let trimmed = s.trim();
-                let selectors: Vec<RowSelector> = trimmed.into();
-                match selection_strategy {
-                    RowSelectionStrategy::Mask => {
-                        RowSelectionCursor::new_mask_from_selectors(selectors)
-                    }
-                    RowSelectionStrategy::Selectors => RowSelectionCursor::new_selectors(selectors),
-                }
-            })
+            .map(|s| build_cursor(s.trim(), selection_strategy))
             .unwrap_or(RowSelectionCursor::new_all());
 
         ReadPlan {
@@ -324,6 +318,24 @@ impl ReadPlanBuilder {
     }
 }
 
+/// Lower a [`RowSelection`] to the cursor form requested by the resolved strategy.
+fn build_cursor(selection: RowSelection, strategy: RowSelectionStrategy) -> RowSelectionCursor {
+    match (strategy, selection.into_inner()) {
+        (RowSelectionStrategy::Mask, RowSelectionInner::Mask(mask)) => {
+            RowSelectionCursor::new_mask_from_buffer(mask)
+        }
+        (RowSelectionStrategy::Mask, RowSelectionInner::Selectors(selectors)) => {
+            RowSelectionCursor::new_mask_from_selectors(selectors)
+        }
+        (RowSelectionStrategy::Selectors, RowSelectionInner::Selectors(selectors)) => {
+            RowSelectionCursor::new_selectors(selectors)
+        }
+        (RowSelectionStrategy::Selectors, RowSelectionInner::Mask(mask)) => {
+            RowSelectionCursor::new_selectors(mask_to_selectors(&mask))
+        }
+    }
+}
+
 /// Builder for [`ReadPlan`] that applies a limit and offset to the read plan
 ///
 /// See [`ReadPlanBuilder::limited`] to create this builder.