Support limit in StreamingTableExec (apache#10309)

* Support limit in StreamingTableExec

* Update Display and add test

* refine limit test

* Add unit test for limit

* clippy

---------

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

Author: lewiszlw

datafusion/core/src/datasource/stream.rs

@@ -242,7 +242,7 @@ impl TableProvider for StreamTable {
         _state: &SessionState,
         projection: Option<&Vec<usize>>,
         _filters: &[Expr],
-        _limit: Option<usize>,
+        limit: Option<usize>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         let projected_schema = match projection {
             Some(p) => {
@@ -258,6 +258,7 @@ impl TableProvider for StreamTable {
             projection,
             projected_schema,
             true,
+            limit,
         )?))
     }
 

datafusion/core/src/datasource/streaming.rs

@@ -88,15 +88,15 @@ impl TableProvider for StreamingTable {
         _state: &SessionState,
         projection: Option<&Vec<usize>>,
         _filters: &[Expr],
-        _limit: Option<usize>,
+        limit: Option<usize>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        // TODO: push limit down
         Ok(Arc::new(StreamingTableExec::try_new(
             self.schema.clone(),
             self.partitions.clone(),
             projection,
             None,
             self.infinite,
+            limit,
         )?))
     }
 }

datafusion/core/src/physical_optimizer/projection_pushdown.rs

@@ -259,6 +259,7 @@ fn try_swapping_with_streaming_table(
         Some(new_projections.as_ref()),
         lex_orderings,
         streaming_table.is_infinite(),
+        streaming_table.limit(),
     )
     .map(|e| Some(Arc::new(e) as _))
 }
@@ -1860,6 +1861,7 @@ mod tests {
             ]
             .into_iter(),
             true,
