From 87bc874f07695c9f59a983be5ef590f29afdb198 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Sat, 25 Apr 2026 13:32:12 +0800 Subject: [PATCH 1/4] refactor: add a take_n_true to BooleanArray --- arrow-array/src/array/boolean_array.rs | 126 +++++++++++++++++++- parquet/src/arrow/arrow_reader/read_plan.rs | 64 +--------- 2 files changed, 127 insertions(+), 63 deletions(-) diff --git a/arrow-array/src/array/boolean_array.rs b/arrow-array/src/array/boolean_array.rs index 22a1ba7653ac..6f36f7f0221c 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::bit_chunk_iterator::UnalignedBitChunk; @@ -608,6 +608,62 @@ 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(); + if n == 0 { + if !self.has_true() { + return self; + } + return BooleanArray::new(BooleanBuffer::new_unset(len), self.nulls); + } + + // `set_indices` scans 64 bits at a time via `trailing_zeros`, so locating + // the n-th set bit 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 (last_kept, has_more) = match self.nulls.as_ref() { + None => { + let mut iter = self.values.set_indices(); + match iter.nth(n - 1) { + Some(i) => (i, iter.next().is_some()), + None => return self, + } + } + Some(nulls) => { + let mut iter = self.values.set_indices().filter(|&i| nulls.is_valid(i)); + match iter.nth(n - 1) { + Some(i) => (i, iter.next().is_some()), + None => return self, + } + } + }; + if !has_more { + return self; + } + + let mut builder = BooleanBufferBuilder::new(len); + builder.append_buffer(&self.values.slice(0, last_kept + 1)); + builder.append_n(len - last_kept - 1, false); + BooleanArray::new(builder.finish(), self.nulls) + } + /// Deconstruct this array into its constituent parts pub fn into_parts(self) -> (BooleanBuffer, Option) { (self.values, self.nulls) @@ -1630,4 +1686,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 = (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::::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 ac2e105ecf4f..f178e8aae866 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; @@ -248,7 +248,7 @@ impl ReadPlanBuilder { 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); + let truncated = filter.take_n_true(needed); filters.push(truncated); break; } @@ -409,35 +409,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 +472,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 = (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; From 7a3793463c39d95c7fa09e194e518bf8ac4fdeaf Mon Sep 17 00:00:00 2001 From: Huaijin Date: Wed, 1 Jul 2026 09:49:17 +0800 Subject: [PATCH 2/4] apply suggestion --- arrow-array/src/array/boolean_array.rs | 41 ++++-------- parquet/src/arrow/arrow_reader/read_plan.rs | 72 ++++++++++++++++++--- 2 files changed, 75 insertions(+), 38 deletions(-) diff --git a/arrow-array/src/array/boolean_array.rs b/arrow-array/src/array/boolean_array.rs index 08f409160df4..9245f5b33700 100644 --- a/arrow-array/src/array/boolean_array.rs +++ b/arrow-array/src/array/boolean_array.rs @@ -579,40 +579,23 @@ impl BooleanArray { /// ``` pub fn take_n_true(self, n: usize) -> BooleanArray { let len = self.len(); - if n == 0 { - if !self.has_true() { - return self; - } - return BooleanArray::new(BooleanBuffer::new_unset(len), self.nulls); - } - // `set_indices` scans 64 bits at a time via `trailing_zeros`, so locating - // the n-th set bit 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 (last_kept, has_more) = match self.nulls.as_ref() { - None => { - let mut iter = self.values.set_indices(); - match iter.nth(n - 1) { - Some(i) => (i, iter.next().is_some()), - None => return self, - } - } - Some(nulls) => { - let mut iter = self.values.set_indices().filter(|&i| nulls.is_valid(i)); - match iter.nth(n - 1) { - Some(i) => (i, iter.next().is_some()), - None => return self, - } - } + // 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), }; - if !has_more { + let Some(end) = end else { return self; - } + }; let mut builder = BooleanBufferBuilder::new(len); - builder.append_buffer(&self.values.slice(0, last_kept + 1)); - builder.append_n(len - last_kept - 1, false); + builder.append_buffer(&self.values.slice(0, end)); + builder.append_n(len - end, false); BooleanArray::new(builder.finish(), self.nulls) } diff --git a/parquet/src/arrow/arrow_reader/read_plan.rs b/parquet/src/arrow/arrow_reader/read_plan.rs index f348b42090f6..b1dd146995e9 100644 --- a/parquet/src/arrow/arrow_reader/read_plan.rs +++ b/parquet/src/arrow/arrow_reader/read_plan.rs @@ -238,22 +238,21 @@ impl ReadPlanBuilder { filter.len() )); } - let filter = match filter.null_count() { - 0 => filter, - _ => 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 = filter.take_n_true(needed); + Some(limit) if limit - matched_rows <= filter.len() => { + let truncated = filter.take_n_true(limit - matched_rows); + matched_rows += truncated.true_count(); + let truncated = prep_filter_for_row_selection(truncated); filters.push(truncated); - break; + if matched_rows >= limit { + break; + } } _ => { matched_rows += filter.true_count(); + let filter = prep_filter_for_row_selection(filter); filters.push(filter); } } @@ -409,6 +408,13 @@ impl LimitedReadPlanBuilder { } } +fn prep_filter_for_row_selection(filter: BooleanArray) -> BooleanArray { + match filter.null_count() { + 0 => filter, + _ => prep_null_mask_filter(&filter), + } +} + /// A plan reading specific rows from a Parquet Row Group. /// /// See [`ReadPlanBuilder`] to create `ReadPlan`s @@ -523,4 +529,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 = (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::()) + }); + + 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); + } } From 32c4d1c39ee47c29675211bc8320cf5598f0a0fd Mon Sep 17 00:00:00 2001 From: Huaijin Date: Wed, 1 Jul 2026 09:59:29 +0800 Subject: [PATCH 3/4] update --- parquet/src/arrow/arrow_reader/read_plan.rs | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/parquet/src/arrow/arrow_reader/read_plan.rs b/parquet/src/arrow/arrow_reader/read_plan.rs index b1dd146995e9..2552191a947f 100644 --- a/parquet/src/arrow/arrow_reader/read_plan.rs +++ b/parquet/src/arrow/arrow_reader/read_plan.rs @@ -238,13 +238,17 @@ impl ReadPlanBuilder { filter.len() )); } + let filter = match filter.null_count() { + 0 => filter, + _ => prep_null_mask_filter(&filter), + }; + processed_rows += input_rows; match limit { Some(limit) if limit - matched_rows <= filter.len() => { let truncated = filter.take_n_true(limit - matched_rows); matched_rows += truncated.true_count(); - let truncated = prep_filter_for_row_selection(truncated); filters.push(truncated); if matched_rows >= limit { break; @@ -252,7 +256,6 @@ impl ReadPlanBuilder { } _ => { matched_rows += filter.true_count(); - let filter = prep_filter_for_row_selection(filter); filters.push(filter); } } @@ -408,13 +411,6 @@ impl LimitedReadPlanBuilder { } } -fn prep_filter_for_row_selection(filter: BooleanArray) -> BooleanArray { - match filter.null_count() { - 0 => filter, - _ => prep_null_mask_filter(&filter), - } -} - /// A plan reading specific rows from a Parquet Row Group. /// /// See [`ReadPlanBuilder`] to create `ReadPlan`s From 20b3d4b6447bd6497b9cc64a4a45d90488d36af4 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Wed, 1 Jul 2026 22:12:44 +0800 Subject: [PATCH 4/4] added comment --- parquet/src/arrow/arrow_reader/read_plan.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/parquet/src/arrow/arrow_reader/read_plan.rs b/parquet/src/arrow/arrow_reader/read_plan.rs index 2552191a947f..71b64802b0c4 100644 --- a/parquet/src/arrow/arrow_reader/read_plan.rs +++ b/parquet/src/arrow/arrow_reader/read_plan.rs @@ -240,6 +240,8 @@ 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), };