refactor: add take_n_true method to BooleanArray

arrow-array/src/array/boolean_array.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use crate::array::print_long_array;
-use crate::builder::BooleanBuilder;
+use crate::builder::{BooleanBufferBuilder, BooleanBuilder};
 use crate::iterator::BooleanIter;
 use crate::{Array, ArrayAccessor, ArrayRef, Scalar};
 use arrow_buffer::{BooleanBuffer, Buffer, MutableBuffer, NullBuffer, bit_util};
@@ -560,6 +560,45 @@ impl BooleanArray {
         }
     }
 
+    /// Returns a new [`BooleanArray`] of the same length where only the first
+    /// `n` non-null `true` positions remain `true`; any `true` positions
+    /// beyond the first `n` are replaced with `false`. The null buffer is
+    /// preserved unchanged.
+    ///
+    /// If this array has at most `n` non-null `true` values, `self` is
+    /// returned unchanged.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use arrow_array::BooleanArray;
+    /// let a = BooleanArray::from(vec![true, false, true, true, false, true]);
+    /// // Keep only the first 2 `true` positions; later trues become false.
+    /// let r = a.take_n_true(2);
+    /// assert_eq!(r, BooleanArray::from(vec![true, false, true, false, false, false]));
+    /// ```
+    pub fn take_n_true(self, n: usize) -> BooleanArray {
+        let len = self.len();
+        // `set_indices` scans 64 bits at a time via `trailing_zeros`, so locating
+        // the first set bit beyond the retained prefix is cheaper than visiting
+        // every bit. When a null buffer is present, skip set bits whose
+        // corresponding entry is null so only non-null trues count toward `n`
+        // (matching `true_count` semantics).
+        let mut iter = self.values.set_indices();
+        let end = match self.nulls.as_ref() {
+            Some(nulls) => iter.filter(|&i| nulls.is_valid(i)).nth(n),
+            None => iter.nth(n),
+        };
+        let Some(end) = end else {
+            return self;
+        };
+
+        let mut builder = BooleanBufferBuilder::new(len);
+        builder.append_buffer(&self.values.slice(0, end));
+        builder.append_n(len - end, false);
+        BooleanArray::new(builder.finish(), self.nulls)
+    }
+
     /// Deconstruct this array into its constituent parts
     pub fn into_parts(self) -> (BooleanBuffer, Option<NullBuffer>) {
         (self.values, self.nulls)
@@ -1582,4 +1621,72 @@ mod tests {
         let result = left.bitwise_bin_op_mut_or_clone(&right, |a, b| a & b);
         assert_eq!(result, expected);
     }
+
+    #[test]
+    fn test_take_n_true_keeps_first_n_matches() {
+        let a = BooleanArray::from(vec![true, false, true, true, false, true, true]);
+        // true positions: 0, 2, 3, 5, 6
+        let r = a.clone().take_n_true(3);
+        assert_eq!(r.len(), a.len());
+        assert_eq!(r.true_count(), 3);
+        let out: Vec<bool> = (0..r.len()).map(|i| r.value(i)).collect();
+        assert_eq!(
+            out,
+            vec![true, false, true, true, false, false, false],
+            "first three trues should survive, the rest become false"
+        );
+    }
+
+    #[test]
+    fn test_take_n_true_passes_through_when_already_small_enough() {
+        let a = BooleanArray::from(vec![true, false, true, false]);
+        let r = a.clone().take_n_true(5);
+        assert_eq!(r.len(), a.len());
+        assert_eq!(r.true_count(), 2);
+        assert_eq!(r, a);
+    }
+
+    #[test]
+    fn test_take_n_true_zero_returns_all_false() {
+        let a = BooleanArray::from(vec![true, true, true]);
+        let r = a.take_n_true(0);
+        assert_eq!(r.len(), 3);
+        assert_eq!(r.true_count(), 0);
+    }
+
+    #[test]
+    fn test_take_n_true_preserves_nulls_and_skips_them() {
+        // Non-null trues: positions 0, 3, 5. Null at 2 must not count toward `n`.
+        let a = BooleanArray::from(vec![
+            Some(true),
+            Some(false),
+            None,
+            Some(true),
+            Some(false),
+            Some(true),
+        ]);
+        assert_eq!(a.true_count(), 3);
+        let len = a.len();
+
+        let r = a.take_n_true(2);
+        assert_eq!(r.len(), len);
+        assert_eq!(r.true_count(), 2);
+        // Null buffer is preserved unchanged.
+        assert_eq!(r.null_count(), 1);
+        assert!(r.is_null(2));
+        // First two non-null trues kept; the third (position 5) becomes false.
+        assert!(r.value(0));
+        assert!(!r.value(1));
+        assert!(r.value(3));
+        assert!(!r.value(4));
+        assert!(!r.value(5));
+    }
+
+    #[test]
+    fn test_take_n_true_empty_array() {
+        let a = BooleanArray::from(Vec::<bool>::new());
+        let r = a.take_n_true(5);
+        assert_eq!(r.len(), 0);
+        assert_eq!(r.true_count(), 0);
+    }
 }

