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";