diff --git a/datafusion/expr/src/expr_rewriter/order_by.rs b/datafusion/expr/src/expr_rewriter/order_by.rs index 0044b6cf6f377..6db95555502da 100644 --- a/datafusion/expr/src/expr_rewriter/order_by.rs +++ b/datafusion/expr/src/expr_rewriter/order_by.rs @@ -21,7 +21,9 @@ use crate::expr::Alias; use crate::expr_rewriter::normalize_col; use crate::{expr::Sort, Cast, Expr, LogicalPlan, TryCast}; -use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_common::tree_node::{ + Transformed, TransformedResult, TreeNode, TreeNodeRecursion, +}; use datafusion_common::{Column, Result}; /// Rewrite sort on aggregate expressions to sort on the column of aggregate output @@ -101,8 +103,18 @@ fn rewrite_in_terms_of_projection( let search_col = Expr::Column(Column::new_unqualified(name)); // look for the column named the same as this expr - if let Some(found) = proj_exprs.iter().find(|a| expr_match(&search_col, a)) { - let found = found.clone(); + let mut found = None; + for proj_expr in &proj_exprs { + proj_expr.apply(|e| { + if expr_match(&search_col, e) { + found = Some(e.clone()); + return Ok(TreeNodeRecursion::Stop); + } + Ok(TreeNodeRecursion::Continue) + })?; + } + + if let Some(found) = found { return Ok(Transformed::yes(match normalized_expr { Expr::Cast(Cast { expr: _, data_type }) => Expr::Cast(Cast { expr: Box::new(found), diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index b9502c1520049..74a5abf95c84a 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -2511,6 +2511,21 @@ fn select_groupby_orderby() { FROM person GROUP BY person.birth_date ORDER BY birth_date; "#; quick_test(sql, expected); + + // Use columnized `avg(age)` in the order by + let sql = r#"SELECT + avg(age) + avg(age), + date_trunc('month', person.birth_date) AS "birth_date" + FROM person GROUP BY person.birth_date ORDER BY avg(age) + avg(age); +"#; + + let expected = + "Sort: avg(person.age) + avg(person.age) ASC NULLS LAST\ + \n Projection: avg(person.age) + avg(person.age), date_trunc(Utf8(\"month\"), person.birth_date) AS birth_date\ + \n Aggregate: groupBy=[[person.birth_date]], aggr=[[avg(person.age)]]\ + \n TableScan: person"; + + quick_test(sql, expected); } fn logical_plan(sql: &str) -> Result { diff --git a/datafusion/sqllogictest/test_files/order.slt b/datafusion/sqllogictest/test_files/order.slt index 0f84171697257..383be5ee9152a 100644 --- a/datafusion/sqllogictest/test_files/order.slt +++ b/datafusion/sqllogictest/test_files/order.slt @@ -384,6 +384,19 @@ ORDER BY time; 2 2022-01-01T01:00:00 3 2022-01-02T00:00:00 +# Tests for https://github.com/apache/datafusion/issues/14459 +query PI +select + date_trunc('minute',time) AS "time", + sum(value) + sum(value) +FROM t +GROUP BY time +ORDER BY sum(value) + sum(value); +---- +2022-01-01T00:00:00 2 +2022-01-01T01:00:00 4 +2022-01-02T00:00:00 6 + ## SORT BY is not supported statement error DataFusion error: This feature is not implemented: SORT BY select * from t SORT BY time;