From f2a254f23cc6aaf77d98ef5f3c76dbc44bb348e0 Mon Sep 17 00:00:00 2001
From: Piotr Findeisen <piotr.findeisen@gmail.com>
Date: Mon, 26 Aug 2024 19:02:18 +0200
Subject: [PATCH] fixup! fixup! fixup! WIP

---
 datafusion/expr/src/logical_plan/plan.rs      | 19 ++++++++++---------
 datafusion/expr/src/logical_plan/tree_node.rs |  2 +-
 datafusion/expr/src/tree_node.rs              |  4 ++--
 .../optimizer/src/common_subexpr_eliminate.rs |  1 -
 datafusion/sql/src/unparser/expr.rs           |  2 +-
 datafusion/sql/src/unparser/rewrite.rs        |  2 +-
 6 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs
index 74f860fa91b4c..debc18e035435 100644
--- a/datafusion/expr/src/logical_plan/plan.rs
+++ b/datafusion/expr/src/logical_plan/plan.rs
@@ -54,6 +54,7 @@ use datafusion_common::{
 // backwards compatibility
 use crate::display::PgJsonVisitor;
 use crate::logical_plan::tree_node::unwrap_arc;
+use crate::tree_node::replace_sort_expressions;
 pub use datafusion_common::display::{PlanType, StringifiedPlan, ToStringifiedPlan};
 pub use datafusion_common::{JoinConstraint, JoinType};
 
@@ -887,15 +888,15 @@ impl LogicalPlan {
                 Aggregate::try_new(Arc::new(inputs.swap_remove(0)), expr, agg_expr)
                     .map(LogicalPlan::Aggregate)
             }
-            LogicalPlan::Sort(Sort { /*fetch,*/ .. }) => {
-                // TODO FIXME
-                // LogicalPlan::Sort(Sort { fetch, .. }) => Ok(LogicalPlan::Sort(Sort {
-                //                 expr,
-                //                 input: Arc::new(inputs.swap_remove(0)),
-                //                 fetch: *fetch,
-                //             })),
-                internal_err!("with_ew_exprs not implemented for Sort")
-            },
+            LogicalPlan::Sort(Sort {
+                expr: sort_expr,
+                fetch,
+                ..
+            }) => Ok(LogicalPlan::Sort(Sort {
+                expr: replace_sort_expressions(sort_expr.clone(), expr),
+                input: Arc::new(inputs.swap_remove(0)),
+                fetch: *fetch,
+            })),
             LogicalPlan::Join(Join {
                 join_type,
                 join_constraint,
diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs
index 297e8e4a01c02..82f49c5f782aa 100644
--- a/datafusion/expr/src/logical_plan/tree_node.rs
+++ b/datafusion/expr/src/logical_plan/tree_node.rs
@@ -491,7 +491,7 @@ impl LogicalPlan {
                     .visit_sibling(|| filter.iter().apply_until_stop(f))
             }
             LogicalPlan::Sort(Sort { expr, .. }) => {
-                expr.iter().apply_until_stop(|sort| f(&*sort.expr))
+                expr.iter().apply_until_stop(|sort| f(&sort.expr))
             }
             LogicalPlan::Extension(extension) => {
                 // would be nice to avoid this copy -- maybe can
diff --git a/datafusion/expr/src/tree_node.rs b/datafusion/expr/src/tree_node.rs
index 8d1107c062330..69748aded531f 100644
--- a/datafusion/expr/src/tree_node.rs
+++ b/datafusion/expr/src/tree_node.rs
@@ -97,7 +97,7 @@ impl TreeNode for Expr {
                     expr_vec.push(f.as_ref());
                 }
                 if let Some(order_by) = order_by {
-                    expr_vec.extend(order_by.into_iter().map(|sort| sort.expr.as_ref()));
+                    expr_vec.extend(order_by.iter().map(|sort| sort.expr.as_ref()));
                 }
                 expr_vec
             }
@@ -109,7 +109,7 @@ impl TreeNode for Expr {
             }) => {
                 let mut expr_vec = args.iter().collect::<Vec<_>>();
                 expr_vec.extend(partition_by);
-                expr_vec.extend(order_by.into_iter().map(|sort| sort.expr.as_ref()));
+                expr_vec.extend(order_by.iter().map(|sort| sort.expr.as_ref()));
                 expr_vec
             }
             Expr::InList(InList { expr, list, .. }) => {
diff --git a/datafusion/optimizer/src/common_subexpr_eliminate.rs b/datafusion/optimizer/src/common_subexpr_eliminate.rs
index 7d9f1eedc13a7..81f29a3c71b2e 100644
--- a/datafusion/optimizer/src/common_subexpr_eliminate.rs
+++ b/datafusion/optimizer/src/common_subexpr_eliminate.rs
@@ -886,7 +886,6 @@ enum ExprMask {
     /// - [`Columns`](Expr::Column)
     /// - [`ScalarVariable`](Expr::ScalarVariable)
     /// - [`Alias`](Expr::Alias)
-    /// - [`Sort`](Expr::Sort)
     /// - [`Wildcard`](Expr::Wildcard)
     /// - [`AggregateFunction`](Expr::AggregateFunction)
     Normal,
diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs
index 53a6e75cc85c3..a61a0e095d3ce 100644
--- a/datafusion/sql/src/unparser/expr.rs
+++ b/datafusion/sql/src/unparser/expr.rs
@@ -253,7 +253,7 @@ impl Unparser<'_> {
                 };
                 let order_by: Vec<ast::OrderByExpr> = order_by
                     .iter()
-                    .map(|expr| sort_to_sql(expr))
+                    .map(sort_to_sql)
                     .collect::<Result<Vec<_>>>()?;
 
                 let start_bound = self.convert_bound(&window_frame.start_bound)?;
diff --git a/datafusion/sql/src/unparser/rewrite.rs b/datafusion/sql/src/unparser/rewrite.rs
index 715d3aca68044..1662579b069cd 100644
--- a/datafusion/sql/src/unparser/rewrite.rs
+++ b/datafusion/sql/src/unparser/rewrite.rs
@@ -86,7 +86,7 @@ pub(super) fn normalize_union_schema(plan: &LogicalPlan) -> Result<LogicalPlan>
 }
 
 /// Rewrite sort expressions that have a UNION plan as their input to remove the table reference.
-fn rewrite_sort_expr_for_union(exprs: Vec<SortExpr>) -> Result<Vec<SortExpr>> {
+fn rewrite_sort_expr_for_union(_exprs: Vec<SortExpr>) -> Result<Vec<SortExpr>> {
     // TODO FIXME
     //  this impl does't compile
     /*