Merge pull request #179 from benclmnt/orderby

Partial order by support

Author: penberg

COMPAT.md

@@ -44,7 +44,7 @@ This document describes the SQLite compatibility status of Limbo:
 | SELECT ... WHERE             | Partial |         |
 | SELECT ... WHERE ... LIKE    | Yes     |         |
 | SELECT ... LIMIT             | Yes     |         |
-| SELECT ... ORDER BY          | No      |         |
+| SELECT ... ORDER BY          | Partial |         |
 | SELECT ... GROUP BY          | No      |         |
 | SELECT ... JOIN              | Partial |         |
 | SELECT ... CROSS JOIN        | Partial |         |

Cargo.lock

@@ -34,7 +34,6 @@ fallible-iterator = "0.3.0"
 libc = "0.2.155"
 log = "0.4.20"
 nix = { version = "0.29.0", features = ["fs"] }
-ordered-multimap = "0.7.1"
 sieve-cache = "0.1.4"
 sqlite3-parser = "0.11.0"
 thiserror = "1.0.61"

core/Cargo.toml

@@ -152,8 +152,8 @@ impl Cursor for BTreeCursor {
         Ok(())
     }
 
-    fn rowid(&self) -> Result<Ref<Option<u64>>> {
-        Ok(self.rowid.borrow())
+    fn rowid(&self) -> Result<Option<u64>> {
+        Ok(*self.rowid.borrow())
     }
 
     fn record(&self) -> Result<Ref<Option<OwnedRecord>>> {

core/btree.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},
@@ -78,6 +78,7 @@ pub fn build_select<'a>(schema: &Schema, select: &'a ast::Select) -> Result<Sele
                 column_info,
                 src_tables: joins,
                 limit: &select.limit,
+                order_by: &select.order_by,
                 exist_aggregation,
                 where_clause,
                 loops: Vec::new(),
@@ -98,6 +99,7 @@ pub fn build_select<'a>(schema: &Schema, select: &'a ast::Select) -> Result<Sele
                 column_info,
                 src_tables: Vec::new(),
                 limit: &select.limit,
+                order_by: &select.order_by,
                 where_clause,
                 exist_aggregation,
                 loops: Vec::new(),
