This is an automated email from the ASF dual-hosted git repository.
alamb 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 3e03c57afd perf: Add APIs for filter with known selected-row count
(#11206)
3e03c57afd is described below
commit 3e03c57afd96232ccb88a73511849fedfeaf43f4
Author: Neil Conway <[email protected]>
AuthorDate: Tue Sep 29 16:29:55 2026 -0400
perf: Add APIs for filter with known selected-row count (#11206)
# Which issue does this PR close?
- Closes #11205.
# Rationale for this change
`filter`, `filter_record_batch`,
`BatchCoalescer::push_batch_with_filter`, and `FilterBuilder` start by
calling `BooleanArray::true_count` on the caller-provided mask. However,
some callers may already know the number of rows selected by the mask --
for example, because they built the mask one row at a time, derived it
from another data structure (e.g., a validity bitmap), or needed to
compute it for some other reason. In those situations, calling
`true_count` requires a redundant full pass over the bitmap.
Some examples from DataFusion where the count is known before the kernel
runs:
- `FilterExec` with a `LIMIT` counts the mask to track how many rows
remain, then `push_batch_with_filter` counts it again.
- `CASE` evaluation filters the remaining rows for each branch; the
count of each new mask is the previous predicate's `count()` subtracted
from the remaining row count.
- `AND`/`OR` short-circuit evaluation counts the left side to decide
whether to pre-select rows, then drops the count and filters the batch.
- Aggregate `FILTER` clauses count the selection to decide whether
filtering is worthwhile, then filter with the same mask.
# What changes are included in this PR?
- Add `FilterBuilder::with_count`, a `pub unsafe fn` that supplies the
number of rows the mask selects so that `build` does not count them. To
make this possible, `FilterBuilder` now counts the mask, picks the
iteration strategy, and applies `optimize` in `build`, instead of in
`new` and `optimize`.
- Add `BatchCoalescer::push_batch_with_filter_builder`, which takes a
`FilterBuilder` instead of a mask. It is a safe method: the count is
supplied when the `FilterBuilder` is created, so `with_count` is the
only new `unsafe` API.
- The docs of `filter`, `filter_record_batch`, and
`push_batch_with_filter` point to `FilterBuilder::with_count` for
callers that already know the count.
- Add tests.
`with_count` is marked `unsafe` because the filter kernels rely on the
count being correct: it sizes the filtered output, and a wrong count
leads to out-of-bounds writes. The count is only checked with
`debug_assert_eq!`.
# Are these changes tested?
Yes. Existing tests pass. `test_push_batch_with_filter_builder` covers
the new `BatchCoalescer` method, using one builder with a
caller-supplied count and one the caller optimized. Doc examples cover
`FilterBuilder::with_count` and
`BatchCoalescer::push_batch_with_filter_builder`.
# Are there any user-facing changes?
Two new public APIs: `FilterBuilder::with_count` and
`BatchCoalescer::push_batch_with_filter_builder`. The work done by
`FilterBuilder::new` and `FilterBuilder::optimize` now happens in
`FilterBuilder::build`; this is not visible through the public API.
# AI usage
Developed with Claude Code (Fable 5.1 and Opus 5.5). I have reviewed and
understand the resulting code.
---
arrow-select/src/coalesce.rs | 83 +++++++++++++++++++++++++++++++++-----
arrow-select/src/filter.rs | 96 ++++++++++++++++++++++++++++++++------------
2 files changed, 143 insertions(+), 36 deletions(-)
diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs
index 929526ef40..863d391829 100644
--- a/arrow-select/src/coalesce.rs
+++ b/arrow-select/src/coalesce.rs
@@ -238,6 +238,10 @@ impl BatchCoalescer {
/// This is semantically equivalent of calling [`Self::push_batch`]
/// with the results from [`crate::filter::filter_record_batch`]
///
+ /// If the number of rows that `filter` selects is already known,
+ /// [`Self::push_batch_with_filter_builder`] with
+ /// [`FilterBuilder::with_count`] avoids counting them again.
+ ///
/// # Example
/// ```
/// # use arrow_array::{record_batch, BooleanArray};
@@ -262,7 +266,45 @@ impl BatchCoalescer {
batch: RecordBatch,
filter: &BooleanArray,
) -> Result<(), ArrowError> {
- self.push_batch_with_filtered_columns(batch, filter)
+ self.push_batch_with_filter_builder(batch, FilterBuilder::new(filter))
+ }
+
+ /// Push a batch into the Coalescer after applying the filter described by
+ /// `filter_builder`.
+ ///
+ /// This is [`Self::push_batch_with_filter`] for a [`FilterBuilder`] the
+ /// caller has already created. For example, callers that already know how
+ /// many rows the filter selects can provide that number with
+ /// [`FilterBuilder::with_count`] so that it is not counted again.
+ ///
+ /// Callers do not need to call [`FilterBuilder::optimize`]: like
+ /// [`Self::push_batch_with_filter`], this optimizes the filter when
`batch`
+ /// has more than one column, or one column for which
+ /// [`FilterBuilder::is_optimize_beneficial`] returns true. A filter the
+ /// caller already optimized stays optimized.
+ ///
+ /// # Example
+ /// ```
+ /// # use arrow_array::{record_batch, BooleanArray};
+ /// # use arrow_select::coalesce::BatchCoalescer;
+ /// # use arrow_select::filter::FilterBuilder;
+ /// let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
+ /// let filter = BooleanArray::from(vec![true, false, true]);
+ /// // SAFETY: the filter selects two rows
+ /// let filter_builder = unsafe {
FilterBuilder::new(&filter).with_count(2) };
+ /// let mut coalescer = BatchCoalescer::new(batch.schema(), 1000);
+ /// coalescer.push_batch_with_filter_builder(batch,
filter_builder).unwrap();
+ /// coalescer.finish_buffered_batch().unwrap();
+ /// let expected_batch = record_batch!(("a", Int32, [1, 3])).unwrap();
+ /// assert_eq!(coalescer.next_completed_batch().unwrap(), expected_batch);
+ /// ```
+ pub fn push_batch_with_filter_builder(
+ &mut self,
+ batch: RecordBatch,
+ filter_builder: FilterBuilder,
+ ) -> Result<(), ArrowError> {
+ let predicate = Self::filter_predicate_for_batch(&batch,
filter_builder);
+ self.push_batch_with_filtered_columns(batch, &predicate)
}
/// Push a batch into the Coalescer after applying a set of indices
@@ -606,10 +648,8 @@ impl BatchCoalescer {
impl BatchCoalescer {
fn filter_predicate_for_batch(
batch: &RecordBatch,
- filter: &BooleanArray,
- selected_count: usize,
+ mut filter_builder: FilterBuilder,
) -> FilterPredicate {
- let mut filter_builder = FilterBuilder::new_with_count(filter,
selected_count);
if batch.num_columns() > 1
|| (batch.num_columns() > 0
&&
FilterBuilder::is_optimize_beneficial(batch.schema_ref().field(0).data_type()))
@@ -622,9 +662,10 @@ impl BatchCoalescer {
fn push_batch_with_filtered_columns(
&mut self,
batch: RecordBatch,
- filter: &BooleanArray,
+ predicate: &FilterPredicate,
) -> Result<(), ArrowError> {
- let filter_len = filter.len();
+ let filter_len = predicate.filter_len();
+ let selected_count = predicate.count();
let batch_num_rows = batch.num_rows();
let batch_num_columns = batch.num_columns();
@@ -634,7 +675,6 @@ impl BatchCoalescer {
)));
}
- let selected_count = filter.true_count();
if selected_count == 0 {
return Ok(());
}
@@ -662,16 +702,14 @@ impl BatchCoalescer {
if should_materialize_filter {
// Use materialized filtering when sparse per-column copying is
unavailable.
- let predicate = Self::filter_predicate_for_batch(&batch, filter,
selected_count);
let filtered_batch = predicate.filter_record_batch(&batch)?;
return self.push_batch(filtered_batch);
}
- let predicate = Self::filter_predicate_for_batch(&batch, filter,
selected_count);
let (_schema, arrays, _num_rows) = batch.into_parts();
for (in_progress, array) in
self.in_progress_arrays.iter_mut().zip(arrays) {
- in_progress.copy_rows_by_filter_from(array, &predicate)?;
+ in_progress.copy_rows_by_filter_from(array, predicate)?;
}
self.buffered_rows += selected_count;
@@ -788,7 +826,7 @@ mod tests {
use arrow_array::types::Int32Type;
use arrow_array::{
BinaryViewArray, Int32Array, Int64Array, RecordBatchOptions,
StringArray, StringViewArray,
- TimestampNanosecondArray, UInt32Array, UInt64Array, make_array,
+ TimestampNanosecondArray, UInt32Array, UInt64Array, make_array,
record_batch,
};
use arrow_buffer::BooleanBufferBuilder;
use arrow_schema::{DataType, Field, Schema};
@@ -1025,6 +1063,29 @@ mod tests {
.run();
}
+ #[test]
+ fn test_push_batch_with_filter_builder() {
+ let batch = record_batch!(("a", Int32, [1, 2, 3, 4])).unwrap();
+ // The null is not selected, so the filter selects two rows.
+ let filter = BooleanArray::from(vec![Some(true), None, Some(false),
Some(true)]);
+ let expected = record_batch!(("a", Int32, [1, 4])).unwrap();
+
+ // SAFETY: the count matches the filter.
+ let with_count = unsafe { FilterBuilder::new(&filter).with_count(2) };
+ // The coalescer does not optimize a single primitive column itself, so
+ // this covers a predicate that only the caller optimized.
+ let optimized = FilterBuilder::new(&filter).optimize();
+ for filter_builder in [with_count, optimized] {
+ let mut coalescer = BatchCoalescer::new(batch.schema(), 10);
+ coalescer
+ .push_batch_with_filter_builder(batch.clone(), filter_builder)
+ .unwrap();
+ coalescer.finish_buffered_batch().unwrap();
+ assert_eq!(coalescer.next_completed_batch().unwrap(), expected);
+ assert!(coalescer.next_completed_batch().is_none());
+ }
+ }
+
#[test]
fn test_coalesce_non_null() {
Test::new("coalesce_non_null")
diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs
index ed9b0da637..9baccf745c 100644
--- a/arrow-select/src/filter.rs
+++ b/arrow-select/src/filter.rs
@@ -181,6 +181,9 @@ pub fn prep_null_mask_filter(filter: &BooleanArray) ->
BooleanArray {
/// In contrast to this function, it is then the responsibility of the caller
/// to use [FilterBuilder::optimize] if appropriate.
///
+/// If the number of rows that `predicate` selects is already known,
+/// [`FilterBuilder::with_count`] avoids counting them again.
+///
/// # See also
/// * [`FilterBuilder`] for more control over the filtering process.
/// * [`filter_record_batch`] to filter a [`RecordBatch`]
@@ -223,6 +226,9 @@ pub fn filter(values: &dyn Array, predicate: &BooleanArray)
-> Result<ArrayRef,
/// calling [FilterPredicate::filter_record_batch].
/// In contrast to this function, it is then the responsibility of the caller
/// to use [FilterBuilder::optimize] if appropriate.
+///
+/// If the number of rows that `predicate` selects is already known,
+/// [`FilterBuilder::with_count`] avoids counting them again.
pub fn filter_record_batch(
record_batch: &RecordBatch,
predicate: &BooleanArray,
@@ -248,33 +254,62 @@ pub fn filter_record_batch(
#[derive(Debug)]
pub struct FilterBuilder {
filter: BooleanArray,
- count: usize,
- strategy: IterationStrategy,
+ /// The number of rows `filter` selects, if provided by
[`Self::with_count`]
+ count: Option<usize>,
+ optimize: bool,
}
impl FilterBuilder {
/// Create a new [`FilterBuilder`] that can be used to construct a
[`FilterPredicate`]
pub fn new(filter: &BooleanArray) -> Self {
- Self::new_with_count(filter, filter.true_count())
- }
-
- pub(crate) fn new_with_count(filter: &BooleanArray, count: usize) -> Self {
let filter = match filter.null_count() {
0 => filter.clone(),
_ => prep_null_mask_filter(filter),
};
- let strategy = IterationStrategy::default_strategy(filter.len(),
count);
-
Self {
filter,
- count,
- strategy,
+ count: None,
+ optimize: false,
}
}
- /// Compute an optimized representation of the provided `filter` mask that
can be
- /// applied to an array more quickly.
+ /// Set the number of rows that the filter selects, so that [`Self::build`]
+ /// does not have to count them.
+ ///
+ /// Callers that build a mask row by row, or derive it from a validity
+ /// buffer with a cached null count, often already hold this number.
+ ///
+ /// # Safety
+ ///
+ /// `count` must equal [`BooleanArray::true_count`] of the filter passed to
+ /// [`Self::new`]: the number of `true` values that are not null.
+ ///
+ /// # Example
+ /// ```
+ /// # use arrow_array::{BooleanArray, Int32Array};
+ /// # use arrow_select::filter::FilterBuilder;
+ /// let values = Int32Array::from(vec![1, 2, 3, 4]);
+ /// let mask = BooleanArray::from(vec![Some(true), None, Some(true),
Some(false)]);
+ /// // The null is not selected, so the mask selects two rows.
+ /// // SAFETY: the count matches the mask.
+ /// let predicate = unsafe { FilterBuilder::new(&mask).with_count(2)
}.build();
+ /// assert_eq!(predicate.count(), 2);
+ /// let filtered = predicate.filter(&values).unwrap();
+ /// assert_eq!(filtered.as_ref(), &Int32Array::from(vec![1, 3]));
+ /// ```
+ pub unsafe fn with_count(mut self, count: usize) -> Self {
+ debug_assert_eq!(
+ count,
+ self.filter.true_count(),
+ "count must match the number of rows the filter selects"
+ );
+ self.count = Some(count);
+ self
+ }
+
+ /// Compute an optimized representation of the provided `filter` mask in
+ /// [`Self::build`], so that it can be applied to an array more quickly.
///
/// When filtering multiple arrays (e.g. a [`RecordBatch`] or a
/// [`StructArray`] with multiple fields), optimizing the filter can
provide
@@ -284,17 +319,7 @@ impl FilterBuilder {
/// than the original mask, so it is often faster to filter a single array,
/// without filter optimization.
pub fn optimize(mut self) -> Self {
- match self.strategy {
- IterationStrategy::SlicesIterator => {
- let slices = SlicesIterator::new(&self.filter).collect();
- self.strategy = IterationStrategy::Slices(slices)
- }
- IterationStrategy::IndexIterator => {
- let indices = IndexIterator::new(&self.filter,
self.count).collect();
- self.strategy = IterationStrategy::Indices(indices)
- }
- _ => {}
- }
+ self.optimize = true;
self
}
@@ -316,10 +341,26 @@ impl FilterBuilder {
/// Construct the final `FilterPredicate`
pub fn build(self) -> FilterPredicate {
+ let count = self.count.unwrap_or_else(|| self.filter.true_count());
+ let mut strategy =
IterationStrategy::default_strategy(self.filter.len(), count);
+ if self.optimize {
+ match strategy {
+ IterationStrategy::SlicesIterator => {
+ let slices = SlicesIterator::new(&self.filter).collect();
+ strategy = IterationStrategy::Slices(slices)
+ }
+ IterationStrategy::IndexIterator => {
+ let indices = IndexIterator::new(&self.filter,
count).collect();
+ strategy = IterationStrategy::Indices(indices)
+ }
+ _ => {}
+ }
+ }
+
FilterPredicate {
filter: self.filter,
- count: self.count,
- strategy: self.strategy,
+ count,
+ strategy,
}
}
}
@@ -483,6 +524,11 @@ impl FilterPredicate {
self.count
}
+ /// Length of the filter mask, including rows that are not selected
+ pub(crate) fn filter_len(&self) -> usize {
+ self.filter.len()
+ }
+
/// Return a [`FilterSelection`] for iterating over the rows selected by
/// this [`FilterPredicate`].
pub(crate) fn selection(&self) -> FilterSelection<'_> {
