dqkqd commented on code in PR #18831:
URL: https://github.com/apache/datafusion/pull/18831#discussion_r2552794225
##########
datafusion/sql/src/select.rs:
##########
@@ -988,11 +1012,42 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
None
};
+ // Rewrite the ORDER BY expressions to use the columns produced by the
+ // aggregation. If an ORDER BY expression matches a SELECT expression
+ // (ignoring aliases), use the SELECT's output column name to avoid
+ // duplication when the SELECT expression has an alias.
+ let order_by_post_aggr = order_by_exprs
+ .iter()
+ .map(|sort_expr| {
+ let rewritten_expr =
+ rebase_expr(&sort_expr.expr, &aggr_projection_exprs,
input)?;
+
+ // Check if this ORDER BY expression matches any aliased
SELECT expression
+ // If so, use the SELECT's alias instead of the raw expression
+ let final_expr = select_exprs_post_aggr
+ .iter()
+ .find_map(|select_expr| {
+ // Only consider aliased expressions
+ if let Expr::Alias(alias) = select_expr {
+ if alias.expr.as_ref() == &rewritten_expr {
+ // Use the alias name
+ return
Some(Expr::Column(alias.name.clone().into()));
+ }
+ }
+ None
+ })
+ .unwrap_or(rewritten_expr);
+
+ Ok(sort_expr.with_expr(final_expr))
+ })
+ .collect::<Result<Vec<SortExpr>>>()?;
+
Review Comment:
Should we add a `check_columns_satisfy_exprs` here for better UX?
```rust
check_columns_satisfy_exprs(
&order_by_post_aggr,
std::slice::from_ref(&order_by_post_aggr),
CheckColumnsSatisfyExprsPurpose::Aggregate(
CheckColumnsMustReferenceAggregatePurpose::OrderBy,
),
)?;
```
Might need a testcase too.
```
> SELECT column2 FROM foo GROUP BY column2 ORDER BY column1
Error during planning: Column in ORDER BY must be in GROUP BY or an
aggregate function
```
--
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]