@@ -112,14 +114,15 @@ pub fn translate_expr(
     select: &Select,
     expr: &ast::Expr,
     target_register: usize,
+    cursor_hint: Option<usize>,
 ) -> Result<usize> {
     match expr {
         ast::Expr::Between { .. } => todo!(),
         ast::Expr::Binary(e1, op, e2) => {
             let e1_reg = program.alloc_register();
             let e2_reg = program.alloc_register();
-            let _ = translate_expr(program, select, e1, e1_reg)?;
-            let _ = translate_expr(program, select, e2, e2_reg)?;
+            let _ = translate_expr(program, select, e1, e1_reg, cursor_hint)?;
+            let _ = translate_expr(program, select, e2, e2_reg, cursor_hint)?;
 
             match op {
                 ast::Operator::NotEquals => {
@@ -250,7 +253,13 @@ pub fn translate_expr(
                             // whenever a not null check succeeds, we jump to the end of the series
                             let label_coalesce_end = program.allocate_label();
                             for (index, arg) in args.iter().enumerate() {
-                                let reg = translate_expr(program, select, arg, target_register)?;
+                                let reg = translate_expr(
+                                    program,
+                                    select,
+                                    arg,
+                                    target_register,
+                                    cursor_hint,
+                                )?;
                                 if index < args.len() - 1 {
                                     program.emit_insn_with_label_dependency(
                                         Insn::NotNull {
@@ -282,7 +291,7 @@ pub fn translate_expr(
                             };
                             for arg in args {
                                 let reg = program.alloc_register();
-                                let _ = translate_expr(program, select, &arg, reg)?;
+                                let _ = translate_expr(program, select, &arg, reg, cursor_hint)?;
                                 match arg {
                                     ast::Expr::Literal(_) => program.mark_last_insn_constant(),
                                     _ => {}
@@ -315,7 +324,7 @@ pub fn translate_expr(
                             };
 
                             let regs = program.alloc_register();
-                            translate_expr(program, select, &args[0], regs)?;
+                            translate_expr(program, select, &args[0], regs, cursor_hint)?;
                             program.emit_insn(Insn::Function {
                                 start_reg: regs,
                                 dest: target_register,
@@ -356,7 +365,7 @@ pub fn translate_expr(
 
                             for arg in args.iter() {
                                 let reg = program.alloc_register();
-                                translate_expr(program, select, arg, reg)?;
+                                translate_expr(program, select, arg, reg, cursor_hint)?;
                                 if let ast::Expr::Literal(_) = arg {
                                     program.mark_last_insn_constant();
                                 }
@@ -378,8 +387,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) = resolve_ident_table(program, &ident.0, select)?;
-            if col.primary_key {
+            let (idx, col_type, cursor_id, is_primary_key) =
+                resolve_ident_table(program, &ident.0, select, cursor_hint)?;
+            if is_primary_key {
                 program.emit_insn(Insn::RowId {
                     cursor_id,
                     dest: target_register,
@@ -391,7 +401,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!(),
@@ -439,8 +449,9 @@ pub fn translate_expr(
         ast::Expr::NotNull(_) => todo!(),
         ast::Expr::Parenthesized(_) => todo!(),
         ast::Expr::Qualified(tbl, ident) => {
-            let (idx, col, cursor_id) = resolve_ident_qualified(program, &tbl.0, &ident.0, select)?;
-            if col.primary_key {
+            let (idx, col_type, cursor_id, is_primary_key) =
+                resolve_ident_qualified(program, &tbl.0, &ident.0, select, cursor_hint)?;
+            if is_primary_key {
                 program.emit_insn(Insn::RowId {
                     cursor_id,
                     dest: target_register,
@@ -452,7 +463,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!(),
@@ -563,7 +574,8 @@ pub fn resolve_ident_qualified<'a>(
     table_name: &String,
     ident: &String,
     select: &'a Select,
-) -> Result<(usize, &'a Column, usize)> {
+    cursor_hint: Option<usize>,
+) -> Result<(usize, Type, usize, bool)> {
     for join in &select.src_tables {
         match join.table {
             Table::BTree(ref table) => {
@@ -579,8 +591,8 @@ pub fn resolve_ident_qualified<'a>(
                         .find(|(_, col)| col.name == *ident);
                     if res.is_some() {
                         let (idx, col) = res.unwrap();
-                        let cursor_id = program.resolve_cursor_id(&table_identifier);
-                        return Ok((idx, col, cursor_id));
+                        let cursor_id = program.resolve_cursor_id(&table_identifier, cursor_hint);
+                        return Ok((idx, col.ty, cursor_id, col.primary_key));
                     }
                 }
             }
@@ -598,7 +610,8 @@ pub fn resolve_ident_table<'a>(
     program: &ProgramBuilder,
     ident: &String,
     select: &'a Select,
-) -> Result<(usize, &'a Column, usize)> {
+    cursor_hint: Option<usize>,
+) -> Result<(usize, Type, usize, bool)> {
     let mut found = Vec::new();
     for join in &select.src_tables {
         match join.table {
@@ -611,11 +624,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();
-                    let cursor_id = program.resolve_cursor_id(&table_identifier);
-                    found.push((idx, col, cursor_id));
+                    (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_type, cursor_id, is_primary_key));
                 }
             }
             Table::Pseudo(_) => todo!(),
@@ -631,8 +662,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,
         })

core/expr.rs

@@ -56,6 +56,7 @@ impl ToString for SingleRowFunc {
     }
 }
 
+#[derive(Debug)]
 pub enum Func {
     Agg(AggFunc),
     SingleRow(SingleRowFunc),

core/function.rs

@@ -4,6 +4,7 @@ mod expr;
 mod function;
 mod io;
 mod pager;
+mod pseudo;
 mod schema;
 mod select;
 mod sorter;

core/lib.rs

@@ -0,0 +1,65 @@
+use anyhow::Result;
+use std::cell::{Ref, RefCell};
+
+use crate::types::{Cursor, CursorResult, OwnedRecord, OwnedValue};
+
+pub struct PseudoCursor {
+    current: RefCell<Option<OwnedRecord>>,
+}
+
+impl PseudoCursor {
+    pub fn new() -> Self {
+        Self {
+            current: RefCell::new(None),
+        }
+    }
+}
+
+impl Cursor for PseudoCursor {
+    fn is_empty(&self) -> bool {
+        self.current.borrow().is_none()
+    }
+
+    fn rewind(&mut self) -> Result<CursorResult<()>> {
+        *self.current.borrow_mut() = None;
+        Ok(CursorResult::Ok(()))
+    }
+
+    fn next(&mut self) -> Result<CursorResult<()>> {
+        *self.current.borrow_mut() = None;
+        Ok(CursorResult::Ok(()))
+    }
+
+    fn wait_for_completion(&mut self) -> Result<()> {
+        Ok(())
+    }
+
+    fn rowid(&self) -> Result<Option<u64>> {
+        let x = self
+            .current
+            .borrow()
+            .as_ref()
+            .map(|record| match record.values[0] {
+                OwnedValue::Integer(rowid) => rowid as u64,
+                _ => panic!("Expected integer value"),
+            });
+        Ok(x)
+    }
+
+    fn record(&self) -> Result<Ref<Option<OwnedRecord>>> {
+        Ok(self.current.borrow())
+    }
+
+    fn insert(&mut self, record: &OwnedRecord) -> Result<()> {
+        *self.current.borrow_mut() = Some(record.clone());
+        Ok(())
+    }
+
+    fn get_null_flag(&self) -> bool {
+        false
+    }
+
+    fn set_null_flag(&mut self, _null_flag: bool) {
+        // Do nothing
+    }
+}