Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion datafusion/sql/src/expr/order_by.rs
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,13 @@ impl<S: ContextProvider> 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)?;

Copy link
Copy Markdown
Contributor

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 BY expressions are now inferred, but the corresponding DISTINCT ON expressions are still planned through sql_expr_to_logical_expr without inference (datafusion/sql/src/select.rs:185-195). Since DistinctOn::with_sort_expr requires structural equality, these expressions no longer compare equal.

For example:

SELECT DISTINCT ON (
    CASE WHEN age < $1 THEN 'young' ELSE 'old' END
) first_name
FROM person
ORDER BY CASE WHEN age < $1 THEN 'young' ELSE 'old' END;

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 ON expressions after alias substitution and normalization? It would also be great to add this query as a regression test.

expr
}
};
sort_expr_vec.push(make_sort_expr(expr, asc, nulls_first));
Expand Down
12 changes: 10 additions & 2 deletions datafusion/sql/src/select.rs
Original file line number Diff line number Diff line change
Expand Up @@ -219,7 +219,10 @@ impl<S: ContextProvider> 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()?;

Expand Down Expand Up @@ -248,6 +251,8 @@ impl<S: ContextProvider> 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::<Result<Vec<Expr>>>()?
Expand Down Expand Up @@ -286,7 +291,10 @@ impl<S: ContextProvider> 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()?;

Expand Down
79 changes: 79 additions & 0 deletions datafusion/sql/tests/sql_integration.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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() {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The 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 $1 is inferred as the type of person.age. That would make the placeholder inference contract explicit instead of only verifying it indirectly through expression equality.

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";
Expand Down
Loading