Merge 'translate: implement Sequence opcode and fix sort order' from Preston Thorpe

This PR implements the `Sequence` and `SequenceTest` opcodes, although
does not yet add plumbing to emit the latter.
SQLite has two distinct mechanisms that determine the final row order
with aggregates:
Traversal order of GROUP BY, and ORDER BY tiebreaking. When ORDER BY
contains only aggregate expressions and/or constants, SQLite has no
extra tiebreak key, but when ORDER BY mixes aggregate and non-aggregate
terms, SQLite adds an implicit, stable row `sequence` so “ties” respect
the input order.
This PR also fixes an issue with a query like the following:
```sql
SELECT u.first_name, COUNT(*) AS c
FROM users u
JOIN orders o ON o.user_id = u.id
GROUP BY u.first_name
ORDER BY c DESC;
```
Because ORDER BY has only an aggregate (COUNT(*) DESC) and no non-
aggregate terms, SQLite traverses the group key (u.first_name) in DESC
order in this case, so ties on c naturally appear with group keys in
descending order.
Previously tursodb would return the group key sorted in ASC order,
because it was used in all cases as the default

Closes #3287

core/translate/emitter.rs

@@ -307,6 +307,8 @@ pub fn emit_query<'a>(
             &plan.result_columns,
             &plan.order_by,
             &plan.table_references,
+            plan.group_by.is_some(),
+            &plan.aggregates,
         )?;
     }
 

core/translate/group_by.rs

@@ -1,4 +1,4 @@
-use turso_parser::ast;
+use turso_parser::ast::{self, SortOrder};
 
 use super::{
     emitter::TranslateCtx,
@@ -7,9 +7,16 @@ use super::{
     plan::{Distinctness, GroupBy, SelectPlan},
     result_row::emit_select_result,
 };
-use crate::translate::aggregation::{translate_aggregation_step, AggArgumentSource};
-use crate::translate::expr::{walk_expr, WalkControl};
 use crate::translate::plan::ResultSetColumn;
+use crate::translate::{
+    aggregation::{translate_aggregation_step, AggArgumentSource},
+    plan::Aggregate,
+};
+use crate::translate::{
+    emitter::Resolver,
+    expr::{walk_expr, WalkControl},
+    optimizer::Optimizable,
+};
 use crate::{
     schema::PseudoCursorType,
     translate::collate::CollationSeq,
@@ -231,6 +238,71 @@ pub fn init_group_by<'a>(
     Ok(())
 }
 
+/// Returns whether an ORDER BY expression should be treated as an
+/// aggregate-position term for the purposes of tie-ordering.
+///
+/// We classify an ORDER BY term as "aggregate or constant" when:
+/// it is syntactically equivalent to one of the finalized aggregate
+/// expressions for this SELECT (`COUNT(*)`, `SUM(col)`, `MAX(price)`), or
+/// it is a constant literal
+///
+/// Why this matters:
+/// When ORDER BY consists only of aggregates and/or constants, SQLite relies
+/// on the stability of the ORDER BY sorter to preserve the traversal order
+/// of groups established by GROUP BY iteration, and no extra tiebreak
+/// `Sequence` column is appended
+pub fn is_orderby_agg_or_const(resolver: &Resolver, e: &ast::Expr, aggs: &[Aggregate]) -> bool {
+    if aggs
+        .iter()
+        .any(|agg| exprs_are_equivalent(&agg.original_expr, e))
+    {
+        return true;
+    }
+    e.is_constant(resolver)
+}
+
+/// Computes the traversal order of GROUP BY keys so that the final
+/// ORDER BY matches SQLite's tie-breaking semantics.
+///
+/// If there are no GROUP BY keys or no ORDER BY terms, all keys default to ascending.
+///
+/// If *every* ORDER BY term is an aggregate or a constant then we mirror the
+/// direction of the first ORDER BY term across all GROUP BY keys.
+///
+/// Otherwise (mixed ORDER BY: at least one non-aggregate, non-constant term),
+/// we try to mirror explicit directions for any GROUP BY expression that
+/// appears in ORDER BY, and the remaining keys default to `ASC`.
+pub fn compute_group_by_sort_order(
+    group_by_exprs: &[ast::Expr],
+    order_by: &[(Box<ast::Expr>, SortOrder)],
+    aggs: &[Aggregate],
+    resolver: &Resolver,
+) -> Vec<SortOrder> {
+    let groupby_len = group_by_exprs.len();
+    if groupby_len == 0 || order_by.is_empty() {
+        return vec![SortOrder::Asc; groupby_len];
+    }
+    let only_agg_or_const = order_by
+        .iter()
+        .all(|(e, _)| is_orderby_agg_or_const(resolver, e, aggs));
+
+    if only_agg_or_const {
+        let first_direction = order_by[0].1;
+        return vec![first_direction; groupby_len];
+    }
+
+    let mut result = vec![SortOrder::Asc; groupby_len];
+    for (idx, groupby_expr) in group_by_exprs.iter().enumerate() {
+        if let Some((_, direction)) = order_by
+            .iter()
+            .find(|(expr, _)| exprs_are_equivalent(expr, groupby_expr))
+        {
+            result[idx] = *direction;
+        }
+    }
+    result
+}
+
 fn collect_non_aggregate_expressions<'a>(
     non_aggregate_expressions: &mut Vec<(&'a ast::Expr, bool)>,
     group_by: &'a GroupBy,
@@ -736,15 +808,7 @@ pub fn group_by_emit_row_phase<'a>(
             )?;
         }
         false => {
-            order_by_sorter_insert(
-                program,
-                &t_ctx.resolver,
-                t_ctx
-                    .meta_sort
-                    .as_ref()
-                    .expect("sort metadata must exist for ORDER BY"),
-                plan,
-            )?;
+            order_by_sorter_insert(program, t_ctx, plan)?;
         }
     }
 

