[arrow-datafusion] branch main updated: generate new projection plan in inline_table_scan instead of discarding (#5371)

[jakevin]

This is an automated email from the ASF dual-hosted git repository.

jakevin pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/main by this push:
     new 095f93484 generate new projection plan in inline_table_scan instead of 
discarding (#5371)
095f93484 is described below

commit 095f9348455880e6169c09a7c02edfe231aac59a
Author: jakevin <jakevin...@gmail.com>
AuthorDate: Thu Feb 23 21:46:28 2023 +0800

    generate new projection plan in inline_table_scan instead of discarding 
(#5371)
    
    * fix inline_table_scan
    
    * add test
---
 datafusion/optimizer/src/inline_table_scan.rs | 45 +++++++++++++++++++++++++--
 1 file changed, 42 insertions(+), 3 deletions(-)

diff --git a/datafusion/optimizer/src/inline_table_scan.rs 
b/datafusion/optimizer/src/inline_table_scan.rs
index 39e7d4384..6b5839919 100644
--- a/datafusion/optimizer/src/inline_table_scan.rs
+++ b/datafusion/optimizer/src/inline_table_scan.rs
@@ -49,11 +49,14 @@ impl OptimizerRule for InlineTableScan {
                 source,
                 table_name,
                 filters,
+                projection,
                 ..
             }) if filters.is_empty() => {
                 if let Some(sub_plan) = source.get_logical_plan() {
+                    let projection_exprs =
+                        generate_projection_expr(projection, sub_plan)?;
                     let plan = LogicalPlanBuilder::from(sub_plan.clone())
-                        .project(vec![Expr::Wildcard])?
+                        .project(projection_exprs)?
                         .alias(table_name)?;
                     Ok(Some(plan.build()?))
                 } else {
@@ -73,6 +76,23 @@ impl OptimizerRule for InlineTableScan {
     }
 }
 
+fn generate_projection_expr(
+    projection: &Option<Vec<usize>>,
+    sub_plan: &LogicalPlan,
+) -> Result<Vec<Expr>> {
+    let mut exprs = vec![];
+    if let Some(projection) = projection {
+        for i in projection {
+            exprs.push(Expr::Column(
+                sub_plan.schema().fields()[*i].qualified_column(),
+            ));
+        }
+    } else {
+        exprs.push(Expr::Wildcard);
+    }
+    Ok(exprs)
+}
+
 #[cfg(test)]
 mod tests {
     use std::{sync::Arc, vec};
@@ -91,7 +111,10 @@ mod tests {
         }
 
         fn schema(&self) -> arrow::datatypes::SchemaRef {
-            Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, 
false)]))
+            Arc::new(Schema::new(vec![
+                Field::new("a", DataType::Int64, false),
+                Field::new("b", DataType::Int64, false),
+            ]))
         }
 
         fn supports_filter_pushdown(
@@ -150,9 +173,25 @@ mod tests {
         let plan = scan.filter(col("x.a").eq(lit(1)))?.build()?;
         let expected = "Filter: x.a = Int32(1)\
         \n  SubqueryAlias: x\
-        \n    Projection: y.a\
+        \n    Projection: y.a, y.b\
         \n      TableScan: y";
 
         assert_optimized_plan_eq(Arc::new(InlineTableScan::new()), &plan, 
expected)
     }
+
+    #[test]
+    fn inline_table_scan_with_projection() -> datafusion_common::Result<()> {
+        let scan = LogicalPlanBuilder::scan(
+            "x".to_string(),
+            Arc::new(CustomSource::new()),
+            Some(vec![0]),
+        )?;
+
+        let plan = scan.build()?;
+        let expected = "SubqueryAlias: x\
+        \n  Projection: y.a\
+        \n    TableScan: y";
+
+        assert_optimized_plan_eq(Arc::new(InlineTableScan::new()), &plan, 
expected)
+    }
 }