fix #18683 by enabling sort by aggregate

Author: tshauck

datafusion/sql/src/select.rs

@@ -41,7 +41,7 @@ use datafusion_expr::utils::{
 };
 use datafusion_expr::{
     Aggregate, Expr, Filter, GroupingSet, LogicalPlan, LogicalPlanBuilder,
-    LogicalPlanBuilderOptions, Partitioning,
+    LogicalPlanBuilderOptions, Partitioning, SortExpr,
 };
 
 use indexmap::IndexMap;
@@ -219,11 +219,13 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
             .transpose()?;
 
         // The outer expressions we will search through for aggregates.
-        // Aggregates may be sourced from the SELECT list or from the HAVING expression.
+        // Aggregates may be sourced from the SELECT list, HAVING, QUALIFY, or ORDER BY.
         let aggr_expr_haystack = select_exprs
             .iter()
             .chain(having_expr_opt.iter())
-            .chain(qualify_expr_opt.iter());
+            .chain(qualify_expr_opt.iter())
+            .chain(order_by_rex.iter().map(|s| &s.expr));
+
         // All of the aggregate expressions (deduplicated).
         let aggr_exprs = find_aggregate_exprs(aggr_expr_haystack);
 
@@ -233,19 +235,21 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
             mut select_exprs_post_aggr,
             having_expr_post_aggr,
             qualify_expr_post_aggr,
+            order_by_rex,
         ) = if !group_by_exprs.is_empty() || !aggr_exprs.is_empty() {
             self.aggregate(
                 &base_plan,
                 &select_exprs,
                 having_expr_opt.as_ref(),
                 qualify_expr_opt.as_ref(),
+                &order_by_rex,
                 &group_by_exprs,
                 &aggr_exprs,
             )?
         } else {
             match having_expr_opt {
                 Some(having_expr) => return plan_err!("HAVING clause references: {having_expr} must appear in the GROUP BY clause or be used in an aggregate function"),
-                None => (base_plan.clone(), select_exprs.clone(), having_expr_opt, qualify_expr_opt)
+                None => (base_plan.clone(), select_exprs.clone(), having_expr_opt, qualify_expr_opt, order_by_rex)
             }
         };
 
@@ -872,16 +876,25 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
     ///   the aggregate
     /// * `qualify_expr_post_aggr`  - The "qualify" expression rewritten to reference a column from
     ///   the aggregate
-    #[allow(clippy::type_complexity)]
+    /// * `order_by_post_aggr`     - The ORDER BY expressions rewritten to reference columns from
+    ///   the aggregate
+    #[allow(clippy::type_complexity, clippy::too_many_arguments)]
     fn aggregate(
         &self,
         input: &LogicalPlan,
         select_exprs: &[Expr],
         having_expr_opt: Option<&Expr>,
         qualify_expr_opt: Option<&Expr>,
+        order_by_exprs: &[SortExpr],
         group_by_exprs: &[Expr],
         aggr_exprs: &[Expr],
-    ) -> Result<(LogicalPlan, Vec<Expr>, Option<Expr>, Option<Expr>)> {
+    ) -> Result<(
+        LogicalPlan,
+        Vec<Expr>,
+        Option<Expr>,
+        Option<Expr>,
+        Vec<SortExpr>,
+    )> {
         // create the aggregate plan
         let options =
             LogicalPlanBuilderOptions::new().with_add_implicit_group_by_exprs(true);
@@ -988,11 +1001,23 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
             None
         };
 
+        // Rewrite the ORDER BY expressions to use the columns produced by the
+        // aggregation.
+        let order_by_post_aggr = order_by_exprs
+            .iter()
+            .map(|sort_expr| {
+                let rewritten_expr =
+                    rebase_expr(&sort_expr.expr, &aggr_projection_exprs, input)?;
+                Ok(sort_expr.with_expr(rewritten_expr))
+            })
+            .collect::<Result<Vec<SortExpr>>>()?;
+
         Ok((
             plan,
             select_exprs_post_aggr,
             having_expr_post_aggr,
             qualify_expr_post_aggr,
+            order_by_post_aggr,
         ))
     }
 

datafusion/sql/tests/sql_integration.rs

@@ -3257,6 +3257,55 @@ Sort: avg(person.age) + avg(person.age) ASC NULLS LAST
     );
 }
 
+#[test]
+fn select_groupby_orderby_aggregate_on_non_selected_column() {
+    let sql = "SELECT state FROM person GROUP BY state ORDER BY MIN(age)";
+    let plan = logical_plan(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+Projection: person.state
+  Sort: min(person.age) ASC NULLS LAST
+    Projection: person.state, min(person.age)
+      Aggregate: groupBy=[[person.state]], aggr=[[min(person.age)]]
+        TableScan: person
+"#
+    );
+}
+
+#[test]
+fn select_groupby_orderby_multiple_aggregates_on_non_selected_columns() {
+    let sql =
+        "SELECT state FROM person GROUP BY state ORDER BY MIN(age), MAX(salary) DESC";
+    let plan = logical_plan(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+Projection: person.state
+  Sort: min(person.age) ASC NULLS LAST, max(person.salary) DESC NULLS FIRST
+    Projection: person.state, min(person.age), max(person.salary)
+      Aggregate: groupBy=[[person.state]], aggr=[[min(person.age), max(person.salary)]]
+        TableScan: person
+"#
+    );
+}
+
+#[test]
+fn select_groupby_orderby_aggregate_on_non_selected_column_original_issue() {
+    let sql = "SELECT id FROM person GROUP BY id ORDER BY min(age)";
+    let plan = logical_plan(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+Projection: person.id
+  Sort: min(person.age) ASC NULLS LAST
+    Projection: person.id, min(person.age)
+      Aggregate: groupBy=[[person.id]], aggr=[[min(person.age)]]
+        TableScan: person
+"#
+    );
+}
+
 fn logical_plan(sql: &str) -> Result<LogicalPlan> {
     logical_plan_with_options(sql, ParserOptions::default())
 }