fix: order by expr rewrite fix (#14486)

* fix: rewrite order by on compound expr

* fmt

* revert transform

* test fix

Author: akoshchiy

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),

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<LogicalPlan> {

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;