From c25ee4d6487577bfdd68a37452c2cda5a864131a Mon Sep 17 00:00:00 2001 From: Braedon Wooding Date: Sat, 1 Aug 2026 17:06:58 +1000 Subject: [PATCH] fix: infer placeholder types in GROUP BY, HAVING, QUALIFY and ORDER BY The SELECT list is planned by sql_to_expr, which infers placeholder types. These four clauses are planned by sql_expr_to_logical_expr, which does not, so the same expression written in both places does not compare equal. The result is that a grouping key containing a placeholder is never matched against the identical SELECT expression, and the columns inside it are reported as ungrouped. The same query with literals in place of the placeholder plans fine. QUALIFY fails differently, on a duplicate field name, because the typed and untyped spellings print alike but are not equal. Adds a planner test per clause. --- datafusion/sql/src/expr/order_by.rs | 8 ++- datafusion/sql/src/select.rs | 12 +++- datafusion/sql/tests/sql_integration.rs | 79 +++++++++++++++++++++++++ 3 files changed, 96 insertions(+), 3 deletions(-) diff --git a/datafusion/sql/src/expr/order_by.rs b/datafusion/sql/src/expr/order_by.rs index faecfbcfecc05..0067a1ebd708c 100644 --- a/datafusion/sql/src/expr/order_by.rs +++ b/datafusion/sql/src/expr/order_by.rs @@ -109,7 +109,13 @@ impl SqlToRel<'_, S> { )) } e => { - self.sql_expr_to_logical_expr(e, order_by_schema, planner_context)? + let expr = self.sql_expr_to_logical_expr( + e, + order_by_schema, + planner_context, + )?; + let (expr, _) = expr.infer_placeholder_types(order_by_schema)?; + expr } }; sort_expr_vec.push(make_sort_expr(expr, asc, nulls_first)); diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index bdab013144462..e41792ce4c14e 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -219,7 +219,10 @@ impl SqlToRel<'_, S> { // SELECT c1, MAX(c2) AS m FROM t GROUP BY c1 HAVING MAX(c2) > 10; // let having_expr = resolve_aliases_to_exprs(having_expr, &alias_map)?; - normalize_col(having_expr, &projected_plan) + let having_expr = normalize_col(having_expr, &projected_plan)?; + let (having_expr, _) = + having_expr.infer_placeholder_types(&combined_schema)?; + Ok(having_expr) }) .transpose()?; @@ -248,6 +251,8 @@ impl SqlToRel<'_, S> { base_plan.schema(), std::slice::from_ref(&group_by_expr), )?; + let (group_by_expr, _) = + group_by_expr.infer_placeholder_types(&combined_schema)?; Ok(group_by_expr) }) .collect::>>()? @@ -286,7 +291,10 @@ impl SqlToRel<'_, S> { // select row_number() over (PARTITION BY id) as rk from users qualify row_number() over (PARTITION BY id) > 1; // let qualify_expr = resolve_aliases_to_exprs(qualify_expr, &alias_map)?; - normalize_col(qualify_expr, &projected_plan) + let qualify_expr = normalize_col(qualify_expr, &projected_plan)?; + let (qualify_expr, _) = + qualify_expr.infer_placeholder_types(&combined_schema)?; + Ok(qualify_expr) }) .transpose()?; diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index a4bf0db910774..4b676bc84ddee 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -1504,6 +1504,85 @@ fn select_aggregate_with_group_by_with_having_using_count_star_not_in_select() { ); } +/// An expression containing a placeholder, written in both the SELECT list and +/// the GROUP BY, has to be recognised as one expression the way its literal +/// equivalent is. Otherwise the columns inside it read as ungrouped, because the +/// SELECT list has its placeholder types inferred and the grouping key does not. +#[test] +fn select_aggregate_with_group_by_placeholder_expression() { + let sql = "SELECT CASE WHEN age < $1 THEN 'young' ELSE 'old' END, count(*) + FROM person + GROUP BY CASE WHEN age < $1 THEN 'young' ELSE 'old' END"; + let plan = logical_plan(sql).unwrap(); + assert_snapshot!( + plan, + @r#" + Projection: CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END, count(*) + Aggregate: groupBy=[[CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END]], aggr=[[count(*)]] + TableScan: person + "# + ); +} + +/// The same, for a grouping expression repeated in HAVING. +#[test] +fn select_aggregate_with_having_placeholder_expression() { + let sql = "SELECT CASE WHEN age < $1 THEN 'young' ELSE 'old' END, count(*) + FROM person + GROUP BY CASE WHEN age < $1 THEN 'young' ELSE 'old' END + HAVING CASE WHEN age < $1 THEN 'young' ELSE 'old' END = 'young'"; + let plan = logical_plan(sql).unwrap(); + assert_snapshot!( + plan, + @r#" + Projection: CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END, count(*) + Filter: CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END = Utf8("young") + Aggregate: groupBy=[[CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END]], aggr=[[count(*)]] + TableScan: person + "# + ); +} + +/// The same, for a grouping expression repeated in ORDER BY. +#[test] +fn select_aggregate_with_order_by_placeholder_expression() { + let sql = "SELECT CASE WHEN age < $1 THEN 'young' ELSE 'old' END, count(*) + FROM person + GROUP BY CASE WHEN age < $1 THEN 'young' ELSE 'old' END + ORDER BY CASE WHEN age < $1 THEN 'young' ELSE 'old' END"; + let plan = logical_plan(sql).unwrap(); + assert_snapshot!( + plan, + @r#" + Sort: CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END ASC NULLS LAST + Projection: CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END, count(*) + Aggregate: groupBy=[[CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END]], aggr=[[count(*)]] + TableScan: person + "# + ); +} + +/// The same, for a window expression repeated in QUALIFY. Here the two spellings +/// of the window expression collide by name instead, since they print alike but +/// do not compare equal. +#[test] +fn select_window_with_qualify_placeholder_expression() { + let sql = "SELECT first_name, + row_number() OVER (PARTITION BY CASE WHEN age < $1 THEN 'young' ELSE 'old' END) + FROM person + QUALIFY row_number() OVER (PARTITION BY CASE WHEN age < $1 THEN 'young' ELSE 'old' END) = 1"; + let plan = logical_plan(sql).unwrap(); + assert_snapshot!( + plan, + @r#" + Projection: person.first_name, row_number() PARTITION BY [CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING + Filter: row_number() PARTITION BY [CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING = Int64(1) + WindowAggr: windowExpr=[[row_number() PARTITION BY [CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]] + TableScan: person + "# + ); +} + #[test] fn select_binary_expr() { let sql = "SELECT age + salary from person";