Move generate_series projection logic into LazyMemoryStream

Author: mkleen

datafusion/core/tests/execution/coop.rs

@@ -148,7 +148,7 @@ fn make_lazy_exec_with_range(
     let generator: Arc<RwLock<dyn LazyBatchGenerator>> = Arc::new(RwLock::new(gen));
 
     // Create a LazyMemoryExec with one partition using our generator
-    let mut exec = LazyMemoryExec::try_new(schema, vec![generator]).unwrap();
+    let mut exec = LazyMemoryExec::try_new(schema, None, vec![generator]).unwrap();
 
     exec.add_ordering(vec![PhysicalSortExpr::new(
         Arc::new(Column::new(column_name, 0)),

datafusion/functions-table/src/generate_series.rs

@@ -237,7 +237,6 @@ impl GenerateSeriesTable {
     pub fn as_generator(
         &self,
         batch_size: usize,
-        projection: Option<Vec<usize>>,
     ) -> Result<Arc<RwLock<dyn LazyBatchGenerator>>> {
         let generator: Arc<RwLock<dyn LazyBatchGenerator>> = match &self.args {
             GenSeriesArgs::ContainsNull { name } => Arc::new(RwLock::new(Empty { name })),
@@ -256,7 +255,6 @@ impl GenerateSeriesTable {
                 batch_size,
                 include_end: *include_end,
                 name,
-                projection,
             })),
             GenSeriesArgs::TimestampArgs {
                 start,
@@ -297,7 +295,6 @@ impl GenerateSeriesTable {
                     batch_size,
                     include_end: *include_end,
                     name,
-                    projection,
                 }))
             }
             GenSeriesArgs::DateArgs {
@@ -327,7 +324,6 @@ impl GenerateSeriesTable {
                 batch_size,
                 include_end: *include_end,
                 name,
-                projection,
             })),
         };
 
@@ -345,7 +341,6 @@ pub struct GenericSeriesState<T: SeriesValue> {
     current: T,
     include_end: bool,
     name: &'static str,
-    projection: Option<Vec<usize>>,
 }
 
 impl<T: SeriesValue> GenericSeriesState<T> {
@@ -401,11 +396,7 @@ impl<T: SeriesValue> LazyBatchGenerator for GenericSeriesState<T> {
 
         let array = self.current.create_array(buf)?;
         let batch = RecordBatch::try_new(Arc::clone(&self.schema), vec![array])?;
-        let projected = match self.projection.as_ref() {
-            Some(projection) => batch.project(projection)?,
-            None => batch,
-        };
-        Ok(Some(projected))
+        Ok(Some(batch))
     }
 }
 
@@ -481,14 +472,13 @@ impl TableProvider for GenerateSeriesTable {
         _limit: Option<usize>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         let batch_size = state.config_options().execution.batch_size;
-        let schema = match projection {
-            Some(projection) => Arc::new(self.schema.project(projection)?),
-            None => self.schema(),
-        };
-
-        let generator = self.as_generator(batch_size, projection.cloned())?;
+        let generator = self.as_generator(batch_size)?;
 
-        Ok(Arc::new(LazyMemoryExec::try_new(schema, vec![generator])?))
+        Ok(Arc::new(LazyMemoryExec::try_new(
+            self.schema(),
+            projection.cloned(),
+            vec![generator],
+        )?))
     }
 }
 

datafusion/physical-plan/src/memory.rs

