Add coalescing into FilterExec

Author: alamb

datafusion/physical-plan/src/filter.rs

@@ -44,6 +44,7 @@ use datafusion_physical_expr::{
     analyze, split_conjunction, AnalysisContext, ConstExpr, ExprBoundaries, PhysicalExpr,
 };
 
+use crate::coalesce::{BatchCoalescer, CoalescerState};
 use futures::stream::{Stream, StreamExt};
 use log::trace;
 
@@ -279,10 +280,12 @@ impl ExecutionPlan for FilterExec {
         trace!("Start FilterExec::execute for partition {} of context session_id {} and task_id {:?}", partition, context.session_id(), context.task_id());
         let baseline_metrics = BaselineMetrics::new(&self.metrics, partition);
         Ok(Box::pin(FilterExecStream {
-            schema: self.input.schema(),
+            done: false,
             predicate: Arc::clone(&self.predicate),
             input: self.input.execute(partition, context)?,
             baseline_metrics,
+            // TODO use actual target batch size, for now hardcode the default size
+            coalescer: BatchCoalescer::new(self.input.schema(), 8192, self.fetch()),
         }))
     }
 
@@ -337,14 +340,16 @@ fn collect_new_statistics(
 /// The FilterExec streams wraps the input iterator and applies the predicate expression to
 /// determine which rows to include in its output batches
 struct FilterExecStream {
-    /// Output schema, which is the same as the input schema for this operator
-    schema: SchemaRef,
+    /// Is the sstream done?
+    done: bool,
     /// The expression to filter on. This expression must evaluate to a boolean value.
     predicate: Arc<dyn PhysicalExpr>,
     /// The input partition to filter.
     input: SendableRecordBatchStream,
     /// runtime metrics recording
     baseline_metrics: BaselineMetrics,
+    /// Build up output batches incrementally
+    coalescer: BatchCoalescer,
 }
 
 pub fn batch_filter(
@@ -374,22 +379,47 @@ impl Stream for FilterExecStream {
         mut self: Pin<&mut Self>,
         cx: &mut Context<'_>,
     ) -> Poll<Option<Self::Item>> {
+        if self.done {
+            return Poll::Ready(None);
+        }
         let poll;
         loop {
             match ready!(self.input.poll_next_unpin(cx)) {
                 Some(Ok(batch)) => {
-                    let timer = self.baseline_metrics.elapsed_compute().timer();
+                    // clone timer so we can borrow self mutably
+                    let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+                    let _timer = elapsed_compute.timer(); // records time on drop
                     let filtered_batch = batch_filter(&batch, &self.predicate)?;
-                    timer.done();
-                    // skip entirely filtered batches
-                    if filtered_batch.num_rows() == 0 {
-                        continue;
-                    }
-                    poll = Poll::Ready(Some(Ok(filtered_batch)));
+                    match self.coalescer.push_batch(filtered_batch) {
+                        CoalescerState::TargetReached => {
+                            let batch = self.coalescer.finish_batch();
+                            poll = Poll::Ready(Some(batch));
+                            break;
+                        }
+                        CoalescerState::Continue => {
+                            // sill need more rows
+                            continue;
+                        }
+                        CoalescerState::LimitReached => {
+                            let batch = self.coalescer.finish_batch();
+                            poll = Poll::Ready(Some(batch));
+                            break;
+                        }
+                    };
+                }
+                // end of input, see if we have any remaining batches
+                None => {
+                    self.done = true;
+                    let maybe_result = if self.coalescer.is_empty() {
+                        None
+                    } else {
+                        Some(self.coalescer.finish_batch())
+                    };
+                    poll = Poll::Ready(maybe_result);
                     break;
                 }
-                value => {
-                    poll = Poll::Ready(value);
+                err => {
+                    poll = Poll::Ready(err);
                     break;
                 }
             }
@@ -405,7 +435,7 @@ impl Stream for FilterExecStream {
 
 impl RecordBatchStream for FilterExecStream {
     fn schema(&self) -> SchemaRef {
-        Arc::clone(&self.schema)
+        self.coalescer.schema()
     }
 }