diff --git a/arrow-array/src/array/boolean_array.rs b/arrow-array/src/array/boolean_array.rs index 0fa5b5d91209..9245f5b33700 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) { (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 = (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 674ae2b8d964..71b64802b0c4 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 = (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 = (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); + } }