Preserve projection for inline scan (#15825)

* Preserve projection for inline scan

* fix

---------

Co-authored-by: Vadim Piven <[email protected]>

Author: jayzhan211

datafusion/core/tests/execution/logical_plan.rs

@@ -19,15 +19,19 @@
 //! create them and depend on them. Test executable semantics of logical plans.
 
 use arrow::array::Int64Array;
-use arrow::datatypes::{DataType, Field};
+use arrow::datatypes::{DataType, Field, Schema};
+use datafusion::datasource::{provider_as_source, ViewTable};
 use datafusion::execution::session_state::SessionStateBuilder;
-use datafusion_common::{Column, DFSchema, Result, ScalarValue, Spans};
+use datafusion_common::{Column, DFSchema, DFSchemaRef, Result, ScalarValue, Spans};
 use datafusion_execution::TaskContext;
 use datafusion_expr::expr::{AggregateFunction, AggregateFunctionParams};
 use datafusion_expr::logical_plan::{LogicalPlan, Values};
-use datafusion_expr::{Aggregate, AggregateUDF, Expr};
+use datafusion_expr::{
+    Aggregate, AggregateUDF, EmptyRelation, Expr, LogicalPlanBuilder, UNNAMED_TABLE,
+};
 use datafusion_functions_aggregate::count::Count;
 use datafusion_physical_plan::collect;
+use insta::assert_snapshot;
 use std::collections::HashMap;
 use std::fmt::Debug;
 use std::ops::Deref;
@@ -96,3 +100,37 @@ where
     };
     element
 }
+
+#[test]
+fn inline_scan_projection_test() -> Result<()> {
+    let name = UNNAMED_TABLE;
+    let column = "a";
+
+    let schema = Schema::new(vec![
+        Field::new("a", DataType::Int32, false),
+        Field::new("b", DataType::Int32, false),
+    ]);
+    let projection = vec![schema.index_of(column)?];
+
+    let provider = ViewTable::new(
+        LogicalPlan::EmptyRelation(EmptyRelation {
+            produce_one_row: false,
+            schema: DFSchemaRef::new(DFSchema::try_from(schema)?),
+        }),
+        None,
+    );
+    let source = provider_as_source(Arc::new(provider));
+
+    let plan = LogicalPlanBuilder::scan(name, source, Some(projection))?.build()?;
+
+    assert_snapshot!(
+        format!("{plan}"),
+        @r"
+    SubqueryAlias: ?table?
+      Projection: a
+        EmptyRelation
+    "
+    );
+
+    Ok(())
+}

datafusion/expr/src/logical_plan/builder.rs

@@ -501,6 +501,21 @@ impl LogicalPlanBuilder {
         if table_scan.filters.is_empty() {
             if let Some(p) = table_scan.source.get_logical_plan() {
                 let sub_plan = p.into_owned();
+
+                if let Some(proj) = table_scan.projection {
+                    let projection_exprs = proj
+                        .into_iter()
+                        .map(|i| {
+                            Expr::Column(Column::from(
+                                sub_plan.schema().qualified_field(i),
+                            ))
+                        })
+                        .collect::<Vec<_>>();
+                    return Self::new(sub_plan)
+                        .project(projection_exprs)?
+                        .alias(table_scan.table_name);
+                }
+
                 // Ensures that the reference to the inlined table remains the
                 // same, meaning we don't have to change any of the parent nodes
                 // that reference this table.