This is an automated email from the ASF dual-hosted git repository.
scovich 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 d8d3fa32c1 refactor: add `take_n_true` method to `BooleanArray` (#9823)
d8d3fa32c1 is described below
commit d8d3fa32c107ecae6c77c6f1db7700e2c7eb60ef
Author: Huaijin <[email protected]>
AuthorDate: Wed Jul 1 22:33:23 2026 +0800
refactor: add `take_n_true` method to `BooleanArray` (#9823)
# Which issue does this PR close?
- Closes #9803
# Rationale for this change
- see #9803
# What changes are included in this PR?
move `truncate_filter_after_n_trues` in read_plan.rs to BooleanArray
`take_n_true` method
# Are these changes tested?
yes, already have test case
# Are there any user-facing changes?
---
arrow-array/src/array/boolean_array.rs | 109 ++++++++++++++++++++++++-
parquet/src/arrow/arrow_reader/read_plan.rs | 122 +++++++++++++---------------
2 files changed, 165 insertions(+), 66 deletions(-)
diff --git a/arrow-array/src/array/boolean_array.rs
b/arrow-array/src/array/boolean_array.rs
index 0fa5b5d912..9245f5b337 100644
--- a/arrow-array/src/array/boolean_array.rs
+++ b/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);
+ }
}
diff --git a/parquet/src/arrow/arrow_reader/read_plan.rs
b/parquet/src/arrow/arrow_reader/read_plan.rs
index 674ae2b8d9..71b64802b0 100644
--- a/parquet/src/arrow/arrow_reader/read_plan.rs
+++ b/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);
+ }
}
