-
Notifications
You must be signed in to change notification settings - Fork 2.3k
fix: infer placeholder types in GROUP BY, HAVING, QUALIFY and ORDER BY (fix for #24042) #24043
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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() { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. These snapshot tests demonstrate that planning succeeds, which is great. One small suggestion would be to add a focused structural assertion that |
||
| 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"; | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Nice catch applying placeholder inference here. I think this introduces a regression for
DISTINCT ON, though.ORDER BYexpressions are now inferred, but the correspondingDISTINCT ONexpressions are still planned throughsql_expr_to_logical_exprwithout inference (datafusion/sql/src/select.rs:185-195). SinceDistinctOn::with_sort_exprrequires structural equality, these expressions no longer compare equal.For example:
This now fails with
SELECT DISTINCT ON expressions must match initial ORDER BY expressions, whereas before this change both expressions were untyped and matched.Could we also infer placeholder types for the
DISTINCT ONexpressions after alias substitution and normalization? It would also be great to add this query as a regression test.