Improve unproject_sort_expr to handle arbitrary expressions (#16127)

* Add failing test to demonstrate problem

* Improve `unproject_sort_expr` to handle arbitrary expressions (#83)

* Remove redundant return

Author: phillipleblanc

datafusion/sql/src/unparser/plan.rs

@@ -509,7 +509,7 @@ impl Unparser<'_> {
                     .expr
                     .iter()
                     .map(|sort_expr| {
-                        unproject_sort_expr(sort_expr, agg, sort.input.as_ref())
+                        unproject_sort_expr(sort_expr.clone(), agg, sort.input.as_ref())
                     })
                     .collect::<Result<Vec<_>>>()?;
 

datafusion/sql/src/unparser/utils.rs

@@ -270,51 +270,58 @@ fn find_window_expr<'a>(
         .find(|expr| expr.schema_name().to_string() == column_name)
 }
 
-/// Transforms a Column expression into the actual expression from aggregation or projection if found.
+/// Transforms all Column expressions in a sort expression into the actual expression from aggregation or projection if found.
 /// This is required because if an ORDER BY expression is present in an Aggregate or Select, it is replaced
 /// with a Column expression (e.g., "sum(catalog_returns.cr_net_loss)"). We need to transform it back to
 /// the actual expression, such as sum("catalog_returns"."cr_net_loss").
 pub(crate) fn unproject_sort_expr(
-    sort_expr: &SortExpr,
+    mut sort_expr: SortExpr,
     agg: Option<&Aggregate>,
     input: &LogicalPlan,
 ) -> Result<SortExpr> {
-    let mut sort_expr = sort_expr.clone();
-
-    // Remove alias if present, because ORDER BY cannot use aliases
-    if let Expr::Alias(alias) = &sort_expr.expr {
-        sort_expr.expr = *alias.expr.clone();
-    }
-
-    let Expr::Column(ref col_ref) = sort_expr.expr else {
-        return Ok(sort_expr);
-    };
+    sort_expr.expr = sort_expr
+        .expr
+        .transform(|sub_expr| {
+            match sub_expr {
+                // Remove alias if present, because ORDER BY cannot use aliases
+                Expr::Alias(alias) => Ok(Transformed::yes(*alias.expr)),
+                Expr::Column(col) => {
+                    if col.relation.is_some() {
+                        return Ok(Transformed::no(Expr::Column(col)));
+                    }
 
-    if col_ref.relation.is_some() {
-        return Ok(sort_expr);
-    };
+                    // In case of aggregation there could be columns containing aggregation functions we need to unproject
+                    if let Some(agg) = agg {
+                        if agg.schema.is_column_from_schema(&col) {
+                            return Ok(Transformed::yes(unproject_agg_exprs(
+                                Expr::Column(col),
+                                agg,
+                                None,
+                            )?));
+                        }
+                    }
 
-    // In case of aggregation there could be columns containing aggregation functions we need to unproject
-    if let Some(agg) = agg {
-        if agg.schema.is_column_from_schema(col_ref) {
-            let new_expr = unproject_agg_exprs(sort_expr.expr, agg, None)?;
-            sort_expr.expr = new_expr;
-            return Ok(sort_expr);
-        }
-    }
+                    // If SELECT and ORDER BY contain the same expression with a scalar function, the ORDER BY expression will
+                    // be replaced by a Column expression (e.g., "substr(customer.c_last_name, Int64(0), Int64(5))"), and we need
+                    // to transform it back to the actual expression.
+                    if let LogicalPlan::Projection(Projection { expr, schema, .. }) =
+                        input
+                    {
+                        if let Ok(idx) = schema.index_of_column(&col) {
+                            if let Some(Expr::ScalarFunction(scalar_fn)) = expr.get(idx) {
+                                return Ok(Transformed::yes(Expr::ScalarFunction(
+                                    scalar_fn.clone(),
+                                )));
+                            }
+                        }
+                    }
 
-    // If SELECT and ORDER BY contain the same expression with a scalar function, the ORDER BY expression will
-    // be replaced by a Column expression (e.g., "substr(customer.c_last_name, Int64(0), Int64(5))"), and we need
-    // to transform it back to the actual expression.
-    if let LogicalPlan::Projection(Projection { expr, schema, .. }) = input {
-        if let Ok(idx) = schema.index_of_column(col_ref) {
-            if let Some(Expr::ScalarFunction(scalar_fn)) = expr.get(idx) {
-                sort_expr.expr = Expr::ScalarFunction(scalar_fn.clone());
+                    Ok(Transformed::no(Expr::Column(col)))
+                }
+                _ => Ok(Transformed::no(sub_expr)),
             }
-        }
-        return Ok(sort_expr);
-    }
-
+        })
+        .map(|e| e.data)?;
     Ok(sort_expr)
 }
 

datafusion/sql/tests/cases/plan_to_sql.rs

@@ -1827,6 +1827,51 @@ fn test_order_by_to_sql_3() {
     );
 }
 
+#[test]
+fn test_complex_order_by_with_grouping() -> Result<()> {
+    let state = MockSessionState::default().with_aggregate_function(grouping_udaf());
+
+    let context = MockContextProvider { state };
+    let sql_to_rel = SqlToRel::new(&context);
+
+    // This SQL is based on a simplified version of the TPC-DS query 36.
+    let statement = Parser::new(&GenericDialect {})
+        .try_with_sql(
+            r#"SELECT
+            j1_id,
+            j1_string,
+            grouping(j1_id) + grouping(j1_string) as lochierarchy
+        FROM
+            j1
+        GROUP BY
+            ROLLUP (j1_id, j1_string)
+        ORDER BY
+            grouping(j1_id) + grouping(j1_string) DESC,
+            CASE
+                WHEN grouping(j1_id) + grouping(j1_string) = 0 THEN j1_id
+            END
+        LIMIT 100"#,
+        )?
+        .parse_statement()?;
+
+    let plan = sql_to_rel.sql_statement_to_plan(statement)?;
+    let unparser = Unparser::default();
+    let sql = unparser.plan_to_sql(&plan)?;
+    insta::with_settings!({
+        filters => vec![
+            // Force a deterministic order for the grouping pairs
+            (r#"grouping\(j1\.(?:j1_id|j1_string)\),\s*grouping\(j1\.(?:j1_id|j1_string)\)"#, "grouping(j1.j1_string), grouping(j1.j1_id)")
+        ],
+    }, {
+        assert_snapshot!(
+            sql,
+            @r#"SELECT j1.j1_id, j1.j1_string, lochierarchy FROM (SELECT j1.j1_id, j1.j1_string, (grouping(j1.j1_id) + grouping(j1.j1_string)) AS lochierarchy, grouping(j1.j1_string), grouping(j1.j1_id) FROM j1 GROUP BY ROLLUP (j1.j1_id, j1.j1_string) ORDER BY (grouping(j1.j1_id) + grouping(j1.j1_string)) DESC NULLS FIRST, CASE WHEN ((grouping(j1.j1_id) + grouping(j1.j1_string)) = 0) THEN j1.j1_id END ASC NULLS LAST) LIMIT 100"#
+        );
+    });
+
+    Ok(())
+}
+
 #[test]
 fn test_aggregation_to_sql() {
     let sql = r#"SELECT id, first_name,