parquet/src/arrow/arrow_reader/read_plan.rs

@@ -26,7 +26,7 @@ use crate::arrow::arrow_reader::{
 };
 use crate::errors::{ParquetError, Result};
 use arrow_array::{Array, BooleanArray};
-use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder};
+use arrow_buffer::BooleanBuffer;
 use arrow_select::filter::prep_null_mask_filter;
 use std::collections::VecDeque;
 
@@ -240,17 +240,21 @@ impl ReadPlanBuilder {
             }
             let filter = match filter.null_count() {
                 0 => filter,
+                // RowSelection::from_filters expects non-null filters. Convert
+                // NULL predicate results to false so they are not selected.
                 _ => prep_null_mask_filter(&filter),
             };
 
             processed_rows += input_rows;
 
             match limit {
-                Some(limit) if matched_rows + filter.true_count() >= limit => {
-                    let needed = limit - matched_rows;
-                    let truncated = truncate_filter_after_n_trues(filter, needed);
+                Some(limit) if limit - matched_rows <= filter.len() => {
+                    let truncated = filter.take_n_true(limit - matched_rows);
+                    matched_rows += truncated.true_count();
                     filters.push(truncated);
-                    break;
+                    if matched_rows >= limit {
+                        break;
+                    }
                 }
                 _ => {
                     matched_rows += filter.true_count();
@@ -409,35 +413,6 @@ impl LimitedReadPlanBuilder {
     }
 }
 
-/// Produce a new `BooleanArray` of the same length as `filter` in which only
-/// the first `n` `true` positions from `filter` remain `true`; any `true`
-/// positions beyond the first `n` are replaced with `false`.
-///
-/// `filter` must not contain nulls (callers apply [`prep_null_mask_filter`]
-/// first). If `filter` has at most `n` `true` values, a clone is returned.
-fn truncate_filter_after_n_trues(filter: BooleanArray, n: usize) -> BooleanArray {
-    if filter.true_count() <= n {
-        return filter;
-    }
-    let len = filter.len();
-    if n == 0 {
-        return BooleanArray::new(BooleanBuffer::new_unset(len), None);
-    }
-    // `set_indices` scans 64 bits at a time via `trailing_zeros`, so locating
-    // the `n`-th set bit is cheaper than visiting every bit. Everything up to
-    // and including that position is copied verbatim; the rest is zeroed.
-    let values = filter.values();
-    let last_kept = values
-        .set_indices()
-        .nth(n - 1)
-        .expect("n - 1 < true_count, checked above");
-
-    let mut builder = BooleanBufferBuilder::new(len);
-    builder.append_buffer(&values.slice(0, last_kept + 1));
-    builder.append_n(len - last_kept - 1, false);
-    BooleanArray::new(builder.finish(), None)
-}
-
 /// A plan reading specific rows from a Parquet Row Group.
 ///
 /// See [`ReadPlanBuilder`] to create `ReadPlan`s
@@ -501,37 +476,6 @@ mod tests {
         );
     }
 
-    #[test]
-    fn truncate_filter_after_n_trues_keeps_first_n_matches() {
-        let f = BooleanArray::from(vec![true, false, true, true, false, true, true]);
-        // true positions: 0, 2, 3, 5, 6
-        let t = truncate_filter_after_n_trues(f.clone(), 3);
-        assert_eq!(t.len(), f.len());
-        assert_eq!(t.true_count(), 3);
-        let out: Vec<bool> = (0..t.len()).map(|i| t.value(i)).collect();
-        assert_eq!(
-            out,
-            vec![true, false, true, true, false, false, false],
-            "first three trues should survive, the rest become false"
-        );
-    }
-
-    #[test]
-    fn truncate_filter_after_n_trues_passes_through_when_already_small_enough() {
-        let f = BooleanArray::from(vec![true, false, true, false]);
-        let t = truncate_filter_after_n_trues(f.clone(), 5);
-        assert_eq!(t.len(), f.len());
-        assert_eq!(t.true_count(), 2);
-    }
-
-    #[test]
-    fn truncate_filter_after_n_trues_zero_returns_all_false() {
-        let f = BooleanArray::from(vec![true, true, true]);
-        let t = truncate_filter_after_n_trues(f, 0);
-        assert_eq!(t.len(), 3);
-        assert_eq!(t.true_count(), 0);
-    }
-
     #[test]
     fn with_predicate_options_limit_pads_tail_when_no_prior_selection() {
         use crate::arrow::ProjectionMask;
@@ -583,4 +527,52 @@ mod tests {
             "selection must span the full row group, not only the prefix evaluated before the limit"
         );
     }
+
+    #[test]
+    fn with_predicate_options_limit_handles_null_filters() {
+        use crate::arrow::ProjectionMask;
+        use crate::arrow::array_reader::StructArrayReader;
+        use crate::arrow::array_reader::test_util::make_int32_page_reader;
+        use crate::arrow::arrow_reader::ArrowPredicateFn;
+        use arrow_schema::{DataType as ArrowType, Field, Fields};
+
+        const TOTAL_ROWS: usize = 100;
+        const LIMIT: usize = 10;
+
+        let data: Vec<i32> = (0..TOTAL_ROWS as i32).collect();
+        let levels = vec![0; TOTAL_ROWS];
+        let leaf = make_int32_page_reader(&data, &levels, &levels, 0, 0);
+        let struct_type = ArrowType::Struct(Fields::from(vec![Field::new(
+            "c0",
+            ArrowType::Int32,
+            false,
+        )]));
+        let struct_reader = StructArrayReader::new(struct_type, vec![leaf], 0, 0, false);
+
+        let mut predicate = ArrowPredicateFn::new(ProjectionMask::all(), |batch| {
+            Ok((0..batch.num_rows())
+                .map(|i| match i % 4 {
+                    0 | 2 => Some(true),
+                    1 => None,
+                    _ => Some(false),
+                })
+                .collect::<BooleanArray>())
+        });
+
+        let builder = ReadPlanBuilder::new(16)
+            .with_predicate_options(
+                PredicateOptions::new(Box::new(struct_reader), &mut predicate)
+                    .with_limit(LIMIT, TOTAL_ROWS),
+            )
+            .unwrap();
+
+        let selection = builder
+            .selection()
+            .expect("limit-driven early break must produce a selection");
+
+        assert_eq!(selection.row_count(), LIMIT);
+
+        let total: usize = selection.iter().map(|s| s.row_count).sum();
+        assert_eq!(total, TOTAL_ROWS);
+    }
 }