From eca196a54b265d93e3e51d90d47f71b742d2c808 Mon Sep 17 00:00:00 2001
From: Jussi Saurio <jussi.saurio@gmail.com>
Date: Fri, 14 Feb 2025 16:53:37 +0200
Subject: [PATCH] Support numeric column references in GROUP BY

---
 core/translate/select.rs | 56 ++++++++++++++++++++++++----------------
 testing/groupby.test     |  4 +++
 2 files changed, 38 insertions(+), 22 deletions(-)

diff --git a/core/translate/select.rs b/core/translate/select.rs
index 9401f4f0b..fe7656f16 100644
--- a/core/translate/select.rs
+++ b/core/translate/select.rs
@@ -300,6 +300,7 @@ pub fn prepare_select_plan<'a>(
 
             if let Some(mut group_by) = group_by {
                 for expr in group_by.exprs.iter_mut() {
+                    replace_column_number_with_copy_of_column_expr(expr, &plan.result_columns)?;
                     bind_column_references(
                         expr,
                         &plan.table_references,
@@ -343,35 +344,21 @@ pub fn prepare_select_plan<'a>(
             if let Some(order_by) = select.order_by {
                 let mut key = Vec::new();
 
-                for o in order_by {
-                    // if the ORDER BY expression is a number, interpret it as an 1-indexed column number
-                    // otherwise, interpret it normally as an expression
-                    let mut expr = if let ast::Expr::Literal(ast::Literal::Numeric(num)) = &o.expr {
-                        let column_number = num.parse::<usize>()?;
-                        if column_number == 0 {
-                            crate::bail_parse_error!("invalid column index: {}", column_number);
-                        }
-                        let maybe_result_column = columns.get(column_number - 1);
-                        match maybe_result_column {
-                            Some(ResultColumn::Expr(e, _)) => e.clone(),
-                            None => {
-                                crate::bail_parse_error!("invalid column index: {}", column_number)
-                            }
-                            _ => todo!(),
-                        }
-                    } else {
-                        o.expr
-                    };
+                for mut o in order_by {
+                    replace_column_number_with_copy_of_column_expr(
+                        &mut o.expr,
+                        &plan.result_columns,
+                    )?;
 
                     bind_column_references(
-                        &mut expr,
+                        &mut o.expr,
                         &plan.table_references,
                         Some(&plan.result_columns),
                     )?;
-                    resolve_aggregates(&expr, &mut plan.aggregates);
+                    resolve_aggregates(&o.expr, &mut plan.aggregates);
 
                     key.push((
-                        expr,
+                        o.expr,
                         o.order.map_or(Direction::Ascending, |o| match o {
                             ast::SortOrder::Asc => Direction::Ascending,
                             ast::SortOrder::Desc => Direction::Descending,
@@ -392,6 +379,31 @@ pub fn prepare_select_plan<'a>(
     }
 }
 
+/// Replaces a column number in an ORDER BY or GROUP BY expression with a copy of the column expression.
+/// For example, in SELECT u.first_name, count(1) FROM users u GROUP BY 1 ORDER BY 2,
+/// the column number 1 is replaced with u.first_name and the column number 2 is replaced with count(1).
+fn replace_column_number_with_copy_of_column_expr(
+    order_by_or_group_by_expr: &mut ast::Expr,
+    columns: &[ResultSetColumn],
+) -> Result<()> {
+    if let ast::Expr::Literal(ast::Literal::Numeric(num)) = order_by_or_group_by_expr {
+        let column_number = num.parse::<usize>()?;
+        if column_number == 0 {
+            crate::bail_parse_error!("invalid column index: {}", column_number);
+        }
+        let maybe_result_column = columns.get(column_number - 1);
+        match maybe_result_column {
+            Some(ResultSetColumn { expr, .. }) => {
+                *order_by_or_group_by_expr = expr.clone();
+            }
+            None => {
+                crate::bail_parse_error!("invalid column index: {}", column_number)
+            }
+        };
+    }
+    Ok(())
+}
+
 fn count_plan_required_cursors(plan: &SelectPlan) -> usize {
     let num_table_cursors: usize = plan
         .table_references
diff --git a/testing/groupby.test b/testing/groupby.test
index c370c3df2..9fd6e51bf 100644
--- a/testing/groupby.test
+++ b/testing/groupby.test
@@ -181,3 +181,7 @@ do_execsql_test column_alias_in_group_by_order_by_having {
   select first_name as fn, count(1) as fn_count from users where fn in ('Wanda', 'Whitney', 'William') group by fn having fn_count > 10 order by fn_count;
 } {Whitney|11
 William|111}
+
+do_execsql_test group_by_column_number {
+  select u.first_name, count(1) from users u group by 1 limit 1;
+} {Aaron|41}