From 9aca7c97d7a8d4a1128e710b7e251e0fe750ea82 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 16 Oct 2025 11:44:30 +0800 Subject: [PATCH] fix: part cols not in projection (#7090) * fix: part cols not in projection Signed-off-by: discord9 * test: table scan with projection Signed-off-by: discord9 * Update src/query/src/dist_plan/analyzer.rs Co-authored-by: Yingwen Signed-off-by: discord9 --------- Signed-off-by: discord9 Co-authored-by: Yingwen --- src/query/src/dist_plan/analyzer.rs | 14 ++++++------ src/query/src/dist_plan/analyzer/test.rs | 29 ++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 7 deletions(-) diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index 9c216f74df..34e035644b 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 1d15530c8d..e9528ccf00 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -1630,3 +1630,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()); +}