Skip to content

Commit 2ed2e3a

Browse files
committed
Preserve projection for inline scan
1 parent 2f7cc9b commit 2ed2e3a

File tree

2 files changed

+34
-3
lines changed

2 files changed

+34
-3
lines changed

datafusion/core/tests/execution/logical_plan.rs

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,14 @@
2020
2121
use arrow::array::Int64Array;
2222
use arrow::datatypes::{DataType, Field};
23+
use arrow_schema::Schema;
24+
use datafusion::datasource::{provider_as_source, ViewTable};
2325
use datafusion::execution::session_state::SessionStateBuilder;
24-
use datafusion_common::{Column, DFSchema, Result, ScalarValue, Spans};
26+
use datafusion_common::{Column, DFSchema, DFSchemaRef, Result, ScalarValue, Spans};
2527
use datafusion_execution::TaskContext;
2628
use datafusion_expr::expr::{AggregateFunction, AggregateFunctionParams};
2729
use datafusion_expr::logical_plan::{LogicalPlan, Values};
28-
use datafusion_expr::{Aggregate, AggregateUDF, Expr};
30+
use datafusion_expr::{Aggregate, AggregateUDF, EmptyRelation, Expr, LogicalPlanBuilder, UNNAMED_TABLE};
2931
use datafusion_functions_aggregate::count::Count;
3032
use datafusion_physical_plan::collect;
3133
use std::collections::HashMap;
@@ -96,3 +98,32 @@ where
9698
};
9799
element
98100
}
101+
102+
#[test]
103+
fn inline_scan_projection_test() -> Result<()> {
104+
let name = UNNAMED_TABLE;
105+
let column = "a";
106+
107+
let schema = Schema::new(vec![
108+
Field::new("a", DataType::Int32, false),
109+
Field::new("b", DataType::Int32, false),
110+
]);
111+
let projection = vec![schema.index_of(column)?];
112+
113+
let provider = ViewTable::new(
114+
LogicalPlan::EmptyRelation(EmptyRelation {
115+
produce_one_row: false,
116+
schema: DFSchemaRef::new(DFSchema::try_from(schema)?),
117+
}),
118+
None,
119+
);
120+
let source = provider_as_source(Arc::new(provider));
121+
122+
let plan = LogicalPlanBuilder::scan(name, source, Some(projection))?.build()?;
123+
assert_eq!(
124+
format!("{plan}"),
125+
format!("TableScan: {name} projection=[{column}]")
126+
);
127+
128+
Ok(())
129+
}

datafusion/expr/src/logical_plan/builder.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -498,7 +498,7 @@ impl LogicalPlanBuilder {
498498
TableScan::try_new(table_name, table_source, projection, filters, fetch)?;
499499

500500
// Inline TableScan
501-
if table_scan.filters.is_empty() {
501+
if table_scan.projection.is_none() && table_scan.filters.is_empty() {
502502
if let Some(p) = table_scan.source.get_logical_plan() {
503503
let sub_plan = p.into_owned();
504504
// Ensures that the reference to the inlined table remains the

0 commit comments

Comments
 (0)