core/translate/main_loop.rs

@@ -860,15 +860,7 @@ fn emit_loop_source(
             Ok(())
         }
         LoopEmitTarget::OrderBySorter => {
-            order_by_sorter_insert(
-                program,
-                &t_ctx.resolver,
-                t_ctx
-                    .meta_sort
-                    .as_ref()
-                    .expect("sort metadata must exist for ORDER BY"),
-                plan,
-            )?;
+            order_by_sorter_insert(program, t_ctx, plan)?;
 
             if let Distinctness::Distinct { ctx } = &plan.distinctness {
                 let distinct_ctx = ctx.as_ref().expect("distinct context must exist");

core/translate/order_by.rs

@@ -3,7 +3,7 @@ use turso_parser::ast::{self, SortOrder};
 use crate::{
     emit_explain,
     schema::PseudoCursorType,
-    translate::collate::CollationSeq,
+    translate::{collate::CollationSeq, group_by::is_orderby_agg_or_const, plan::Aggregate},
     util::exprs_are_equivalent,
     vdbe::{
         builder::{CursorType, ProgramBuilder},
@@ -13,7 +13,7 @@ use crate::{
 };
 
 use super::{
-    emitter::{Resolver, TranslateCtx},
+    emitter::TranslateCtx,
     expr::translate_expr,
     plan::{Distinctness, ResultSetColumn, SelectPlan, TableReferences},
     result_row::{emit_offset, emit_result_row_and_limit},
@@ -30,6 +30,11 @@ pub struct SortMetadata {
     // This vector holds the indexes of the result columns in the ORDER BY sorter.
     // This vector must be the same length as the result columns.
     pub remappings: Vec<OrderByRemapping>,
+    /// Whether we append an extra ascending "Sequence" key to the ORDER BY sort keys.
+    /// This is used *only* when a GROUP BY is present *and* ORDER BY is not purely
+    /// aggregates/constants, so that rows that tie on ORDER BY terms are output in
+    /// the same relative order the underlying row stream produced them.
+    pub has_sequence: bool,
 }
 
 /// Initialize resources needed for ORDER BY processing
@@ -39,12 +44,21 @@ pub fn init_order_by(
     result_columns: &[ResultSetColumn],
     order_by: &[(Box<ast::Expr>, SortOrder)],
     referenced_tables: &TableReferences,
+    has_group_by: bool,
+    aggregates: &[Aggregate],
 ) -> Result<()> {
     let sort_cursor = program.alloc_cursor_id(CursorType::Sorter);
+    let only_aggs = order_by
+        .iter()
+        .all(|(e, _)| is_orderby_agg_or_const(&t_ctx.resolver, e, aggregates));
+
+    // only emit sequence column if we have GROUP BY and ORDER BY is not only aggregates or constants
+    let has_sequence = has_group_by && !only_aggs;
     t_ctx.meta_sort = Some(SortMetadata {
         sort_cursor,
         reg_sorter_data: program.alloc_register(),
-        remappings: order_by_deduplicate_result_columns(order_by, result_columns),
+        remappings: order_by_deduplicate_result_columns(order_by, result_columns, has_sequence),
+        has_sequence,
     });
 
     /*
@@ -54,7 +68,7 @@ pub fn init_order_by(
      * then the collating sequence of the column is used to determine sort order.
      * If the expression is not a column and has no COLLATE clause, then the BINARY collating sequence is used.
      */
-    let collations = order_by
+    let mut collations = order_by
         .iter()
         .map(|(expr, _)| match expr.as_ref() {
             ast::Expr::Collate(_, collation_name) => {
@@ -72,10 +86,25 @@ pub fn init_order_by(
             _ => Ok(Some(CollationSeq::default())),
         })
         .collect::<Result<Vec<_>>>()?;
+
+    if has_sequence {
+        // sequence column uses BINARY collation
+        collations.push(Some(CollationSeq::default()));
+    }
+
+    let key_len = order_by.len() + if has_sequence { 1 } else { 0 };
+
     program.emit_insn(Insn::SorterOpen {
         cursor_id: sort_cursor,
-        columns: order_by.len(),
-        order: order_by.iter().map(|(_, direction)| *direction).collect(),
+        columns: key_len,
+        order: {
+            let mut ord: Vec<SortOrder> = order_by.iter().map(|(_, d)| *d).collect();
+            if has_sequence {
+                // sequence is ascending tiebreaker
+                ord.push(SortOrder::Asc);
+            }
+            ord
+        },
         collations,
     });
     Ok(())
@@ -98,9 +127,12 @@ pub fn emit_order_by(
         sort_cursor,
         reg_sorter_data,
         ref remappings,
+        has_sequence,
     } = *t_ctx.meta_sort.as_ref().unwrap();
-    let sorter_column_count =
-        order_by.len() + remappings.iter().filter(|r| !r.deduplicated).count();
+
+    let sorter_column_count = order_by.len()
+        + if has_sequence { 1 } else { 0 }
+        + remappings.iter().filter(|r| !r.deduplicated).count();
 
     // TODO: we need to know how many indices used for sorting
     // to emit correct explain output.
@@ -142,10 +174,11 @@ pub fn emit_order_by(
     let start_reg = t_ctx.reg_result_cols_start.unwrap();
     for i in 0..result_columns.len() {
         let reg = start_reg + i;
-        let column_idx = remappings
+        let remapping = remappings
             .get(i)
-            .expect("remapping must exist for all result columns")
-            .orderby_sorter_idx;
+            .expect("remapping must exist for all result columns");
+
+        let column_idx = remapping.orderby_sorter_idx;
         program.emit_column_or_rowid(cursor_id, column_idx, reg);
     }
 
@@ -170,10 +203,11 @@ pub fn emit_order_by(
 /// Emits the bytecode for inserting a row into an ORDER BY sorter.
 pub fn order_by_sorter_insert(
     program: &mut ProgramBuilder,
-    resolver: &Resolver,
-    sort_metadata: &SortMetadata,
+    t_ctx: &TranslateCtx,
     plan: &SelectPlan,
 ) -> Result<()> {
+    let resolver = &t_ctx.resolver;
+    let sort_metadata = t_ctx.meta_sort.as_ref().expect("sort metadata must exist");
     let order_by = &plan.order_by;
     let order_by_len = order_by.len();
     let result_columns = &plan.result_columns;
@@ -185,19 +219,49 @@ pub fn order_by_sorter_insert(
 
     // The ORDER BY sorter has the sort keys first, then the result columns.
     let orderby_sorter_column_count =
-        order_by_len + result_columns.len() - result_columns_to_skip_len;
+        order_by_len + if sort_metadata.has_sequence { 1 } else { 0 } + result_columns.len()
+            - result_columns_to_skip_len;
+
     let start_reg = program.alloc_registers(orderby_sorter_column_count);
     for (i, (expr, _)) in order_by.iter().enumerate() {
         let key_reg = start_reg + i;
-        translate_expr(
-            program,
-            Some(&plan.table_references),
-            expr,
-            key_reg,
-            resolver,
-        )?;
+
+        // Check if this ORDER BY expression matches a finalized aggregate
+        if let Some(agg_idx) = plan
+            .aggregates
+            .iter()
+            .position(|agg| exprs_are_equivalent(&agg.original_expr, expr))
+        {
+            // This ORDER BY expression is an aggregate, so copy from register
+            let agg_start_reg = t_ctx
+                .reg_agg_start
+                .expect("aggregate registers must be initialized");
+            let src_reg = agg_start_reg + agg_idx;
+            program.emit_insn(Insn::Copy {
+                src_reg,
+                dst_reg: key_reg,
+                extra_amount: 0,
+            });
+        } else {
+            // Not an aggregate, translate normally
+            translate_expr(
+                program,
+                Some(&plan.table_references),
+                expr,
+                key_reg,
+                resolver,
+            )?;
+        }
     }
     let mut cur_reg = start_reg + order_by_len;
+    if sort_metadata.has_sequence {
+        program.emit_insn(Insn::Sequence {
+            cursor_id: sort_metadata.sort_cursor,
+            target_reg: cur_reg,
+        });
+        cur_reg += 1;
+    }
+
     for (i, rc) in result_columns.iter().enumerate() {
         // If the result column is an exact duplicate of a sort key, we skip it.
         if sort_metadata
@@ -251,12 +315,12 @@ pub fn order_by_sorter_insert(
             let reordered_start_reg = program.alloc_registers(result_columns.len());
 
             for (select_idx, _rc) in result_columns.iter().enumerate() {
-                let src_reg = sort_metadata
+                let remapping = sort_metadata
                     .remappings
                     .get(select_idx)
-                    .map(|r| start_reg + r.orderby_sorter_idx)
                     .expect("remapping must exist for all result columns");
 
+                let src_reg = start_reg + remapping.orderby_sorter_idx;
                 let dst_reg = reordered_start_reg + select_idx;
 
                 program.emit_insn(Insn::Copy {
@@ -336,9 +400,13 @@ pub struct OrderByRemapping {
 pub fn order_by_deduplicate_result_columns(
     order_by: &[(Box<ast::Expr>, SortOrder)],
     result_columns: &[ResultSetColumn],
+    has_sequence: bool,
 ) -> Vec<OrderByRemapping> {
     let mut result_column_remapping: Vec<OrderByRemapping> = Vec::new();
     let order_by_len = order_by.len();
+    // `sequence_offset` shifts the base index where non-deduped SELECT columns begin,
+    // because Sequence sits after ORDER BY keys but before result columns.
+    let sequence_offset = if has_sequence { 1 } else { 0 };
 
     let mut i = 0;
     for rc in result_columns.iter() {
@@ -356,7 +424,7 @@ pub fn order_by_deduplicate_result_columns(
             // index comes after all ORDER BY entries (hence the +order_by_len). The
             // counter `i` tracks how many such non-duplicate result columns we've seen.
             result_column_remapping.push(OrderByRemapping {
-                orderby_sorter_idx: i + order_by_len,
+                orderby_sorter_idx: order_by_len + sequence_offset + i,
                 deduplicated: false,
             });
             i += 1;

core/translate/select.rs

@@ -4,7 +4,9 @@ use super::plan::{
     Search, TableReferences, WhereTerm, Window,
 };
 use crate::schema::Table;
+use crate::translate::emitter::Resolver;
 use crate::translate::expr::{bind_and_rewrite_expr, ParamState};
+use crate::translate::group_by::compute_group_by_sort_order;
 use crate::translate::optimizer::optimize_plan;
 use crate::translate::plan::{GroupBy, Plan, ResultSetColumn, SelectPlan};
 use crate::translate::planner::{
@@ -19,7 +21,7 @@ use crate::{schema::Schema, vdbe::builder::ProgramBuilder, Result};
 use crate::{Connection, SymbolTable};
 use std::sync::Arc;
 use turso_parser::ast::ResultColumn;
-use turso_parser::ast::{self, CompoundSelect, Expr, SortOrder};
+use turso_parser::ast::{self, CompoundSelect, Expr};
 
 pub struct TranslateSelectResult {
     pub program: ProgramBuilder,
@@ -424,7 +426,7 @@ fn prepare_one_select_plan(
                 }
 
                 plan.group_by = Some(GroupBy {
-                    sort_order: Some((0..group_by.exprs.len()).map(|_| SortOrder::Asc).collect()),
+                    sort_order: None,
                     exprs: group_by.exprs.iter().map(|expr| *expr.clone()).collect(),
                     having: if let Some(having) = group_by.having {
                         let mut predicates = vec![];
@@ -487,6 +489,16 @@ fn prepare_one_select_plan(
                 key.push((o.expr, o.order.unwrap_or(ast::SortOrder::Asc)));
             }
             plan.order_by = key;
+            if let Some(group_by) = &mut plan.group_by {
+                // now that we have resolved the ORDER BY expressions and aggregates, we can
+                // compute the necessary sort order for the GROUP BY clause
+                group_by.sort_order = Some(compute_group_by_sort_order(
+                    &group_by.exprs,
+                    &plan.order_by,
+                    &plan.aggregates,
+                    &Resolver::new(schema, syms),
+                ));
+            }
 
             // Parse the LIMIT/OFFSET clause
             (plan.limit, plan.offset) = limit.map_or(Ok((None, None)), |mut l| {

core/translate/window.rs

@@ -1013,15 +1013,7 @@ fn emit_return_buffered_rows(
             )?;
         }
         false => {
-            order_by_sorter_insert(
-                program,
-                &t_ctx.resolver,
-                t_ctx
-                    .meta_sort
-                    .as_ref()
-                    .expect("sort metadata must exist for ORDER BY"),
-                plan,
-            )?;
+            order_by_sorter_insert(program, t_ctx, plan)?;
         }
     }
 

core/vdbe/execute.rs

@@ -5354,6 +5354,49 @@ pub fn op_function(
     Ok(InsnFunctionStepResult::Step)
 }
 
+pub fn op_sequence(
+    program: &Program,
+    state: &mut ProgramState,
+    insn: &Insn,
+    pager: &Arc<Pager>,
+    mv_store: Option<&Arc<MvStore>>,
+) -> Result<InsnFunctionStepResult> {
+    load_insn!(
+        Sequence {
+            cursor_id,
+            target_reg
+        },
+        insn
+    );
+    let cursor = state.get_cursor(*cursor_id).as_sorter_mut();
+    let seq_num = cursor.next_sequence();
+    state.registers[*target_reg] = Register::Value(Value::Integer(seq_num));
+    state.pc += 1;
+    Ok(InsnFunctionStepResult::Step)
+}
+
+pub fn op_sequence_test(
+    program: &Program,
+    state: &mut ProgramState,
+    insn: &Insn,
+    pager: &Arc<Pager>,
+    mv_store: Option<&Arc<MvStore>>,
+) -> Result<InsnFunctionStepResult> {
+    load_insn!(
+        SequenceTest {
+            cursor_id,
+            target_pc,
+            value_reg
+        },
+        insn
+    );
+    let cursor = state.get_cursor(*cursor_id).as_sorter_mut();
+    if cursor.seq_beginning() {
+        state.pc = target_pc.as_offset_int();
+    }
+    Ok(InsnFunctionStepResult::Step)
+}
+
 pub fn op_init_coroutine(
     program: &Program,
     state: &mut ProgramState,

core/vdbe/explain.rs

@@ -1759,6 +1759,24 @@ pub fn insn_to_row(
                 0,
                 String::new(),
             ),
+        Insn::Sequence{ cursor_id, target_reg} => (
+                "Sequence",
+                *cursor_id as i32,
+                *target_reg as i32,
+                0,
+                Value::build_text(""),
+                0,
+                String::new(),
+          ),
+        Insn::SequenceTest{ cursor_id, target_pc, value_reg } => (
+            "SequenceTest",
+              *cursor_id as i32,
+            target_pc.as_debug_int(),
+            *value_reg as i32,
+            Value::build_text(""),
+            0,
+            String::new(),
+        ),
         }
 }
 

core/vdbe/insn.rs

@@ -1126,6 +1126,20 @@ pub enum Insn {
         target_pc: BranchOffset,
     },
 
+    /// Find the next available sequence number for cursor P1. Write the sequence number into register P2.
+    /// The sequence number on the cursor is incremented after this instruction.
+    Sequence {
+        cursor_id: CursorID,
+        target_reg: usize,
+    },
+
+    /// P1 is a sorter cursor. If the sequence counter is currently zero, jump to P2. Regardless of whether or not the jump is taken, increment the the sequence value.
+    SequenceTest {
+        cursor_id: CursorID,
+        target_pc: BranchOffset,
+        value_reg: usize,
+    },
+
     // OP_Explain
     Explain {
         p1: usize,         // P1: address of instruction
@@ -1295,6 +1309,8 @@ impl InsnVariants {
             InsnVariants::IfNeg => execute::op_if_neg,
             InsnVariants::Explain => execute::op_noop,
             InsnVariants::OpenDup => execute::op_open_dup,
+            InsnVariants::Sequence => execute::op_sequence,
+            InsnVariants::SequenceTest => execute::op_sequence_test,
         }
     }
 }

core/vdbe/sorter.rs

@@ -88,6 +88,7 @@ pub struct Sorter {
     insert_state: InsertState,
     /// State machine for [Sorter::init_chunk_heap]
     init_chunk_heap_state: InitChunkHeapState,
+    seq_count: i64,
 }
 
 impl Sorter {
@@ -125,6 +126,7 @@ impl Sorter {
             sort_state: SortState::Start,
             insert_state: InsertState::Start,
             init_chunk_heap_state: InitChunkHeapState::Start,
+            seq_count: 0,
         }
     }
 
@@ -136,6 +138,21 @@ impl Sorter {
         self.current.is_some()
     }
 
+    /// Get current sequence count and increment it
+    pub fn next_sequence(&mut self) -> i64 {
+        let current = self.seq_count;
+        self.seq_count += 1;
+        current
+    }
+
+    /// Test if at beginning of sequence (count == 0) and increment
+    /// Returns true if this was the first call (seq_count was 0)
+    pub fn seq_beginning(&mut self) -> bool {
+        let was_zero = self.seq_count == 0;
+        self.seq_count += 1;
+        was_zero
+    }
+
     // We do the sorting here since this is what is called by the SorterSort instruction
     pub fn sort(&mut self) -> Result<IOResult<()>> {
         loop {
@@ -578,6 +595,7 @@ struct SortableImmutableRecord {
     record: ImmutableRecord,
     cursor: RecordCursor,
     key_values: RefCell<Vec<RefValue>>,
+    key_len: usize,
     index_key_info: Rc<Vec<KeyInfo>>,
     /// The key deserialization error, if any.
     deserialization_error: RefCell<Option<LimboError>>,
@@ -601,6 +619,7 @@ impl SortableImmutableRecord {
             key_values: RefCell::new(Vec::with_capacity(key_len)),
             index_key_info,
             deserialization_error: RefCell::new(None),
+            key_len,
         })
     }
 
@@ -638,7 +657,7 @@ impl Ord for SortableImmutableRecord {
         let this_key_values_len = self.key_values.borrow().len();
         let other_key_values_len = other.key_values.borrow().len();
 
-        for i in 0..self.cursor.serial_types.len() {
+        for i in 0..self.key_len {
             // Lazily deserialize the key values if they haven't been deserialized already.
             if i >= this_key_values_len {
                 self.try_deserialize_key(i);