Skip to content

Commit 82181ac

Browse files
authored
fix: update schema's data type for LogicalPlan::Values after placeholder substitution (apache#18740)
## Which issue does this PR close? - Closes apache#18102. ## Rationale for this change This prepare statements below (no type declaration) fails to execute. ```sql DataFusion CLI v51.0.0 > PREPARE my_plan AS SELECT a, b FROM (VALUES ($1, $2)) AS t(a, b); 0 row(s) fetched. Elapsed 0.005 seconds. > EXECUTE my_plan(1, 2); Arrow error: Invalid argument error: column types must match schema types, expected Null but found Int64 at column index 0 > ``` To make it works, schema's data type should be updated after replacing parameters. But calling `recompute_schema` alone does not suffice because we skip recomputing schema for `LogicalPlan::Values`: https://github.com/apache/datafusion/blob/97af468f3da63582f4bacef7d0d1fba4313aa7c4/datafusion/expr/src/logical_plan/plan.rs#L635-L638 ## What changes are included in this PR? - Update schema for `LogicalPlan::Values` manually. ## Are these changes tested? Yes. ## Are there any user-facing changes? No.
1 parent e64de0f commit 82181ac

File tree

2 files changed

+31
-3
lines changed

2 files changed

+31
-3
lines changed

datafusion/expr/src/logical_plan/plan.rs

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1476,13 +1476,29 @@ impl LogicalPlan {
14761476
Ok(transformed_expr.update_data(|expr| original_name.restore(expr)))
14771477
}
14781478
})?
1479-
// always recompute the schema to ensure the changed in the schema's field should be
1480-
// poplulated to the plan's parent
1481-
.map_data(|plan| plan.recompute_schema())
1479+
.map_data(|plan| plan.update_schema_data_type())
14821480
})
14831481
.map(|res| res.data)
14841482
}
14851483

1484+
/// Recompute schema fields' data type after replacing params, ensuring fields data type can be
1485+
/// updated according to the new parameters.
1486+
///
1487+
/// Unlike `recompute_schema()`, this method rebuilds VALUES plans entirely to properly infer
1488+
/// types types from literal values after placeholder substitution.
1489+
fn update_schema_data_type(self) -> Result<LogicalPlan> {
1490+
match self {
1491+
// Build `LogicalPlan::Values` from the values for type inference.
1492+
// We can't use `recompute_schema` because it skips recomputing for
1493+
// `LogicalPlan::Values`.
1494+
LogicalPlan::Values(Values { values, schema: _ }) => {
1495+
LogicalPlanBuilder::values(values)?.build()
1496+
}
1497+
// other plans can just use `recompute_schema` directly.
1498+
plan => plan.recompute_schema(),
1499+
}
1500+
}
1501+
14861502
/// Walk the logical plan, find any `Placeholder` tokens, and return a set of their names.
14871503
pub fn get_parameter_names(&self) -> Result<HashSet<String>> {
14881504
let mut param_names = HashSet::new();

datafusion/sqllogictest/test_files/prepare.slt

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -361,3 +361,15 @@ SET datafusion.explain.logical_plan_only=false;
361361

362362
statement ok
363363
DEALLOCATE my_plan
364+
365+
366+
statement ok
367+
PREPARE my_plan AS SELECT a, b FROM (VALUES ($1, $2)) AS t(a, b);
368+
369+
query II
370+
EXECUTE my_plan(1, 2)
371+
----
372+
1 2
373+
374+
statement ok
375+
DEALLOCATE my_plan

0 commit comments

Comments
 (0)