From 82aa9fa3c908f98e0a909839ef6c2a05220084e0 Mon Sep 17 00:00:00 2001 From: linfeng Date: Thu, 23 Jul 2026 00:12:50 +0800 Subject: [PATCH 1/3] feat: support filtered batches in LimitedBatchCoalescer --- datafusion/physical-plan/src/coalesce/mod.rs | 125 ++++++++++++++++++- datafusion/physical-plan/src/filter.rs | 6 +- 2 files changed, 126 insertions(+), 5 deletions(-) diff --git a/datafusion/physical-plan/src/coalesce/mod.rs b/datafusion/physical-plan/src/coalesce/mod.rs index ea1a87d091481..640e727b50257 100644 --- a/datafusion/physical-plan/src/coalesce/mod.rs +++ b/datafusion/physical-plan/src/coalesce/mod.rs @@ -15,9 +15,10 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::RecordBatch; +use arrow::array::{Array, BooleanArray, RecordBatch}; use arrow::compute::BatchCoalescer; use arrow::datatypes::SchemaRef; +use arrow::error::ArrowError; use datafusion_common::{Result, assert_or_internal_err}; /// Concatenate multiple [`RecordBatch`]es and apply a limit @@ -120,6 +121,69 @@ impl LimitedBatchCoalescer { Ok(PushBatchStatus::Continue) } + /// Pushes the next [`RecordBatch`] into the coalescer after applying `filter`. + 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 filter.len() > batch.num_rows() { + return Err(ArrowError::InvalidArgumentError(format!( + "Filter predicate of length {} is larger than target array of length {}", + filter.len(), + batch.num_rows() + )) + .into()); + } + + let Some(fetch) = self.fetch else { + self.inner.push_batch_with_filter(batch, filter)?; + return Ok(PushBatchStatus::Continue); + }; + + if self.total_rows >= fetch { + return Ok(PushBatchStatus::LimitReached); + } + + let remaining_rows = fetch - self.total_rows; + let mut selected_rows = 0; + let mut filter_len = None; + if filter.null_count() == 0 { + for index in filter.values().set_indices() { + selected_rows += 1; + if selected_rows == remaining_rows { + filter_len = Some(index + 1); + break; + } + } + } else { + for (index, value) in filter.iter().enumerate() { + if value == Some(true) { + selected_rows += 1; + if selected_rows == remaining_rows { + filter_len = Some(index + 1); + break; + } + } + } + } + + if let Some(filter_len) = filter_len { + let filter = filter.slice(0, filter_len); + self.inner.push_batch_with_filter(batch, &filter)?; + self.total_rows += remaining_rows; + Ok(PushBatchStatus::LimitReached) + } else { + self.inner.push_batch_with_filter(batch, filter)?; + self.total_rows += selected_rows; + Ok(PushBatchStatus::Continue) + } + } + /// Return true if there is no data buffered pub fn is_empty(&self) -> bool { self.inner.is_empty() @@ -150,7 +214,8 @@ mod tests { use std::ops::Range; use std::sync::Arc; - use arrow::array::UInt32Array; + use arrow::array::{BooleanArray, UInt32Array}; + use arrow::buffer::{BooleanBuffer, NullBuffer}; use arrow::compute::concat_batches; use arrow::datatypes::{DataType, Field, Schema}; @@ -225,6 +290,62 @@ mod tests { .run() } + #[test] + fn test_coalesce_with_filter_without_fetch() { + let batch = uint32_batch(0..6); + let filter = BooleanArray::from(vec![true, false, true, false, false, true]); + let mut coalescer = LimitedBatchCoalescer::new(batch.schema(), 20, None); + + let status = coalescer.push_batch_with_filter(batch, &filter).unwrap(); + assert_eq!(status, PushBatchStatus::Continue); + + coalescer.finish().unwrap(); + let output = coalescer.next_completed_batch().unwrap(); + let output = output + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(output, &UInt32Array::from(vec![0, 2, 5])); + } + + #[test] + fn test_coalesce_with_filter_and_fetch() { + let first_batch = uint32_batch(0..5); + let second_batch = uint32_batch(5..11); + let mut coalescer = LimitedBatchCoalescer::new(first_batch.schema(), 20, Some(3)); + + let first_filter = BooleanArray::new( + BooleanBuffer::from(vec![true, true, false, false, false]), + Some(NullBuffer::from(vec![true, false, true, true, true])), + ); + assert_eq!( + coalescer + .push_batch_with_filter(first_batch, &first_filter) + .unwrap(), + PushBatchStatus::Continue + ); + + let second_filter = + BooleanArray::from(vec![false, false, true, false, true, true, true]) + .slice(1, 6); + assert_eq!( + coalescer + .push_batch_with_filter(second_batch, &second_filter) + .unwrap(), + PushBatchStatus::LimitReached + ); + + coalescer.finish().unwrap(); + let output = coalescer.next_completed_batch().unwrap(); + let output = output + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(output, &UInt32Array::from(vec![0, 6, 8])); + } + /// Test for [`LimitedBatchCoalescer`] /// /// Pushes the input batches to the coalescer and verifies that the resulting diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index d367be16eb6ed..c056b32f9f71d 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -1265,9 +1265,9 @@ 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 0ef08017fd23b36a87f9a320a077041aa76e2da4 Mon Sep 17 00:00:00 2001 From: linfeng Date: Thu, 23 Jul 2026 21:47:43 +0800 Subject: [PATCH 2/3] expand filtered coalescer coverage and docs --- datafusion/physical-plan/src/coalesce/mod.rs | 23 ++++++++++++ datafusion/physical-plan/src/filter.rs | 39 ++++++++++++++++++++ 2 files changed, 62 insertions(+) diff --git a/datafusion/physical-plan/src/coalesce/mod.rs b/datafusion/physical-plan/src/coalesce/mod.rs index 640e727b50257..1076a72a519d3 100644 --- a/datafusion/physical-plan/src/coalesce/mod.rs +++ b/datafusion/physical-plan/src/coalesce/mod.rs @@ -122,6 +122,17 @@ impl LimitedBatchCoalescer { } /// Pushes the next [`RecordBatch`] into the coalescer after applying `filter`. + /// + /// Only non-null `true` values count toward `fetch`. When the limit is reached, + /// the filter is truncated after the final selected row. + /// + /// # Returns + /// Returns [`PushBatchStatus::LimitReached`] when `fetch` is reached and + /// [`PushBatchStatus::Continue`] otherwise. + /// + /// # Errors + /// Returns an error if called after [`Self::finish`], if `filter` is longer + /// than the batch, or if the internal filtered push operation fails. pub fn push_batch_with_filter( &mut self, batch: RecordBatch, @@ -344,6 +355,18 @@ mod tests { .downcast_ref::() .unwrap(); assert_eq!(output, &UInt32Array::from(vec![0, 6, 8])); + + let oversized_batch = uint32_batch(0..5); + let oversized_filter = BooleanArray::from(vec![true; 6]); + let mut coalescer = + LimitedBatchCoalescer::new(oversized_batch.schema(), 20, Some(3)); + let error = coalescer + .push_batch_with_filter(oversized_batch, &oversized_filter) + .unwrap_err(); + assert_eq!( + error.to_string(), + "Arrow error: Invalid argument error: Filter predicate of length 6 is larger than target array of length 5" + ); } /// Test for [`LimitedBatchCoalescer`] diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index c056b32f9f71d..0d51f2821a9fb 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -1376,6 +1376,7 @@ mod tests { use crate::test; use crate::test::exec::StatisticsExec; use arrow::datatypes::{Field, Schema, UnionFields, UnionMode}; + use datafusion_common::assert_batches_eq; #[tokio::test] async fn collect_columns_predicates() -> Result<()> { @@ -2206,6 +2207,44 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_filter_with_projection_and_fetch() -> Result<()> { + let a = vec![5, 11, 12, 13]; + let b = vec![10, 20, 30, 40]; + let c = vec![50, 110, 120, 130]; + let input = test::build_table_scan_i32(("a", &a), ("b", &b), ("c", &c)); + + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(10)))), + )); + + let filter = FilterExecBuilder::new(predicate, input) + .apply_projection(Some(vec![0, 2]))? + .with_fetch(Some(2)) + .build()?; + + let batches = crate::execution_plan::collect( + Arc::new(filter), + Arc::new(TaskContext::default()), + ) + .await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| a | c |", + "+----+-----+", + "| 11 | 110 |", + "| 12 | 120 |", + "+----+-----+", + ], + &batches + ); + + Ok(()) + } + #[tokio::test] async fn test_builder_without_projection() -> Result<()> { let schema = Arc::new(Schema::new(vec![ From 23beb6a0848fe06c0d6b5f2735593554e3a1cf52 Mon Sep 17 00:00:00 2001 From: linfeng <33561138+lyne7-sc@users.noreply.github.com> Date: Fri, 24 Jul 2026 12:39:35 +0800 Subject: [PATCH 3/3] fix ci --- datafusion/physical-plan/src/coalesce/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/coalesce/mod.rs b/datafusion/physical-plan/src/coalesce/mod.rs index 1076a72a519d3..2bbd6b0879b6b 100644 --- a/datafusion/physical-plan/src/coalesce/mod.rs +++ b/datafusion/physical-plan/src/coalesce/mod.rs @@ -364,7 +364,7 @@ mod tests { .push_batch_with_filter(oversized_batch, &oversized_filter) .unwrap_err(); assert_eq!( - error.to_string(), + error.strip_backtrace(), "Arrow error: Invalid argument error: Filter predicate of length 6 is larger than target array of length 5" ); }