diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index 9c216f74df1c..34e035644b64 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -544,13 +544,13 @@ impl PlanRewriter { return Ok((c.clone(), BTreeSet::new())); } let index = - plan.schema().index_of_column_by_name(None, &c).ok_or_else(|| { - datafusion_common::DataFusionError::Internal( - format!( - "PlanRewriter: maybe_set_partitions: column {c} not found in schema of plan: {plan}" - ), - ) - })?; + if let Some(c) = plan.schema().index_of_column_by_name(None, &c){ + c + } else { + // the `projection` field of `TableScan` doesn't contain the partition columns, + // this is similar to not having a alias, hence return empty alias set + return Ok((c.clone(), BTreeSet::new())) + }; let column = plan.schema().columns().get(index).cloned().ok_or_else(|| { datafusion_common::DataFusionError::Internal(format!( "PlanRewriter: maybe_set_partitions: column index {index} out of bounds in schema of plan: {plan}" diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index c9c874dbc2d9..93e41cfb04dc 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -1629,3 +1629,32 @@ fn test_last_value_no_order_by() { .join("\n"); assert_eq!(expected, result.to_string()); } + +#[test] +fn test_table_scan_projection() { + init_default_ut_logging(); + let test_table = TestTable::table_with_name(0, "t".to_string()); + let table_provider = Arc::new(DfTableProviderAdapter::new(test_table)); + let table_source = Arc::new(DefaultTableSource::new(table_provider.clone() as _)); + let ctx = SessionContext::new(); + ctx.register_table(TableReference::bare("t"), table_provider.clone() as _) + .unwrap(); + + let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, Some(vec![3]), vec![]) + .unwrap() + .build() + .unwrap(); + + let config = ConfigOptions::default(); + let result = DistPlannerAnalyzer {} + .analyze(plan.clone(), &config) + .unwrap(); + let expected = [ + "Projection: t.ts", + " MergeScan [is_placeholder=false, remote_input=[", + "TableScan: t projection=[ts]", + "]]", + ] + .join("\n"); + assert_eq!(expected, result.to_string()); +}