kosiew commented on code in PR #23260:
URL: https://github.com/apache/datafusion/pull/23260#discussion_r3548646336
##########
datafusion/expr/src/logical_plan/plan.rs:
##########
@@ -1672,34 +1672,68 @@ impl LogicalPlan {
let mut param_types: HashMap<String, Option<FieldRef>> =
HashMap::new();
self.apply_with_subqueries(|plan| {
+ let row_count_field = Arc::new(Field::new("", DataType::Int64,
true));
plan.apply_expressions(|expr| {
- expr.apply(|expr| {
- if let Expr::Placeholder(Placeholder { id, field }) = expr
{
- let prev = param_types.get(id);
- match (prev, field) {
- (Some(Some(prev)), Some(field)) => {
- check_metadata_with_storage_equal(
- (field.data_type(),
Some(field.metadata())),
- (prev.data_type(), Some(prev.metadata())),
- "parameter",
- &format!(": Conflicting types for id
{id}"),
- )?;
- }
- (_, Some(field)) => {
- param_types.insert(id.clone(),
Some(Arc::clone(field)));
- }
- _ => {
- param_types.insert(id.clone(), None);
- }
- }
- }
- Ok(TreeNodeRecursion::Continue)
- })
+ if matches!(plan, LogicalPlan::Limit(_)) {
+ let expr = Self::infer_limit_row_count_parameter_fields(
+ expr,
+ &row_count_field,
+ )?;
+ Self::collect_parameter_fields(&expr, &mut param_types)
+ } else {
+ Self::collect_parameter_fields(expr, &mut param_types)
+ }
})
})
.map(|_| param_types)
}
+ fn infer_limit_row_count_parameter_fields(
+ expr: &Expr,
+ row_count_field: &FieldRef,
+ ) -> Result<Expr> {
+ let empty_schema = DFSchema::empty();
+ let (expr, _) = expr.clone().infer_placeholder_types(&empty_schema)?;
+
+ if let Expr::Placeholder(Placeholder { id, field: None }) = expr {
+ return Ok(Expr::Placeholder(Placeholder::new_with_field(
+ id,
+ Some(Arc::clone(row_count_field)),
+ )));
+ }
+
+ Ok(expr)
+ }
+
+ fn collect_parameter_fields(
+ expr: &Expr,
+ param_types: &mut HashMap<String, Option<FieldRef>>,
+ ) -> Result<TreeNodeRecursion> {
+ expr.apply(|expr| {
+ if let Expr::Placeholder(Placeholder { id, field }) = expr {
+ let field = field.as_ref().map(Arc::clone);
+ let prev = param_types.get(id);
+ match (prev, &field) {
+ (Some(Some(prev)), Some(field)) => {
+ check_metadata_with_storage_equal(
+ (field.data_type(), Some(field.metadata())),
+ (prev.data_type(), Some(prev.metadata())),
+ "parameter",
+ &format!(": Conflicting types for id {id}"),
+ )?;
+ }
+ (_, Some(field)) => {
+ param_types.insert(id.clone(),
Some(Arc::clone(field)));
+ }
+ _ => {
Review Comment:
`collect_parameter_fields` can currently erase the `Int64` type inferred
from a LIMIT or OFFSET placeholder when the same parameter is later seen in an
untyped context.
For example, `SELECT $1 FROM person LIMIT $1` records `$1` as `Int64` from
the row-count expression first because `apply_with_subqueries` visits the
`Limit` node before its child plans. Later, the projection's bare `$1` reaches
this catch-all arm and overwrites that entry with `None`. As a result, the
metadata API still reports the row-count parameter as untyped, which breaks the
invariant this PR is trying to enforce when the placeholder is reused in an
unconstrained expression.
Could we keep an existing `Some(field)` when the new occurrence is `None`,
and only insert `None` if the parameter is not already known? Please also add a
regression test for a LIMIT or OFFSET placeholder reused in an
otherwise-untyped expression.
##########
datafusion/expr/src/logical_plan/plan.rs:
##########
@@ -1672,34 +1672,68 @@ impl LogicalPlan {
let mut param_types: HashMap<String, Option<FieldRef>> =
HashMap::new();
self.apply_with_subqueries(|plan| {
+ let row_count_field = Arc::new(Field::new("", DataType::Int64,
true));
Review Comment:
Small polish suggestion: consider hoisting `row_count_field` outside the
`apply_with_subqueries` closure so the `Int64` field is allocated once per
`get_parameter_fields` call instead of once per visited plan node.
This is not blocking, but I think it would also make the row-count
placeholder invariant a little easier to spot while reading the code.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
