Support select <columns> for order by operation

benclmnt · benclmnt · commit e7b164a23fc4 · 2024-07-20T00:26:52.000+08:00
diff --git a/core/btree.rs b/core/btree.rs
@@ -4,7 +4,7 @@ use crate::types::{Cursor, CursorResult, OwnedRecord};
 
 use anyhow::Result;
 
-use std::cell::{Ref, RefCell};
+use std::cell::RefCell;
 use std::rc::Rc;
 
 pub struct MemPage {
diff --git a/core/expr.rs b/core/expr.rs
@@ -3,7 +3,7 @@ use sqlite3_parser::ast::{self, Expr, UnaryOperator};
 
 use crate::{
     function::{Func, SingleRowFunc},
-    schema::{Column, Schema, Table},
+    schema::{Schema, Table, Type},
     select::{ColumnInfo, Select, SrcTable},
     util::normalize_ident,
     vdbe::{BranchOffset, Insn, ProgramBuilder},
@@ -286,9 +286,8 @@ pub fn translate_expr(
                             for arg in args {
                                 let reg = program.alloc_register();
                                 let _ = translate_expr(program, select, arg, reg, cursor_hint)?;
-                                match arg {
-                                    ast::Expr::Literal(_) => program.mark_last_insn_constant(),
-                                    _ => {}
+                                if let ast::Expr::Literal(_) = arg {
+                                    program.mark_last_insn_constant()
                                 }
                             }
                             program.emit_insn(Insn::Function {
@@ -374,9 +373,9 @@ pub fn translate_expr(
         ast::Expr::FunctionCallStar { .. } => todo!(),
         ast::Expr::Id(ident) => {
             // let (idx, col) = table.unwrap().get_column(&ident.0).unwrap();
-            let (idx, col, cursor_id) =
+            let (idx, col_type, cursor_id, is_primary_key) =
                 resolve_ident_table(program, &ident.0, select, cursor_hint)?;
-            if col.primary_key {
+            if is_primary_key {
                 program.emit_insn(Insn::RowId {
                     cursor_id,
                     dest: target_register,
@@ -388,7 +387,7 @@ pub fn translate_expr(
                     cursor_id,
                 });
             }
-            maybe_apply_affinity(col, target_register, program);
+            maybe_apply_affinity(col_type, target_register, program);
             Ok(target_register)
         }
         ast::Expr::InList { .. } => todo!(),
@@ -436,9 +435,9 @@ pub fn translate_expr(
         ast::Expr::NotNull(_) => todo!(),
         ast::Expr::Parenthesized(_) => todo!(),
         ast::Expr::Qualified(tbl, ident) => {
-            let (idx, col, cursor_id) =
+            let (idx, col_type, cursor_id, is_primary_key) =
                 resolve_ident_qualified(program, &tbl.0, &ident.0, select, cursor_hint)?;
-            if col.primary_key {
+            if is_primary_key {
                 program.emit_insn(Insn::RowId {
                     cursor_id,
                     dest: target_register,
@@ -450,7 +449,7 @@ pub fn translate_expr(
                     cursor_id,
                 });
             }
-            maybe_apply_affinity(col, target_register, program);
+            maybe_apply_affinity(col_type, target_register, program);
             Ok(target_register)
         }
         ast::Expr::Raise(_, _) => todo!(),
@@ -562,7 +561,7 @@ pub fn resolve_ident_qualified<'a>(
     ident: &String,
     select: &'a Select,
     cursor_hint: Option<usize>,
-) -> Result<(usize, &'a Column, usize)> {
+) -> Result<(usize, Type, usize, bool)> {
     for join in &select.src_tables {
         match join.table {
             Table::BTree(ref table) => {
@@ -579,7 +578,7 @@ pub fn resolve_ident_qualified<'a>(
                     if res.is_some() {
                         let (idx, col) = res.unwrap();
                         let cursor_id = program.resolve_cursor_id(&table_identifier, cursor_hint);
-                        return Ok((idx, col, cursor_id));
+                        return Ok((idx, col.ty, cursor_id, col.primary_key));
                     }
                 }
             }
@@ -598,7 +597,7 @@ pub fn resolve_ident_table<'a>(
     ident: &String,
     select: &'a Select,
     cursor_hint: Option<usize>,
-) -> Result<(usize, &'a Column, usize)> {
+) -> Result<(usize, Type, usize, bool)> {
     let mut found = Vec::new();
     for join in &select.src_tables {
         match join.table {
@@ -611,11 +610,29 @@ pub fn resolve_ident_table<'a>(
                     .columns
                     .iter()
                     .enumerate()
-                    .find(|(_, col)| col.name == *ident);
+                    .find(|(_, col)| col.name == *ident)
+                    .map(|(idx, col)| (idx, col.ty, col.primary_key));
+                let mut idx;
+                let mut col_type;
+                let mut is_primary_key;
                 if res.is_some() {
-                    let (idx, col) = res.unwrap();
+                    (idx, col_type, is_primary_key) = res.unwrap();
+                    // overwrite if cursor hint is provided
+                    if let Some(cursor_hint) = cursor_hint {
+                        let cols = &program.cursor_ref[cursor_hint].1;
+                        if let Some(res) = cols.as_ref().and_then(|res| {
+                            res.columns()
+                                .iter()
+                                .enumerate()
+                                .find(|x| x.1.name == *ident)
+                        }) {
+                            idx = res.0;
+                            col_type = res.1.ty;
+                            is_primary_key = res.1.primary_key;
+                        }
+                    }
                     let cursor_id = program.resolve_cursor_id(&table_identifier, cursor_hint);
-                    found.push((idx, col, cursor_id));
+                    found.push((idx, col_type, cursor_id, is_primary_key));
                 }
             }
             Table::Pseudo(_) => todo!(),
@@ -631,8 +648,8 @@ pub fn resolve_ident_table<'a>(
     anyhow::bail!("Parse error: ambiguous column name {}", ident.as_str());
 }
 
-pub fn maybe_apply_affinity(col: &Column, target_register: usize, program: &mut ProgramBuilder) {
-    if col.ty == crate::schema::Type::Real {
+pub fn maybe_apply_affinity(col_type: Type, target_register: usize, program: &mut ProgramBuilder) {
+    if col_type == crate::schema::Type::Real {
         program.emit_insn(Insn::RealAffinity {
             register: target_register,
         })
diff --git a/core/pseudo.rs b/core/pseudo.rs
@@ -1,5 +1,5 @@
 use anyhow::Result;
-use std::cell::{Ref, RefCell};
+use std::cell::RefCell;
 
 use crate::types::{Cursor, CursorResult, OwnedRecord, OwnedValue};
 
@@ -35,13 +35,14 @@ impl Cursor for PseudoCursor {
     }
 
     fn rowid(&self) -> Result<Option<u64>> {
-        let x = self.current.borrow().as_ref().map(|record| {
-            let rowid = match record.values[0] {
+        let x = self
+            .current
+            .borrow()
+            .as_ref()
+            .map(|record| match record.values[0] {
                 OwnedValue::Integer(rowid) => rowid as u64,
                 _ => panic!("Expected integer value"),
-            };
-            rowid
-        });
+            });
         Ok(x)
     }
 
diff --git a/core/schema.rs b/core/schema.rs
@@ -1,4 +1,3 @@
-use crate::types::OwnedRecord;
 use crate::util::normalize_ident;
 use anyhow::Result;
 use core::fmt;
@@ -55,7 +54,7 @@ impl Table {
     pub fn get_name(&self) -> &str {
         match self {
             Table::BTree(table) => &table.name,
-            Table::Pseudo(table) => "pseudo",
+            Table::Pseudo(_) => "",
         }
     }
 
@@ -155,15 +154,11 @@ impl BTreeTable {
 #[derive(Debug)]
 pub struct PseudoTable {
     pub columns: Vec<Column>,
-    pub row: Option<OwnedRecord>,
 }
 
 impl PseudoTable {
     pub fn new() -> Self {
-        Self {
-            columns: vec![],
-            row: None,
-        }
+        Self { columns: vec![] }
     }
 
     pub fn add_column(&mut self, name: &str, ty: Type, primary_key: bool) {
@@ -302,7 +297,7 @@ pub fn _build_pseudo_table(columns: &[ResultColumn]) -> PseudoTable {
     table
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Column {
     pub name: String,
     pub ty: Type,
diff --git a/core/sorter.rs b/core/sorter.rs
@@ -1,6 +1,5 @@
-use crate::types::{Cursor, CursorResult, OwnedRecord, OwnedValue};
+use crate::types::{Cursor, CursorResult, OwnedRecord};
 use anyhow::Result;
-use log::trace;
 use std::{
     cell::RefCell,
     collections::{BTreeMap, VecDeque},
@@ -75,7 +74,6 @@ impl Cursor for Sorter {
     fn insert(&mut self, record: &OwnedRecord) -> Result<()> {
         let key_fields = self.order.len();
         let key = OwnedRecord::new(record.values[0..key_fields].to_vec());
-        trace!("Inserting record with key: {:?}", key);
         self.insert(key, OwnedRecord::new(record.values[key_fields..].to_vec()));
         Ok(())
     }
diff --git a/core/translate.rs b/core/translate.rs
@@ -4,10 +4,11 @@ use std::rc::Rc;
 use crate::expr::{build_select, maybe_apply_affinity, translate_expr};
 use crate::function::{AggFunc, Func};
 use crate::pager::Pager;
-use crate::schema::{Schema, Table};
+use crate::schema::{Column, PseudoTable, Schema, Table};
 use crate::select::{ColumnInfo, LoopInfo, Select, SrcTable};
 use crate::sqlite3_ondisk::{DatabaseHeader, MIN_PAGE_CACHE_SIZE};
 use crate::types::{OwnedRecord, OwnedValue};
+use crate::util::normalize_ident;
 use crate::vdbe::{BranchOffset, Insn, Program, ProgramBuilder};
 use crate::where_clause::{
     evaluate_conditions, translate_conditions, translate_where, Inner, Left, QueryConstraint,
@@ -178,7 +179,7 @@ fn translate_select(mut select: Select) -> Result<Program> {
         }
         program.emit_insn(Insn::ResultRow {
             start_reg: register_start,
-            count: count,
+            count,
         });
         emit_limit_insn(&limit_info, &mut program);
     };
@@ -228,8 +229,31 @@ fn translate_sorter(
     limit_info: &Option<LimitInfo>,
 ) -> Result<()> {
     assert!(sort_info.count > 0);
-
-    let pseudo_cursor = program.alloc_cursor_id(None, None);
+    let mut pseudo_columns = Vec::new();
+    for col in select.columns.iter() {
+        match col {
+            ast::ResultColumn::Expr(expr, _) => match expr {
+                ast::Expr::Id(ident) => {
+                    pseudo_columns.push(Column {
+                        name: normalize_ident(&ident.0),
+                        primary_key: false,
+                        ty: crate::schema::Type::Null,
+                    });
+                }
+                _ => {
+                    todo!();
+                }
+            },
+            ast::ResultColumn::Star => {}
+            ast::ResultColumn::TableStar(_) => {}
+        }
+    }
+    let pseudo_cursor = program.alloc_cursor_id(
+        None,
+        Some(Table::Pseudo(Rc::new(PseudoTable {
+            columns: pseudo_columns,
+        }))),
+    );
     let pseudo_content_reg = program.alloc_register();
     program.emit_insn(Insn::OpenPseudo {
         cursor_id: pseudo_cursor,
@@ -530,7 +554,7 @@ fn translate_table_star(
                 dest: col_target_register,
                 cursor_id: table_cursor,
             });
-            maybe_apply_affinity(col, col_target_register, program);
+            maybe_apply_affinity(col.ty, col_target_register, program);
         }
     }
 }
diff --git a/core/vdbe.rs b/core/vdbe.rs
@@ -271,7 +271,7 @@ pub struct ProgramBuilder {
     unresolved_labels: Vec<Vec<InsnReference>>,
     next_insn_label: Option<BranchOffset>,
     // Cursors that are referenced by the program. Indexed by CursorID.
-    cursor_ref: Vec<(Option<String>, Option<Table>)>,
+    pub cursor_ref: Vec<(Option<String>, Option<Table>)>,
     // List of deferred label resolutions. Each entry is a pair of (label, insn_reference).
     deferred_label_resolutions: Vec<(BranchOffset, InsnReference)>,
 }
@@ -303,14 +303,6 @@ impl ProgramBuilder {
         reg
     }
 
-    pub fn drop_register(&mut self) {
-        self.next_free_register -= 1;
-    }
-
-    pub fn drop_registers(&mut self, amount: usize) {
-        self.next_free_register -= amount;
-    }
-
     pub fn next_free_register(&self) -> usize {
         self.next_free_register
     }
@@ -335,17 +327,6 @@ impl ProgramBuilder {
         }
     }
 
-    pub fn last_insn(&self) -> Option<&Insn> {
-        self.insns.last()
-    }
-
-    pub fn last_of_type(&self, typ: std::mem::Discriminant<Insn>) -> Option<&Insn> {
-        self.insns
-            .iter()
-            .rev()
-            .find(|v| std::mem::discriminant(*v) == typ)
-    }
-
     // Emit an instruction that will be put at the end of the program (after Transaction statement).
     // This is useful for instructions that otherwise will be unnecessarily repeated in a loop.
     // Example: In `SELECT * from users where name='John'`, it is unnecessary to set r[1]='John' as we SCAN users table.
@@ -1629,7 +1610,7 @@ fn insn_to_str(program: &Program, addr: InsnReference, insn: &Insn, indent: Stri
                         dest,
                         table
                             .as_ref()
-                            .and_then(|x| Some(x.get_name()))
+                            .map(|x| x.get_name())
                             .unwrap_or(format!("cursor {}", cursor_id).as_str()),
                         table
                             .as_ref()
@@ -1766,7 +1747,7 @@ fn insn_to_str(program: &Program, addr: InsnReference, insn: &Insn, indent: Stri
                     &program.cursor_ref[*cursor_id]
                         .1
                         .as_ref()
-                        .and_then(|x| Some(x.get_name()))
+                        .map(|x| x.get_name())
                         .unwrap_or(format!("cursor {}", cursor_id).as_str())
                 ),
             ),
diff --git a/core/where_clause.rs b/core/where_clause.rs
@@ -518,11 +518,11 @@ fn introspect_expression_for_cursors(
             )?);
         }
         ast::Expr::Id(ident) => {
-            let (_, _, cursor_id) = resolve_ident_table(program, &ident.0, select, cursor_hint)?;
+            let (_, _, cursor_id, _) = resolve_ident_table(program, &ident.0, select, cursor_hint)?;
             cursors.push(cursor_id);
         }
         ast::Expr::Qualified(tbl, ident) => {
-            let (_, _, cursor_id) =
+            let (_, _, cursor_id, _) =
                 resolve_ident_qualified(program, &tbl.0, &ident.0, select, cursor_hint)?;
             cursors.push(cursor_id);
         }
diff --git a/testing/orderby.test b/testing/orderby.test
@@ -23,4 +23,21 @@ do_execsql_test basic-order-by-and-limit {
 9|boots|1.0
 2|cap|82.0
 10|coat|33.0
-1|hat|79.0}
+1|hat|79.0}
+
+do_execsql_test basic-order-by-and-limit-2 {
+    select id, name from products order by name limit 5;
+} {11|accessories
+9|boots
+2|cap
+10|coat
+1|hat}
+
+do_execsql_test basic-order-by-and-limit-3 {
+    select price, name from products where price > 70 order by name;
+} {81.0|accessories
+82.0|cap
+79.0|hat
+78.0|jeans
+82.0|sneakers
+74.0|sweatshirt}

Original file line number	Diff line number	Diff line change
`@@ -518,11 +518,11 @@ fn introspect_expression_for_cursors(`
`518`	`518`	`)?);`
`519`	`519`	`}`
`520`	`520`	`ast::Expr::Id(ident) => {`
`521`		`- let (_, _, cursor_id) = resolve_ident_table(program, &ident.0, select, cursor_hint)?;`
	`521`	`+ let (_, _, cursor_id, _) = resolve_ident_table(program, &ident.0, select, cursor_hint)?;`
`522`	`522`	`cursors.push(cursor_id);`
`523`	`523`	`}`
`524`	`524`	`ast::Expr::Qualified(tbl, ident) => {`
`525`		`- let (_, _, cursor_id) =`
	`525`	`+ let (_, _, cursor_id, _) =`
`526`	`526`	`resolve_ident_qualified(program, &tbl.0, &ident.0, select, cursor_hint)?;`
`527`	`527`	`cursors.push(cursor_id);`
`528`	`528`	`}`