diff --git a/datafusion/functions-table/src/generate_series.rs b/datafusion/functions-table/src/generate_series.rs index d00f3d734d76..c66e652147eb 100644 --- a/datafusion/functions-table/src/generate_series.rs +++ b/datafusion/functions-table/src/generate_series.rs @@ -237,6 +237,7 @@ impl GenerateSeriesTable { pub fn as_generator( &self, batch_size: usize, + projection: Option>, ) -> Result>> { let generator: Arc> = match &self.args { GenSeriesArgs::ContainsNull { name } => Arc::new(RwLock::new(Empty { name })), @@ -255,6 +256,7 @@ impl GenerateSeriesTable { batch_size, include_end: *include_end, name, + projection, })), GenSeriesArgs::TimestampArgs { start, @@ -295,6 +297,7 @@ impl GenerateSeriesTable { batch_size, include_end: *include_end, name, + projection, })) } GenSeriesArgs::DateArgs { @@ -324,6 +327,7 @@ impl GenerateSeriesTable { batch_size, include_end: *include_end, name, + projection, })), }; @@ -341,6 +345,7 @@ pub struct GenericSeriesState { current: T, include_end: bool, name: &'static str, + projection: Option>, } impl GenericSeriesState { @@ -396,7 +401,11 @@ impl LazyBatchGenerator for GenericSeriesState { let array = self.current.create_array(buf)?; let batch = RecordBatch::try_new(Arc::clone(&self.schema), vec![array])?; - Ok(Some(batch)) + let projected = match self.projection.as_ref() { + Some(projection) => batch.project(projection)?, + None => batch, + }; + Ok(Some(projected)) } } @@ -477,7 +486,7 @@ impl TableProvider for GenerateSeriesTable { None => self.schema(), }; - let generator = self.as_generator(batch_size)?; + let generator = self.as_generator(batch_size, projection.cloned())?; Ok(Arc::new(LazyMemoryExec::try_new(schema, vec![generator])?)) } diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index e5f4a1f7d026..0ebbb373f2d1 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -1940,7 +1940,8 @@ impl protobuf::PhysicalPlanNode { }; let table = GenerateSeriesTable::new(Arc::clone(&schema), args); - let generator = table.as_generator(generate_series.target_batch_size as usize)?; + let generator = + table.as_generator(generate_series.target_batch_size as usize, None)?; Ok(Arc::new(LazyMemoryExec::try_new(schema, vec![generator])?)) } diff --git a/datafusion/sqllogictest/test_files/table_functions.slt b/datafusion/sqllogictest/test_files/table_functions.slt index 0159abe8d06b..1bffbc3b3a64 100644 --- a/datafusion/sqllogictest/test_files/table_functions.slt +++ b/datafusion/sqllogictest/test_files/table_functions.slt @@ -188,6 +188,21 @@ SELECT generate_series(1, t1.end) FROM generate_series(3, 5) as t1(end) [1, 2, 3, 4] [1, 2, 3] +# join with projection on generate_series +query I +select g1.value from generate_series(1, 3) g1 CROSS JOIN generate_series(1, 3) g2; +---- +1 +1 +1 +2 +2 +2 +3 +3 +3 + + # Test range table function query I SELECT * FROM range(6)