@@ -153,6 +153,8 @@ pub trait LazyBatchGenerator: Send + Sync + fmt::Debug + fmt::Display {
 pub struct LazyMemoryExec {
     /// Schema representing the data
     schema: SchemaRef,
+    /// Optional projection for which columns to load
+    projection: Option<Vec<usize>>,
     /// Functions to generate batches for each partition
     batch_generators: Vec<Arc<RwLock<dyn LazyBatchGenerator>>>,
     /// Plan properties cache storing equivalence properties, partitioning, and execution mode
@@ -165,6 +167,7 @@ impl LazyMemoryExec {
     /// Create a new lazy memory execution plan
     pub fn try_new(
         schema: SchemaRef,
+        projection: Option<Vec<usize>>,
         generators: Vec<Arc<RwLock<dyn LazyBatchGenerator>>>,
     ) -> Result<Self> {
         let boundedness = generators
@@ -189,6 +192,11 @@ impl LazyMemoryExec {
             })
             .unwrap_or(Boundedness::Bounded);
 
+        let schema = match projection.as_ref() {
+            Some(columns) => Arc::new(schema.project(columns)?),
+            None => schema,
+        };
+
         let cache = PlanProperties::new(
             EquivalenceProperties::new(Arc::clone(&schema)),
             Partitioning::RoundRobinBatch(generators.len()),
@@ -199,6 +207,7 @@ impl LazyMemoryExec {
 
         Ok(Self {
             schema,
+            projection,
             batch_generators: generators,
             cache,
             metrics: ExecutionPlanMetricsSet::new(),
@@ -320,6 +329,7 @@ impl ExecutionPlan for LazyMemoryExec {
 
         let stream = LazyMemoryStream {
             schema: Arc::clone(&self.schema),
+            projection: self.projection.clone(),
             generator: Arc::clone(&self.batch_generators[partition]),
             baseline_metrics,
         };
@@ -338,6 +348,8 @@ impl ExecutionPlan for LazyMemoryExec {
 /// Stream that generates record batches on demand
 pub struct LazyMemoryStream {
     schema: SchemaRef,
+    /// Optional projection for which columns to load
+    projection: Option<Vec<usize>>,
     /// Generator to produce batches
     ///
     /// Note: Idiomatically, DataFusion uses plan-time parallelism - each stream
@@ -361,7 +373,14 @@ impl Stream for LazyMemoryStream {
         let batch = self.generator.write().generate_next_batch();
 
         let poll = match batch {
-            Ok(Some(batch)) => Poll::Ready(Some(Ok(batch))),
+            Ok(Some(batch)) => {
+                // return just the columns requested
+                let batch = match self.projection.as_ref() {
+                    Some(columns) => batch.project(columns)?,
+                    None => batch,
+                };
+                Poll::Ready(Some(Ok(batch)))
+            }
             Ok(None) => Poll::Ready(None),
             Err(e) => Poll::Ready(Some(Err(e))),
         };
@@ -434,8 +453,11 @@ mod lazy_memory_tests {
             schema: Arc::clone(&schema),
         };
 
-        let exec =
-            LazyMemoryExec::try_new(schema, vec![Arc::new(RwLock::new(generator))])?;
+        let exec = LazyMemoryExec::try_new(
+            schema,
+            None,
+            vec![Arc::new(RwLock::new(generator))],
+        )?;
 
         // Test schema
         assert_eq!(exec.schema().fields().len(), 1);
@@ -485,8 +507,11 @@ mod lazy_memory_tests {
             schema: Arc::clone(&schema),
         };
 
-        let exec =
-            LazyMemoryExec::try_new(schema, vec![Arc::new(RwLock::new(generator))])?;
+        let exec = LazyMemoryExec::try_new(
+            schema,
+            None,
+            vec![Arc::new(RwLock::new(generator))],
+        )?;
 
         // Test invalid partition
         let result = exec.execute(1, Arc::new(TaskContext::default()));
@@ -519,8 +544,11 @@ mod lazy_memory_tests {
                 schema: Arc::clone(&schema),
             };
 
-            let exec =
-                LazyMemoryExec::try_new(schema, vec![Arc::new(RwLock::new(generator))])?;
+            let exec = LazyMemoryExec::try_new(
+                schema,
+                None,
+                vec![Arc::new(RwLock::new(generator))],
+            )?;
             let task_ctx = Arc::new(TaskContext::default());
 
             let stream = exec.execute(0, task_ctx)?;

datafusion/proto/src/physical_plan/mod.rs

@@ -1940,10 +1940,13 @@ impl protobuf::PhysicalPlanNode {
         };
 
         let table = GenerateSeriesTable::new(Arc::clone(&schema), args);
-        let generator =
-            table.as_generator(generate_series.target_batch_size as usize, None)?;
+        let generator = table.as_generator(generate_series.target_batch_size as usize)?;
 
-        Ok(Arc::new(LazyMemoryExec::try_new(schema, vec![generator])?))
+        Ok(Arc::new(LazyMemoryExec::try_new(
+            schema,
+            None,
+            vec![generator],
+        )?))
     }
 
     fn try_into_cooperative_physical_plan(