+            None,
         )?;
         let projection = Arc::new(ProjectionExec::try_new(
             vec![

datafusion/core/src/physical_optimizer/replace_with_order_preserving_variants.rs

@@ -1475,6 +1475,7 @@ mod tests {
                 Some(&projection),
                 vec![sort_exprs],
                 true,
+                None,
             )
             .unwrap(),
         )

datafusion/core/src/test/mod.rs

@@ -331,6 +331,7 @@ pub fn stream_exec_ordered(
             None,
             vec![sort_exprs],
             true,
+            None,
         )
         .unwrap(),
     )

datafusion/physical-plan/src/limit.rs

@@ -431,7 +431,7 @@ impl ExecutionPlan for LocalLimitExec {
 }
 
 /// A Limit stream skips `skip` rows, and then fetch up to `fetch` rows.
-struct LimitStream {
+pub struct LimitStream {
     /// The remaining number of rows to skip
     skip: usize,
     /// The remaining number of rows to produce
@@ -446,7 +446,7 @@ struct LimitStream {
 }
 
 impl LimitStream {
-    fn new(
+    pub fn new(
         input: SendableRecordBatchStream,
         skip: usize,
         fetch: Option<usize>,

datafusion/physical-plan/src/streaming.rs

@@ -31,6 +31,8 @@ use datafusion_common::{internal_err, plan_err, Result};
 use datafusion_execution::TaskContext;
 use datafusion_physical_expr::{EquivalenceProperties, LexOrdering};
 
+use crate::limit::LimitStream;
+use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
 use async_trait::async_trait;
 use futures::stream::StreamExt;
 use log::debug;
@@ -58,7 +60,9 @@ pub struct StreamingTableExec {
     projected_schema: SchemaRef,
     projected_output_ordering: Vec<LexOrdering>,
     infinite: bool,
+    limit: Option<usize>,
     cache: PlanProperties,
+    metrics: ExecutionPlanMetricsSet,
 }
 
 impl StreamingTableExec {
@@ -69,6 +73,7 @@ impl StreamingTableExec {
         projection: Option<&Vec<usize>>,
         projected_output_ordering: impl IntoIterator<Item = LexOrdering>,
         infinite: bool,
+        limit: Option<usize>,
     ) -> Result<Self> {
         for x in partitions.iter() {
             let partition_schema = x.schema();
@@ -99,7 +104,9 @@ impl StreamingTableExec {
             projection: projection.cloned().map(Into::into),
             projected_output_ordering,
             infinite,
+            limit,
             cache,
+            metrics: ExecutionPlanMetricsSet::new(),
         })
     }
 
@@ -127,6 +134,10 @@ impl StreamingTableExec {
         self.infinite
     }
 
+    pub fn limit(&self) -> Option<usize> {
+        self.limit
+    }
+
     /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.
     fn compute_properties(
         schema: SchemaRef,
@@ -180,6 +191,9 @@ impl DisplayAs for StreamingTableExec {
                 if self.infinite {
                     write!(f, ", infinite_source=true")?;
                 }
+                if let Some(fetch) = self.limit {
+                    write!(f, ", fetch={fetch}")?;
+                }
 
                 display_orderings(f, &self.projected_output_ordering)?;
 
@@ -224,14 +238,116 @@ impl ExecutionPlan for StreamingTableExec {
         ctx: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
         let stream = self.partitions[partition].execute(ctx);
-        Ok(match self.projection.clone() {
+        let projected_stream = match self.projection.clone() {
             Some(projection) => Box::pin(RecordBatchStreamAdapter::new(
                 self.projected_schema.clone(),
                 stream.map(move |x| {
                     x.and_then(|b| b.project(projection.as_ref()).map_err(Into::into))
                 }),
             )),
             None => stream,
+        };
+        Ok(match self.limit {
+            None => projected_stream,
+            Some(fetch) => {
+                let baseline_metrics = BaselineMetrics::new(&self.metrics, partition);
+                Box::pin(LimitStream::new(
+                    projected_stream,
+                    0,
+                    Some(fetch),
+                    baseline_metrics,
+                ))
+            }
         })
     }
+
+    fn metrics(&self) -> Option<MetricsSet> {
+        Some(self.metrics.clone_inner())
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::collect_partitioned;
+    use crate::streaming::PartitionStream;
+    use crate::test::{make_partition, TestPartitionStream};
+    use arrow::record_batch::RecordBatch;
+
+    #[tokio::test]
+    async fn test_no_limit() {
+        let exec = TestBuilder::new()
+            // make 2 batches, each with 100 rows
+            .with_batches(vec![make_partition(100), make_partition(100)])
+            .build();
+
+        let counts = collect_num_rows(Arc::new(exec)).await;
+        assert_eq!(counts, vec![200]);
+    }
+
+    #[tokio::test]
+    async fn test_limit() {
+        let exec = TestBuilder::new()
+            // make 2 batches, each with 100 rows
+            .with_batches(vec![make_partition(100), make_partition(100)])
+            // limit to only the first 75 rows back
+            .with_limit(Some(75))
+            .build();
+
+        let counts = collect_num_rows(Arc::new(exec)).await;
+        assert_eq!(counts, vec![75]);
+    }
+
+    /// Runs the provided execution plan and returns a vector of the number of
+    /// rows in each partition
+    async fn collect_num_rows(exec: Arc<dyn ExecutionPlan>) -> Vec<usize> {
+        let ctx = Arc::new(TaskContext::default());
+        let partition_batches = collect_partitioned(exec, ctx).await.unwrap();
+        partition_batches
+            .into_iter()
+            .map(|batches| batches.iter().map(|b| b.num_rows()).sum::<usize>())
+            .collect()
+    }
+
+    #[derive(Default)]
+    struct TestBuilder {
+        schema: Option<SchemaRef>,
+        partitions: Vec<Arc<dyn PartitionStream>>,
+        projection: Option<Vec<usize>>,
+        projected_output_ordering: Vec<LexOrdering>,
+        infinite: bool,
+        limit: Option<usize>,
+    }
+
+    impl TestBuilder {
+        fn new() -> Self {
+            Self::default()
+        }
+
+        /// Set the batches for the stream
+        fn with_batches(mut self, batches: Vec<RecordBatch>) -> Self {
+            let stream = TestPartitionStream::new_with_batches(batches);
+            self.schema = Some(stream.schema().clone());
+            self.partitions = vec![Arc::new(stream)];
+            self
+        }
+
+        /// Set the limit for the stream
+        fn with_limit(mut self, limit: Option<usize>) -> Self {
+            self.limit = limit;
+            self
+        }
+
+        fn build(self) -> StreamingTableExec {
+            StreamingTableExec::try_new(
+                self.schema.unwrap(),
+                self.partitions,
+                self.projection.as_ref(),
+                self.projected_output_ordering,
+                self.infinite,
+                self.limit,
+            )
+            .unwrap()
+        }
+    }
 }

datafusion/physical-plan/src/test.rs

@@ -23,9 +23,12 @@ use std::sync::Arc;
 
 use arrow_array::{ArrayRef, Int32Array, RecordBatch};
 use arrow_schema::{DataType, Field, Schema, SchemaRef};
+use datafusion_execution::{SendableRecordBatchStream, TaskContext};
 use futures::{Future, FutureExt};
 
 use crate::memory::MemoryExec;
+use crate::stream::RecordBatchStreamAdapter;
+use crate::streaming::PartitionStream;
 use crate::ExecutionPlan;
 
 pub mod exec;
@@ -121,3 +124,26 @@ pub fn mem_exec(partitions: usize) -> MemoryExec {
     let projection = None;
     MemoryExec::try_new(&data, schema, projection).unwrap()
 }
+
+// construct a stream partition for test purposes
+pub struct TestPartitionStream {
+    pub schema: SchemaRef,
+    pub batches: Vec<RecordBatch>,
+}
+
+impl TestPartitionStream {
+    /// Create a new stream partition with the provided batches
+    pub fn new_with_batches(batches: Vec<RecordBatch>) -> Self {
+        let schema = batches[0].schema();
+        Self { schema, batches }
+    }
+}
+impl PartitionStream for TestPartitionStream {
+    fn schema(&self) -> &SchemaRef {
+        &self.schema
+    }
+    fn execute(&self, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
+        let stream = futures::stream::iter(self.batches.clone().into_iter().map(Ok));
+        Box::pin(RecordBatchStreamAdapter::new(self.schema.clone(), stream))
+    }
+}

datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs

@@ -1512,6 +1512,7 @@ mod tests {
             None,
             orderings,
             is_infinite,
+            None,
         )?) as _;
         Ok(source)
     }

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

@@ -624,6 +624,7 @@ mod tests {
             None,
             Some(sort_exprs),
             infinite_source,
+            None,
         )?))
     }