diff --git a/datafusion/sql/src/unparser/rewrite.rs b/datafusion/sql/src/unparser/rewrite.rs
index 1662579b069cd..b651f3947aeac 100644
--- a/datafusion/sql/src/unparser/rewrite.rs
+++ b/datafusion/sql/src/unparser/rewrite.rs
@@ -24,6 +24,7 @@ use datafusion_common::{
     tree_node::{Transformed, TransformedResult, TreeNode},
     Result,
 };
+use datafusion_expr::tree_node::transform_sort_vec;
 use datafusion_expr::{Expr, LogicalPlan, Projection, Sort, SortExpr};
 use sqlparser::ast::Ident;
 
@@ -86,11 +87,8 @@ 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>> {
-    // TODO FIXME
-    //  this impl does't compile
-    /*
-    let sort_exprs = transform_sort_vec(exprs, |expr| {
+fn rewrite_sort_expr_for_union(exprs: Vec<SortExpr>) -> Result<Vec<SortExpr>> {
+    let sort_exprs = transform_sort_vec(exprs, &mut |expr| {
         expr.transform_up(|expr| {
             if let Expr::Column(mut col) = expr {
                 col.relation = None;
@@ -103,31 +101,6 @@ fn rewrite_sort_expr_for_union(_exprs: Vec<SortExpr>) -> Result<Vec<SortExpr>> {
     .data()?;
 
     Ok(sort_exprs)
-    */
-
-    /*
-        original code:
-
-    fn rewrite_sort_expr_for_union(exprs: Vec<Expr>) -> Result<Vec<Expr>> {
-        let sort_exprs: Vec<Expr> = exprs
-            .into_iter()
-            .map_until_stop_and_collect(|expr| {
-                expr.transform_up(|expr| {
-                    if let Expr::Column(mut col) = expr {
-                        col.relation = None;
-                        Ok(Transformed::yes(Expr::Column(col)))
-                    } else {
-                        Ok(Transformed::no(expr))
-                    }
-                })
-            })
-            .data()?;
-
-        Ok(sort_exprs)
-    }
-         */
-
-    panic!("rewrite_sort_expr_for_union to be implemented");
 }
 
 // Rewrite logic plan for query that order by columns are not in projections