From e3cc3da7ca321db12e96114e90284043db580eed Mon Sep 17 00:00:00 2001
From: jussisaurio <jussi.saurio@gmail.com>
Date: Fri, 26 Jul 2024 14:32:58 +0300
Subject: [PATCH] Implement #216: order by column number

---
 core/translate/select.rs | 22 ++++++++++++++++++++--
 testing/orderby.test     |  5 +++++
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/core/translate/select.rs b/core/translate/select.rs
index 3cd9ebb2d..0b80c4b43 100644
--- a/core/translate/select.rs
+++ b/core/translate/select.rs
@@ -9,7 +9,7 @@ use crate::types::{OwnedRecord, OwnedValue};
 use crate::vdbe::{builder::ProgramBuilder, BranchOffset, Program};
 use crate::Result;
 
-use sqlite3_parser::ast::{self, JoinOperator, JoinType};
+use sqlite3_parser::ast::{self, JoinOperator, JoinType, ResultColumn};
 
 use std::rc::Rc;
 
@@ -305,7 +305,25 @@ pub fn translate_select(mut select: Select) -> Result<Program> {
             let start = program.next_free_register();
             for col in sort_columns.iter() {
                 let target = program.alloc_register();
-                translate_expr(&mut program, &select, &col.expr, target, None)?;
+                // if the ORDER BY expression is a number, interpret it as an 1-indexed column number
+                // otherwise, interpret it normally as an expression
+                let sort_col_expr = if let ast::Expr::Literal(ast::Literal::Numeric(num)) =
+                    &col.expr
+                {
+                    let column_number = num.parse::<usize>()?;
+                    if column_number == 0 {
+                        crate::bail_parse_error!("invalid column index: {}", column_number);
+                    }
+                    let maybe_result_column = select.columns.get(column_number - 1);
+                    match maybe_result_column {
+                        Some(ResultColumn::Expr(expr, _)) => expr,
+                        None => crate::bail_parse_error!("invalid column index: {}", column_number),
+                        _ => todo!(),
+                    }
+                } else {
+                    &col.expr
+                };
+                translate_expr(&mut program, &select, sort_col_expr, target, None)?;
             }
             let (_, result_cols_count) = translate_columns(&mut program, &select, None)?;
             sort_info
diff --git a/testing/orderby.test b/testing/orderby.test
index e5da7072b..ed59db1a4 100644
--- a/testing/orderby.test
+++ b/testing/orderby.test
@@ -45,3 +45,8 @@ do_execsql_test basic-order-by-and-limit-3 {
 do_execsql_test order-by-qualified {
     select u.first_name from users u order by u.first_name limit 1;
 } {Aaron}
+
+do_execsql_test order-by-column-number {
+    select first_name, last_name, age from users order by 3,2 limit 2;
+} {Teresa|Allen|1
+David|Baker|1}