From f471266ad5f392fc1c9c0aa03e60887c1aa7f07a Mon Sep 17 00:00:00 2001 From: RIchard Baah Date: Sat, 19 Sep 2026 11:59:05 -0400 Subject: [PATCH 1/4] coalesce: expose push_batch_with_filter on LimitedBatchCoalescer Arrow's BatchCoalescer already has push_batch_with_filter, which writes surviving rows directly into the coalescer buffer in a single pass. LimitedBatchCoalescer was not exposing it, forcing FilterExec to call filter_record_batch first (materialising a new allocation) and then push_batch (a second copy). Wire push_batch_with_filter through LimitedBatchCoalescer with correct fetch-limit accounting, and update FilterExec to use it, eliminating the intermediate allocation on the hot filter path. --- datafusion/physical-plan/src/coalesce/mod.rs | 53 +++++++++++++++++++- datafusion/physical-plan/src/filter.rs | 4 +- 2 files changed, 53 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/src/coalesce/mod.rs b/datafusion/physical-plan/src/coalesce/mod.rs index ea1a87d091481..4b9a797136c32 100644 --- a/datafusion/physical-plan/src/coalesce/mod.rs +++ b/datafusion/physical-plan/src/coalesce/mod.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::RecordBatch; +use arrow::array::{BooleanArray, RecordBatch}; use arrow::compute::BatchCoalescer; use arrow::datatypes::SchemaRef; use datafusion_common::{Result, assert_or_internal_err}; @@ -120,6 +120,44 @@ impl LimitedBatchCoalescer { Ok(PushBatchStatus::Continue) } + /// Pushes the next [`RecordBatch`] into the coalescer after applying `filter`, + /// avoiding a separate materialization pass compared to calling + /// [`filter_record_batch`] followed by [`Self::push_batch`]. + /// + /// [`filter_record_batch`]: arrow::compute::filter_record_batch + pub fn push_batch_with_filter( + &mut self, + batch: RecordBatch, + filter: &BooleanArray, + ) -> Result { + assert_or_internal_err!( + !self.finished, + "LimitedBatchCoalescer: cannot push batch after finish" + ); + + if let Some(fetch) = self.fetch { + if self.total_rows >= fetch { + return Ok(PushBatchStatus::LimitReached); + } + + let selected_count = filter.true_count(); + if self.total_rows + selected_count >= fetch { + let remaining = fetch - self.total_rows; + let truncated = first_n_true(filter, remaining); + self.total_rows += remaining; + self.inner.push_batch_with_filter(batch, &truncated)?; + return Ok(PushBatchStatus::LimitReached); + } + + self.total_rows += selected_count; + } else { + self.total_rows += filter.true_count(); + } + + self.inner.push_batch_with_filter(batch, filter)?; + Ok(PushBatchStatus::Continue) + } + /// Return true if there is no data buffered pub fn is_empty(&self) -> bool { self.inner.is_empty() @@ -144,6 +182,19 @@ impl LimitedBatchCoalescer { } } +/// Returns a copy of `filter` where only the first `n` `true` values are kept; +/// all subsequent `true` values are replaced with `false`. +fn first_n_true(filter: &BooleanArray, n: usize) -> BooleanArray { + let mut count = 0; + BooleanArray::from_iter(filter.iter().map(|v| match v { + Some(true) if count < n => { + count += 1; + Some(true) + } + _ => Some(false), + })) +} + #[cfg(test)] mod tests { use super::*; diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index 7377ec893621c..8176ce5f6ffff 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -1466,9 +1466,7 @@ impl Stream for FilterExecStream { match as_boolean_array(&array) { Ok(filter_array) => { self.metrics.selectivity.add_total(batch.num_rows()); - // TODO: support push_batch_with_filter in LimitedBatchCoalescer - let batch = filter_record_batch(&batch, filter_array)?; - let state = self.batch_coalescer.push_batch(batch)?; + let state = self.batch_coalescer.push_batch_with_filter(batch, filter_array)?; Ok(state) } Err(_) => { From e3a83ddb1aee52ec7baa0f763cf515ea17a7c009 Mon Sep 17 00:00:00 2001 From: RIchard Baah Date: Sat, 19 Sep 2026 16:27:49 -0400 Subject: [PATCH 2/4] add test --- datafusion/physical-plan/src/filter.rs | 40 +++++++++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index 8176ce5f6ffff..f71d67d04e39c 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -1569,13 +1569,51 @@ pub type EqualAndNonEqual<'a> = #[cfg(test)] mod tests { use super::*; + use crate::common::collect; use crate::empty::EmptyExec; use crate::expressions::*; use crate::statistics::{StatisticsArgs, StatisticsContext}; use crate::test; use crate::test::exec::StatisticsExec; + use arrow::array::Int32Array; use arrow::datatypes::{Field, Schema, UnionFields, UnionMode}; + #[tokio::test] + async fn test_filter_exec_fetch_truncates_within_selected_rows() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from_iter_values(0..10))], + )?; + let input = test::TestMemoryExec::try_new_exec( + &[vec![batch]], + Arc::clone(&schema), + None, + )?; + let predicate = binary(col("i", &schema)?, Operator::GtEq, lit(2i32), &schema)?; + let filter = Arc::new( + FilterExecBuilder::new(predicate, input) + .with_fetch(Some(3)) + .build()?, + ); + + let task_ctx = Arc::new(TaskContext::default()); + let batches = collect(filter.execute(0, task_ctx)?).await?; + let values: Vec = batches + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + assert_eq!(values, vec![2, 3, 4]); + Ok(()) + } + #[test] fn filter_rejects_zero_batch_size() -> Result<()> { let input: Arc = @@ -2346,7 +2384,7 @@ mod tests { // A mathematical lower bound of 5 would exclude a valid input value. let batch = RecordBatch::try_new( input.schema(), - vec![Arc::new(arrow::array::Int32Array::from(vec![i32::MIN]))], + vec![Arc::new(Int32Array::from(vec![i32::MIN]))], ) .unwrap(); let result = predicate.evaluate(&batch).unwrap().into_array(1).unwrap(); From ad2a7f0758cbc869d2d50d33f1ef5284e1d59cfb Mon Sep 17 00:00:00 2001 From: RIchard Baah Date: Sat, 19 Sep 2026 16:51:28 -0400 Subject: [PATCH 3/4] add fast path --- datafusion/physical-plan/src/coalesce/mod.rs | 30 ++++++++++---------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/datafusion/physical-plan/src/coalesce/mod.rs b/datafusion/physical-plan/src/coalesce/mod.rs index 4b9a797136c32..b140ce81a32b9 100644 --- a/datafusion/physical-plan/src/coalesce/mod.rs +++ b/datafusion/physical-plan/src/coalesce/mod.rs @@ -135,25 +135,25 @@ impl LimitedBatchCoalescer { "LimitedBatchCoalescer: cannot push batch after finish" ); - if let Some(fetch) = self.fetch { - if self.total_rows >= fetch { - return Ok(PushBatchStatus::LimitReached); - } + let Some(fetch) = self.fetch else { + self.inner.push_batch_with_filter(batch, filter)?; + return Ok(PushBatchStatus::Continue); + }; - let selected_count = filter.true_count(); - if self.total_rows + selected_count >= fetch { - let remaining = fetch - self.total_rows; - let truncated = first_n_true(filter, remaining); - self.total_rows += remaining; - self.inner.push_batch_with_filter(batch, &truncated)?; - return Ok(PushBatchStatus::LimitReached); - } + if self.total_rows >= fetch { + return Ok(PushBatchStatus::LimitReached); + } - self.total_rows += selected_count; - } else { - self.total_rows += filter.true_count(); + let selected_count = filter.true_count(); + if self.total_rows + selected_count >= fetch { + let remaining = fetch - self.total_rows; + let truncated = first_n_true(filter, remaining); + self.total_rows += remaining; + self.inner.push_batch_with_filter(batch, &truncated)?; + return Ok(PushBatchStatus::LimitReached); } + self.total_rows += selected_count; self.inner.push_batch_with_filter(batch, filter)?; Ok(PushBatchStatus::Continue) } From 6b1623c37f6efefb3d82710011020a82ffd415d0 Mon Sep 17 00:00:00 2001 From: RIchard Baah Date: Sun, 20 Sep 2026 14:34:57 -0400 Subject: [PATCH 4/4] respond to Pr review --- datafusion/physical-plan/src/coalesce/mod.rs | 127 ++++++++++++++++--- 1 file changed, 110 insertions(+), 17 deletions(-) diff --git a/datafusion/physical-plan/src/coalesce/mod.rs b/datafusion/physical-plan/src/coalesce/mod.rs index b140ce81a32b9..70c86368c1477 100644 --- a/datafusion/physical-plan/src/coalesce/mod.rs +++ b/datafusion/physical-plan/src/coalesce/mod.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{BooleanArray, RecordBatch}; -use arrow::compute::BatchCoalescer; +use arrow::array::{Array, BooleanArray, RecordBatch}; +use arrow::compute::{BatchCoalescer, prep_null_mask_filter}; use arrow::datatypes::SchemaRef; use datafusion_common::{Result, assert_or_internal_err}; @@ -147,9 +147,18 @@ impl LimitedBatchCoalescer { let selected_count = filter.true_count(); if self.total_rows + selected_count >= fetch { let remaining = fetch - self.total_rows; - let truncated = first_n_true(filter, remaining); + let mask = match filter.null_count() { + 0 => filter.clone(), + _ => prep_null_mask_filter(filter), + }; + let end = mask + .values() + .set_indices() + .nth(remaining - 1) + .map_or(0, |i| i + 1); self.total_rows += remaining; - self.inner.push_batch_with_filter(batch, &truncated)?; + self.inner + .push_batch_with_filter(batch.slice(0, end), &mask.slice(0, end))?; return Ok(PushBatchStatus::LimitReached); } @@ -182,19 +191,6 @@ impl LimitedBatchCoalescer { } } -/// Returns a copy of `filter` where only the first `n` `true` values are kept; -/// all subsequent `true` values are replaced with `false`. -fn first_n_true(filter: &BooleanArray, n: usize) -> BooleanArray { - let mut count = 0; - BooleanArray::from_iter(filter.iter().map(|v| match v { - Some(true) if count < n => { - count += 1; - Some(true) - } - _ => Some(false), - })) -} - #[cfg(test)] mod tests { use super::*; @@ -276,6 +272,94 @@ mod tests { .run() } + #[test] + fn test_push_batch_with_filter_nulls_and_fetch() { + let batch = uint32_batch(0..8); + let mut coalescer = LimitedBatchCoalescer::new(batch.schema(), 100, Some(3)); + let filter = BooleanArray::from(vec![ + None, + Some(true), + None, + Some(false), + Some(true), + Some(true), + None, + Some(true), + ]); + + assert_eq!( + coalescer.push_batch_with_filter(batch, &filter).unwrap(), + PushBatchStatus::LimitReached, + ); + coalescer.finish().unwrap(); + assert_next_batch_values(&mut coalescer, vec![1, 4, 5]); + } + + #[test] + fn test_push_batch_with_filter_fetch_boundaries() { + let batch1 = uint32_batch(0..4); + let batch2 = uint32_batch(4..8); + let mut coalescer = LimitedBatchCoalescer::new(batch1.schema(), 100, Some(3)); + + assert_eq!( + coalescer + .push_batch_with_filter( + batch1, + &BooleanArray::from(vec![true, false, true, false]), + ) + .unwrap(), + PushBatchStatus::Continue, + ); + assert_eq!( + coalescer + .push_batch_with_filter( + batch2, + &BooleanArray::from(vec![true, true, true, true]), + ) + .unwrap(), + PushBatchStatus::LimitReached, + ); + coalescer.finish().unwrap(); + assert_next_batch_values(&mut coalescer, vec![0, 2, 4]); + + let batch = uint32_batch(0..4); + let mut coalescer = LimitedBatchCoalescer::new(batch.schema(), 100, Some(2)); + assert_eq!( + coalescer + .push_batch_with_filter( + batch, + &BooleanArray::from(vec![true, false, true, false]), + ) + .unwrap(), + PushBatchStatus::LimitReached, + ); + assert_eq!( + coalescer + .push_batch_with_filter( + uint32_batch(4..8), + &BooleanArray::from(vec![true, true, true, true]), + ) + .unwrap(), + PushBatchStatus::LimitReached, + ); + coalescer.finish().unwrap(); + assert_next_batch_values(&mut coalescer, vec![0, 2]); + + let batch = uint32_batch(0..4); + let mut coalescer = LimitedBatchCoalescer::new(batch.schema(), 100, Some(0)); + assert_eq!( + coalescer + .push_batch_with_filter( + batch, + &BooleanArray::from(vec![true, true, true, true]), + ) + .unwrap(), + PushBatchStatus::LimitReached, + ); + coalescer.finish().unwrap(); + assert!(coalescer.next_completed_batch().is_none()); + } + /// Test for [`LimitedBatchCoalescer`] /// /// Pushes the input batches to the coalescer and verifies that the resulting @@ -418,6 +502,15 @@ mod tests { .unwrap() } + fn assert_next_batch_values( + coalescer: &mut LimitedBatchCoalescer, + expected: Vec, + ) { + let output = coalescer.next_completed_batch().unwrap(); + let expected = UInt32Array::from(expected); + assert_eq!(output.column(0).as_ref(), &expected as &dyn Array); + } + fn batch_to_pretty_strings(batch: &RecordBatch) -> String { arrow::util::pretty::pretty_format_batches(std::slice::from_ref(batch)) .unwrap()