fix: part cols not in projection (#7090)

* fix: part cols not in projection

Signed-off-by: discord9 <discord9@163.com>

* test: table scan with projection

Signed-off-by: discord9 <discord9@163.com>

* Update src/query/src/dist_plan/analyzer.rs

Co-authored-by: Yingwen <realevenyag@gmail.com>
Signed-off-by: discord9 <discord9@163.com>

---------

Signed-off-by: discord9 <discord9@163.com>
Co-authored-by: Yingwen <realevenyag@gmail.com>
This commit is contained in:
discord9
2025-10-16 11:44:30 +08:00
committed by GitHub
parent 145c1024d1
commit 9aca7c97d7
2 changed files with 36 additions and 7 deletions

View File

@@ -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}"

View File

@@ -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());
}