FilterExec : expose push_batch_with_filter on LimitedBatchCoalescer

datafusion/physical-plan/src/coalesce/mod.rs

@@ -15,8 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow::array::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};
 
@@ -120,6 +120,53 @@ 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<PushBatchStatus> {
+        assert_or_internal_err!(
+            !self.finished,
+            "LimitedBatchCoalescer: cannot push batch after finish"
+        );
+
+        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 selected_count = filter.true_count();
+        if self.total_rows + selected_count >= fetch {
+            let remaining = fetch - self.total_rows;
+            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.slice(0, end), &mask.slice(0, end))?;
+            return Ok(PushBatchStatus::LimitReached);
+        }
+
+        self.total_rows += selected_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()
@@ -225,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
@@ -367,6 +502,15 @@ mod tests {
         .unwrap()
     }
 
+    fn assert_next_batch_values(
+        coalescer: &mut LimitedBatchCoalescer,
+        expected: Vec<u32>,
+    ) {
+        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()

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(_) => {
@@ -1571,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<i32> = batches
+            .iter()
+            .flat_map(|b| {
+                b.column(0)
+                    .as_any()
+                    .downcast_ref::<Int32Array>()
+                    .unwrap()
+                    .values()
+                    .to_vec()
+            })
+            .collect();
+        assert_eq!(values, vec![2, 3, 4]);
+        Ok(())
+    }
+
     #[test]
     fn filter_rejects_zero_batch_size() -> Result<()> {
         let input: Arc<dyn ExecutionPlan> =
@@ -2348,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();