From 293974e692e63b3151a8c64005aba678fc596014 Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Sun, 30 Mar 2025 15:24:46 -0400 Subject: [PATCH 01/16] Update COMPAT.md --- COMPAT.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/COMPAT.md b/COMPAT.md index 7fff65b87..d81650502 100644 --- a/COMPAT.md +++ b/COMPAT.md @@ -56,7 +56,7 @@ Limbo aims to be fully compatible with SQLite, with opt-in features not supporte | ATTACH DATABASE | No | | | BEGIN TRANSACTION | Partial | Transaction names are not supported. | | COMMIT TRANSACTION | Partial | Transaction names are not supported. | -| CREATE INDEX | No | | +| CREATE INDEX | Yes | | | CREATE TABLE | Partial | | | CREATE TRIGGER | No | | | CREATE VIEW | No | | @@ -461,6 +461,8 @@ Modifiers: | IdxDelete | No | | | IdxGE | Yes | | | IdxInsert | No | | +| IdxInsertAsync | Yes | | +| IdxInsertAwait | Yes | | | IdxLE | Yes | | | IdxLT | Yes | | | IdxRowid | No | | From 4741544dfdadec0892e6e644fa1a4e164e611cc6 Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Sun, 30 Mar 2025 15:25:42 -0400 Subject: [PATCH 02/16] Add query translation for create index --- core/translate/index.rs | 321 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 321 insertions(+) create mode 100644 core/translate/index.rs diff --git a/core/translate/index.rs b/core/translate/index.rs new file mode 100644 index 000000000..6f1a02c3c --- /dev/null +++ b/core/translate/index.rs @@ -0,0 +1,321 @@ +use std::sync::Arc; + +use crate::{ + schema::{BTreeTable, Column, Index, IndexColumn, PseudoTable, Schema}, + types::Record, + util::normalize_ident, + vdbe::{ + builder::{CursorType, ProgramBuilder, QueryMode}, + insn::{IdxInsertFlags, Insn}, + }, + OwnedValue, +}; +use limbo_sqlite3_parser::ast::{self, Expr, Id, SortOrder, SortedColumn}; + +use super::schema::{emit_schema_entry, SchemaEntryType, SQLITE_TABLEID}; + +pub fn translate_create_index( + mode: QueryMode, + unique_if_not_exists: (bool, bool), + idx_name: &str, + tbl_name: &str, + columns: &[SortedColumn], + schema: &Schema, +) -> crate::Result { + let idx_name = normalize_ident(idx_name); + let tbl_name = normalize_ident(tbl_name); + let mut program = ProgramBuilder::new(crate::vdbe::builder::ProgramBuilderOpts { + query_mode: mode, + num_cursors: 5, + approx_num_insns: 40, + approx_num_labels: 5, + }); + + // Check if the index is being created on a valid btree table and + // the name is globally unique in the schema. + if !schema.is_unique_idx_name(&idx_name) { + crate::bail_parse_error!("Error: index with name '{idx_name}' already exists."); + } + let Some(tbl) = schema.tables.get(&tbl_name) else { + crate::bail_parse_error!("Error: table '{tbl_name}' does not exist."); + }; + let Some(tbl) = tbl.btree() else { + crate::bail_parse_error!("Error: table '{tbl_name}' is not a b-tree table."); + }; + let columns = resolve_sorted_columns(&tbl, columns)?; + + // Prologue: + let init_label = program.emit_init(); + let start_offset = program.offset(); + + let idx = Arc::new(Index { + name: idx_name.clone(), + table_name: tbl.name.clone(), + root_page: 0, // we dont have access till its created, after we parse the schema table + columns: columns + .iter() + .map(|c| IndexColumn { + name: c.0 .1.name.as_ref().unwrap().clone(), + order: c.1, + }) + .collect(), + unique: unique_if_not_exists.0, + }); + + // Allocate the necessary cursors. + // + // 1. sqlite_schema_cursor_id - for the sqlite_schema table + // 2. btree_cursor_id - for the index btree + // 3. table_cursor_id - for the table we are creating the index on + // 4. sorter_cursor_id - for the sorter + // 5. pseudo_cursor_id - for the pseudo table to store the sorted index values + let sqlite_table = schema.get_btree_table(SQLITE_TABLEID).unwrap(); + let sqlite_schema_cursor_id = program.alloc_cursor_id( + Some(SQLITE_TABLEID.to_owned()), + CursorType::BTreeTable(sqlite_table.clone()), + ); + let btree_cursor_id = program.alloc_cursor_id( + Some(idx_name.to_owned()), + CursorType::BTreeIndex(idx.clone()), + ); + let table_cursor_id = program.alloc_cursor_id( + Some(tbl_name.to_owned()), + CursorType::BTreeTable(tbl.clone()), + ); + let sorter_cursor_id = program.alloc_cursor_id(None, CursorType::Sorter); + let pseudo_table = PseudoTable::new_with_columns(tbl.columns.clone()); + let pseudo_cursor_id = program.alloc_cursor_id(None, CursorType::Pseudo(pseudo_table.into())); + + // Create a new B-Tree and store the root page index in a register + let root_page_reg = program.alloc_register(); + program.emit_insn(Insn::CreateBtree { + db: 0, + root: root_page_reg, + flags: 2, // index leaf + }); + + // open the sqlite schema table for writing and create a new entry for the index + program.emit_insn(Insn::OpenWriteAsync { + cursor_id: sqlite_schema_cursor_id, + root_page: sqlite_table.root_page, + is_new_idx: false, + }); + program.emit_insn(Insn::OpenWriteAwait {}); + let sql = create_idx_stmt_to_sql(&tbl_name, &idx_name, unique_if_not_exists, &columns); + emit_schema_entry( + &mut program, + sqlite_schema_cursor_id, + SchemaEntryType::Index, + &idx_name, + &tbl_name, + root_page_reg, + Some(sql), + ); + + // determine the order of the columns in the index for the sorter + let order = idx + .columns + .iter() + .map(|c| { + OwnedValue::Integer(match c.order { + SortOrder::Asc => 0, + SortOrder::Desc => 1, + }) + }) + .collect(); + // open the sorter and the pseudo table + program.emit_insn(Insn::SorterOpen { + cursor_id: sorter_cursor_id, + columns: columns.len(), + order: Record::new(order), + }); + let content_reg = program.alloc_register(); + program.emit_insn(Insn::OpenPseudo { + cursor_id: pseudo_cursor_id, + content_reg, + num_fields: columns.len() + 1, + }); + + // open the table we are creating the index on for reading + program.emit_insn(Insn::OpenReadAsync { + cursor_id: table_cursor_id, + root_page: tbl.root_page, + }); + program.emit_insn(Insn::OpenReadAwait {}); + + program.emit_insn(Insn::RewindAsync { + cursor_id: table_cursor_id, + }); + let loop_start_label = program.allocate_label(); + let loop_end_label = program.allocate_label(); + program.emit_insn(Insn::RewindAwait { + cursor_id: table_cursor_id, + pc_if_empty: loop_end_label, + }); + + program.resolve_label(loop_start_label, program.offset()); + + // Loop start: + // Collect index values into start_reg..rowid_reg + // emit MakeRecord (index key + rowid) into record_reg. + // + // Then insert the record into the sorter + let start_reg = program.alloc_registers(columns.len() + 1); + for (i, (col, _)) in columns.iter().enumerate() { + program.emit_insn(Insn::Column { + cursor_id: table_cursor_id, + column: col.0, + dest: start_reg + i, + }); + } + let rowid_reg = start_reg + columns.len(); + program.emit_insn(Insn::RowId { + cursor_id: table_cursor_id, + dest: rowid_reg, + }); + let record_reg = program.alloc_register(); + program.emit_insn(Insn::MakeRecord { + start_reg, + count: columns.len() + 1, + dest_reg: record_reg, + }); + program.emit_insn(Insn::SorterInsert { + cursor_id: sorter_cursor_id, + record_reg, + }); + + program.emit_insn(Insn::NextAsync { + cursor_id: table_cursor_id, + }); + program.emit_insn(Insn::NextAwait { + cursor_id: table_cursor_id, + pc_if_next: loop_start_label, + }); + program.resolve_label(loop_end_label, program.offset()); + + // Open the index btree we created for writing to insert the + // newly sorted index records. + program.emit_insn(Insn::OpenWriteAsync { + cursor_id: btree_cursor_id, + root_page: root_page_reg, + is_new_idx: true, + }); + program.emit_insn(Insn::OpenWriteAwait {}); + + let sorted_loop_start = program.allocate_label(); + let sorted_loop_end = program.allocate_label(); + + // Sort the index records in the sorter + program.emit_insn(Insn::SorterSort { + cursor_id: sorter_cursor_id, + pc_if_empty: sorted_loop_end, + }); + program.resolve_label(sorted_loop_start, program.offset()); + let sorted_record_reg = program.alloc_register(); + program.emit_insn(Insn::SorterData { + pseudo_cursor: pseudo_cursor_id, + cursor_id: sorter_cursor_id, + dest_reg: sorted_record_reg, + }); + + // seek to the end of the index btree to position the cursor for appending + program.emit_insn(Insn::SeekEnd { + cursor_id: btree_cursor_id, + }); + // insert new index record + program.emit_insn(Insn::IdxInsertAsync { + cursor_id: btree_cursor_id, + record_reg: sorted_record_reg, + unpacked_start: None, // TODO: optimize with these to avoid decoding record twice + unpacked_count: None, + flags: IdxInsertFlags::new().use_seek(false), + }); + program.emit_insn(Insn::IdxInsertAwait { + cursor_id: btree_cursor_id, + }); + program.emit_insn(Insn::SorterNext { + cursor_id: sorter_cursor_id, + pc_if_next: sorted_loop_start, + }); + program.resolve_label(sorted_loop_end, program.offset()); + + // End of the outer loop + // + // Keep schema table open to emit ParseSchema, close the other cursors. + program.close_cursors(&[sorter_cursor_id, table_cursor_id, btree_cursor_id]); + + // TODO: SetCookie for schema change + // + // Parse the schema table to get the index root page and add new index to Schema + let parse_schema_where_clause = format!("name = '{}' AND type = 'index'", idx_name); + program.emit_insn(Insn::ParseSchema { + db: sqlite_schema_cursor_id, + where_clause: parse_schema_where_clause, + }); + // Close the final sqlite_schema cursor + program.emit_insn(Insn::Close { + cursor_id: sqlite_schema_cursor_id, + }); + + // Epilogue: + program.emit_halt(); + program.resolve_label(init_label, program.offset()); + program.emit_transaction(true); + program.emit_constant_insns(); + program.emit_goto(start_offset); + + Ok(program) +} + +fn resolve_sorted_columns<'a>( + table: &'a BTreeTable, + cols: &[SortedColumn], +) -> crate::Result> { + let mut resolved = Vec::with_capacity(cols.len()); + for sc in cols { + let ident = normalize_ident(match &sc.expr { + Expr::Id(Id(col_name)) | Expr::Name(ast::Name(col_name)) => col_name, + _ => crate::bail_parse_error!("Error: cannot use expressions in CREATE INDEX"), + }); + let Some(col) = table.get_column(&ident) else { + crate::bail_parse_error!( + "Error: column '{ident}' does not exist in table '{}'", + table.name + ); + }; + resolved.push((col, sc.order.unwrap_or(SortOrder::Asc))); + } + Ok(resolved) +} + +fn create_idx_stmt_to_sql( + tbl_name: &str, + idx_name: &str, + unique_if_not_exists: (bool, bool), + cols: &[((usize, &Column), SortOrder)], +) -> String { + let mut sql = String::new(); + sql.push_str("CREATE "); + if unique_if_not_exists.0 { + sql.push_str("UNIQUE "); + } + sql.push_str("INDEX "); + if unique_if_not_exists.1 { + sql.push_str("IF NOT EXISTS "); + } + sql.push_str(idx_name); + sql.push_str(" ON "); + sql.push_str(tbl_name); + sql.push_str(" ("); + for (i, (col, order)) in cols.iter().enumerate() { + if i > 0 { + sql.push_str(", "); + } + sql.push_str(col.1.name.as_ref().unwrap()); + if *order == SortOrder::Desc { + sql.push_str(" DESC"); + } + } + sql.push(')'); + sql +} From 7567b30d003a72f1ca2bf299f706c40e0d0d8b51 Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Sun, 30 Mar 2025 15:27:11 -0400 Subject: [PATCH 03/16] Add SeekEnd to compat.md --- COMPAT.md | 1 + 1 file changed, 1 insertion(+) diff --git a/COMPAT.md b/COMPAT.md index d81650502..b0cc85d83 100644 --- a/COMPAT.md +++ b/COMPAT.md @@ -550,6 +550,7 @@ Modifiers: | SeekLe | No | | | SeekLt | No | | | SeekRowid | Yes | | +| SeekEnd | Yes | | | Sequence | No | | | SetCookie | No | | | ShiftLeft | Yes | | From 20adedff4cb61ccd98bf431abb729eacfe8fdf2a Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Sun, 30 Mar 2025 15:32:14 -0400 Subject: [PATCH 04/16] Remove Order enum in place of ast::SortOrder --- core/schema.rs | 39 +++++++++++++++------------------------ 1 file changed, 15 insertions(+), 24 deletions(-) diff --git a/core/schema.rs b/core/schema.rs index fbec7627f..dda37d15b 100644 --- a/core/schema.rs +++ b/core/schema.rs @@ -2,7 +2,7 @@ use crate::VirtualTable; use crate::{util::normalize_ident, Result}; use core::fmt; use fallible_iterator::FallibleIterator; -use limbo_sqlite3_parser::ast::{Expr, Literal, TableOptions}; +use limbo_sqlite3_parser::ast::{Expr, Literal, SortOrder, TableOptions}; use limbo_sqlite3_parser::{ ast::{Cmd, CreateTableBody, QualifiedName, ResultColumn, Stmt}, lexer::sql::Parser, @@ -30,6 +30,13 @@ impl Schema { Self { tables, indexes } } + pub fn is_unique_idx_name(&self, name: &str) -> bool { + !self + .indexes + .iter() + .any(|idx| idx.1.iter().any(|i| i.name == name)) + } + pub fn add_btree_table(&mut self, table: Rc) { let name = normalize_ident(&table.name); self.tables.insert(name, Table::BTree(table).into()); @@ -209,7 +216,7 @@ impl BTreeTable { } } -#[derive(Debug)] +#[derive(Debug, Default)] pub struct PseudoTable { pub columns: Vec, } @@ -245,12 +252,6 @@ impl PseudoTable { } } -impl Default for PseudoTable { - fn default() -> Self { - Self::new() - } -} - fn create_table( tbl_name: QualifiedName, body: CreateTableBody, @@ -616,13 +617,7 @@ pub struct Index { #[derive(Debug, Clone)] pub struct IndexColumn { pub name: String, - pub order: Order, -} - -#[derive(Debug, Clone, PartialEq)] -pub enum Order { - Ascending, - Descending, + pub order: SortOrder, } impl Index { @@ -642,11 +637,7 @@ impl Index { .into_iter() .map(|col| IndexColumn { name: normalize_ident(&col.expr.to_string()), - order: match col.order { - Some(limbo_sqlite3_parser::ast::SortOrder::Asc) => Order::Ascending, - Some(limbo_sqlite3_parser::ast::SortOrder::Desc) => Order::Descending, - None => Order::Ascending, - }, + order: col.order.unwrap_or(SortOrder::Asc), }) .collect(); Ok(Index { @@ -685,7 +676,7 @@ impl Index { } Ok(IndexColumn { name: normalize_ident(col_name), - order: Order::Ascending, // Primary key indexes are always ascending + order: SortOrder::Asc, // Primary key indexes are always ascending }) }) .collect::>>()?; @@ -1012,7 +1003,7 @@ mod tests { assert!(index.unique); assert_eq!(index.columns.len(), 1); assert_eq!(index.columns[0].name, "a"); - assert!(matches!(index.columns[0].order, Order::Ascending)); + assert!(matches!(index.columns[0].order, SortOrder::Asc)); Ok(()) } @@ -1029,8 +1020,8 @@ mod tests { assert_eq!(index.columns.len(), 2); assert_eq!(index.columns[0].name, "a"); assert_eq!(index.columns[1].name, "b"); - assert!(matches!(index.columns[0].order, Order::Ascending)); - assert!(matches!(index.columns[1].order, Order::Ascending)); + assert!(matches!(index.columns[0].order, SortOrder::Asc)); + assert!(matches!(index.columns[1].order, SortOrder::Asc)); Ok(()) } From 45a8e5e2267c06744158c4d31e03b26c1f4ca12e Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Sun, 30 Mar 2025 15:34:43 -0400 Subject: [PATCH 05/16] Add close_cursors helper method to program builder --- core/vdbe/builder.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/core/vdbe/builder.rs b/core/vdbe/builder.rs index 058fc8aab..19a71a68d 100644 --- a/core/vdbe/builder.rs +++ b/core/vdbe/builder.rs @@ -131,6 +131,12 @@ impl ProgramBuilder { self.insns.push((insn, function)); } + pub fn close_cursors(&mut self, cursors: &[CursorID]) { + for cursor in cursors { + self.emit_insn(Insn::Close { cursor_id: *cursor }); + } + } + pub fn emit_string8(&mut self, value: String, dest: usize) { self.emit_insn(Insn::String8 { value, dest }); } From b0016a0ee24f6592c1aef0490ec043c64a08f584 Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Sun, 30 Mar 2025 15:41:19 -0400 Subject: [PATCH 06/16] Support create index with SeekEnd and IdxCreate opcode functionality --- core/storage/btree.rs | 160 ++++++++++++++++++++++++++++++++++-- core/translate/insert.rs | 2 + core/translate/main_loop.rs | 6 ++ core/translate/mod.rs | 21 ++++- core/translate/schema.rs | 13 +-- core/vdbe/execute.rs | 88 +++++++++++++++++++- core/vdbe/explain.rs | 34 ++++++++ core/vdbe/insn.rs | 61 +++++++++++++- 8 files changed, 370 insertions(+), 15 deletions(-) diff --git a/core/storage/btree.rs b/core/storage/btree.rs index 558a6eee3..0f094ac5b 100644 --- a/core/storage/btree.rs +++ b/core/storage/btree.rs @@ -12,12 +12,14 @@ use crate::types::{ use crate::{return_corrupt, LimboError, Result}; use std::cell::{Cell, Ref, RefCell}; +use std::cmp::Ordering; use std::pin::Pin; use std::rc::Rc; use super::pager::PageRef; use super::sqlite3_ondisk::{ - write_varint_to_vec, IndexInteriorCell, IndexLeafCell, OverflowCell, DATABASE_HEADER_SIZE, + read_record, write_varint_to_vec, IndexInteriorCell, IndexLeafCell, OverflowCell, + DATABASE_HEADER_SIZE, }; /* @@ -599,8 +601,8 @@ impl BTreeCursor { BTreeCell::TableLeafCell(TableLeafCell { _rowid, _payload, - first_overflow_page, payload_size, + first_overflow_page, }) => { assert!(predicate.is_none()); if let Some(next_page) = first_overflow_page { @@ -814,10 +816,8 @@ impl BTreeCursor { }; let record = self.get_immutable_record(); let record = record.as_ref().unwrap(); - let order = compare_immutable( - &record.get_values().as_slice()[..record.len() - 1], - &index_key.get_values().as_slice()[..], - ); + let without_rowid = &record.get_values().as_slice()[..record.len() - 1]; + let order = without_rowid.cmp(index_key.get_values()); let found = match op { SeekOp::GT => order.is_gt(), SeekOp::GE => order.is_ge(), @@ -1047,6 +1047,65 @@ impl BTreeCursor { } } + pub fn insert_index_key(&mut self, key: &ImmutableRecord) -> Result> { + if let CursorState::None = &self.state { + self.state = CursorState::Write(WriteInfo::new()); + } + + let ret = loop { + let write_state = self.state.mut_write_info().unwrap().state; + match write_state { + WriteState::Start => { + let page = self.stack.top(); + return_if_locked!(page); + page.set_dirty(); + self.pager.add_dirty(page.get().id); + let page = page.get().contents.as_mut().unwrap(); + + assert!(matches!(page.page_type(), PageType::IndexLeaf)); + let cell_idx = self.find_index_cell(page, key); + let mut cell_payload: Vec = Vec::new(); + fill_cell_payload( + page.page_type(), + None, + &mut cell_payload, + key, + self.usable_space() as u16, + self.pager.clone(), + ); + // insert + let overflow = { + debug!( + "insert_index_key(overflow, cell_count={})", + page.cell_count() + ); + insert_into_cell( + page, + cell_payload.as_slice(), + cell_idx, + self.usable_space() as u16, + )?; + page.overflow_cells.len() + }; + let write_info = self.state.mut_write_info().unwrap(); + write_info.state = if overflow > 0 { + WriteState::BalanceStart + } else { + WriteState::Finish + }; + } + WriteState::BalanceStart + | WriteState::BalanceNonRoot + | WriteState::BalanceNonRootWaitLoadPages => { + return_if_io!(self.balance()); + } + WriteState::Finish => break Ok(CursorResult::Ok(())), + } + }; + self.state = CursorState::None; + ret + } + /// Insert a record into the btree. /// If the insert operation overflows the page, it will be split and the btree will be balanced. fn insert_into_page( @@ -1943,6 +2002,74 @@ impl BTreeCursor { cell_idx } + fn find_index_cell(&self, page: &PageContent, key: &ImmutableRecord) -> usize { + let mut cell_idx = 0; + let cell_count = page.cell_count(); + while cell_idx < cell_count { + match page + .cell_get( + cell_idx, + payload_overflow_threshold_max(page.page_type(), self.usable_space() as u16), + payload_overflow_threshold_min(page.page_type(), self.usable_space() as u16), + self.usable_space(), + ) + .unwrap() + { + BTreeCell::IndexInteriorCell(IndexInteriorCell { payload, .. }) + | BTreeCell::IndexLeafCell(IndexLeafCell { payload, .. }) => { + read_record( + payload, + self.get_immutable_record_or_create().as_mut().unwrap(), + ) + .expect("failed to read record"); + let order = compare_immutable( + key.get_values(), + self.get_immutable_record().as_ref().unwrap().get_values(), + ); + match order { + Ordering::Less => { + break; + } + Ordering::Equal => { + break; + } + Ordering::Greater => {} + } + } + _ => unreachable!("Expected Index cell types"), + } + cell_idx += 1; + } + cell_idx + } + + pub fn seek_end(&mut self) -> Result> { + assert!(self.mv_cursor.is_none()); + self.move_to_root(); + loop { + let mem_page = self.stack.top(); + let page_id = mem_page.get().id; + let page = self.pager.read_page(page_id)?; + return_if_locked!(page); + + let contents = page.get().contents.as_ref().unwrap(); + if contents.is_leaf() { + // set cursor just past the last cell to append + self.stack.set_cell_index(contents.cell_count() as i32); + return Ok(CursorResult::Ok(())); + } + + match contents.rightmost_pointer() { + Some(right_most_pointer) => { + self.stack.set_cell_index(contents.cell_count() as i32 + 1); // invalid on interior + let child = self.pager.read_page(right_most_pointer as usize)?; + self.stack.push(child); + } + None => unreachable!("interior page must have rightmost pointer"), + } + } + } + pub fn seek_to_last(&mut self) -> Result> { return_if_io!(self.move_to_rightmost()); let rowid = return_if_io!(self.get_next_record(None)); @@ -2372,6 +2499,27 @@ impl BTreeCursor { self.null_flag } + /// Search for a key in an Index Btree. Looking up indexes that need to be unique, we cannot compare the rowid + pub fn key_exists_in_index(&mut self, key: &ImmutableRecord) -> Result> { + return_if_io!(self.do_seek(SeekKey::IndexKey(key), SeekOp::GE)); + if let Some(record) = self.record().as_ref() { + // get existing record, excluding the rowid + assert!(record.len() > 0); + let existing_key = &record.get_values()[..record.count() - 1]; + let inserted_key_vals = &key.get_values(); + if existing_key + .iter() + .zip(inserted_key_vals.iter()) + .all(|(a, b)| a == b) + { + return Ok(CursorResult::Ok(true)); // duplicate + } + } else { + return Err(LimboError::InvalidArgument("Expected Record key".into())); + } + Ok(CursorResult::Ok(false)) // no matching key found + } + pub fn exists(&mut self, key: &OwnedValue) -> Result> { assert!(self.mv_cursor.is_none()); let int_key = match key { diff --git a/core/translate/insert.rs b/core/translate/insert.rs index 5fda098e6..d6faeb2c8 100644 --- a/core/translate/insert.rs +++ b/core/translate/insert.rs @@ -153,6 +153,7 @@ pub fn translate_insert( program.emit_insn(Insn::OpenWriteAsync { cursor_id, root_page, + is_new_idx: false, }); program.emit_insn(Insn::OpenWriteAwait {}); @@ -169,6 +170,7 @@ pub fn translate_insert( program.emit_insn(Insn::OpenWriteAsync { cursor_id, root_page, + is_new_idx: false, }); program.emit_insn(Insn::OpenWriteAwait {}); diff --git a/core/translate/main_loop.rs b/core/translate/main_loop.rs index 7b51a2328..d9b7d6a36 100644 --- a/core/translate/main_loop.rs +++ b/core/translate/main_loop.rs @@ -103,6 +103,7 @@ pub fn init_loop( program.emit_insn(Insn::OpenWriteAsync { cursor_id, root_page, + is_new_idx: false, }); program.emit_insn(Insn::OpenWriteAwait {}); } @@ -111,6 +112,7 @@ pub fn init_loop( program.emit_insn(Insn::OpenWriteAsync { cursor_id, root_page, + is_new_idx: false, }); program.emit_insn(Insn::OpenWriteAwait {}); } @@ -145,6 +147,7 @@ pub fn init_loop( program.emit_insn(Insn::OpenWriteAsync { cursor_id: table_cursor_id, root_page: table.table.get_root_page(), + is_new_idx: false, }); program.emit_insn(Insn::OpenWriteAwait {}); } @@ -152,6 +155,7 @@ pub fn init_loop( program.emit_insn(Insn::OpenWriteAsync { cursor_id: table_cursor_id, root_page: table.table.get_root_page(), + is_new_idx: false, }); program.emit_insn(Insn::OpenWriteAwait {}); } @@ -178,6 +182,7 @@ pub fn init_loop( program.emit_insn(Insn::OpenWriteAsync { cursor_id: index_cursor_id, root_page: index.root_page, + is_new_idx: false, }); program.emit_insn(Insn::OpenWriteAwait {}); } @@ -185,6 +190,7 @@ pub fn init_loop( program.emit_insn(Insn::OpenWriteAsync { cursor_id: index_cursor_id, root_page: index.root_page, + is_new_idx: false, }); program.emit_insn(Insn::OpenWriteAwait {}); } diff --git a/core/translate/mod.rs b/core/translate/mod.rs index 739ae5f03..cf93b34ba 100644 --- a/core/translate/mod.rs +++ b/core/translate/mod.rs @@ -12,6 +12,7 @@ pub(crate) mod delete; pub(crate) mod emitter; pub(crate) mod expr; pub(crate) mod group_by; +pub(crate) mod index; pub(crate) mod insert; pub(crate) mod main_loop; pub(crate) mod optimizer; @@ -34,6 +35,7 @@ use crate::translate::delete::translate_delete; use crate::vdbe::builder::{ProgramBuilder, ProgramBuilderOpts, QueryMode}; use crate::vdbe::Program; use crate::{bail_parse_error, Connection, Result, SymbolTable}; +use index::translate_create_index; use insert::translate_insert; use limbo_sqlite3_parser::ast::{self, Delete, Insert}; use schema::{translate_create_table, translate_create_virtual_table, translate_drop_table}; @@ -61,7 +63,24 @@ pub fn translate( ast::Stmt::Attach { .. } => bail_parse_error!("ATTACH not supported yet"), ast::Stmt::Begin(tx_type, tx_name) => translate_tx_begin(tx_type, tx_name)?, ast::Stmt::Commit(tx_name) => translate_tx_commit(tx_name)?, - ast::Stmt::CreateIndex { .. } => bail_parse_error!("CREATE INDEX not supported yet"), + ast::Stmt::CreateIndex { + unique, + if_not_exists, + idx_name, + tbl_name, + columns, + .. + } => { + change_cnt_on = true; + translate_create_index( + query_mode, + (unique, if_not_exists), + &idx_name.name.0, + &tbl_name.0, + &columns, + schema, + )? + } ast::Stmt::CreateTable { temporary, if_not_exists, diff --git a/core/translate/schema.rs b/core/translate/schema.rs index 29cf29644..6f87937c9 100644 --- a/core/translate/schema.rs +++ b/core/translate/schema.rs @@ -104,6 +104,7 @@ pub fn translate_create_table( program.emit_insn(Insn::OpenWriteAsync { cursor_id: sqlite_schema_cursor_id, root_page: 1, + is_new_idx: false, }); program.emit_insn(Insn::OpenWriteAwait {}); @@ -155,8 +156,8 @@ pub fn translate_create_table( Ok(program) } -#[derive(Debug)] -enum SchemaEntryType { +#[derive(Debug, Clone, Copy)] +pub enum SchemaEntryType { Table, Index, } @@ -169,9 +170,9 @@ impl SchemaEntryType { } } } -const SQLITE_TABLEID: &str = "sqlite_schema"; +pub const SQLITE_TABLEID: &str = "sqlite_schema"; -fn emit_schema_entry( +pub fn emit_schema_entry( program: &mut ProgramBuilder, sqlite_schema_cursor_id: usize, entry_type: SchemaEntryType, @@ -501,6 +502,7 @@ pub fn translate_create_virtual_table( program.emit_insn(Insn::OpenWriteAsync { cursor_id: sqlite_schema_cursor_id, root_page: 1, + is_new_idx: false, }); program.emit_insn(Insn::OpenWriteAwait {}); @@ -572,7 +574,7 @@ pub fn translate_drop_table( let row_id_reg = program.alloc_register(); // r5 let table_name = "sqlite_schema"; - let schema_table = schema.get_btree_table(&table_name).unwrap(); + let schema_table = schema.get_btree_table(table_name).unwrap(); let sqlite_schema_cursor_id = program.alloc_cursor_id( Some(table_name.to_string()), CursorType::BTreeTable(schema_table.clone()), @@ -580,6 +582,7 @@ pub fn translate_drop_table( program.emit_insn(Insn::OpenWriteAsync { cursor_id: sqlite_schema_cursor_id, root_page: 1, + is_new_idx: false, }); program.emit_insn(Insn::OpenWriteAwait {}); diff --git a/core/vdbe/execute.rs b/core/vdbe/execute.rs index 3c511a0db..89bd26859 100644 --- a/core/vdbe/execute.rs +++ b/core/vdbe/execute.rs @@ -21,7 +21,7 @@ use crate::util::{ checked_cast_text_to_numeric, parse_schema_rows, RoundToPrecision, }; use crate::vdbe::builder::CursorType; -use crate::vdbe::insn::Insn; +use crate::vdbe::insn::{IdxInsertFlags, Insn}; use crate::vector::{vector32, vector64, vector_distance_cos, vector_extract}; use crate::{info, MvCursor, RefValue, Row, StepResult, TransactionState}; @@ -2049,6 +2049,24 @@ pub fn op_idx_ge( Ok(InsnFunctionStepResult::Step) } +pub fn op_seek_end( + program: &Program, + state: &mut ProgramState, + insn: &Insn, + pager: &Rc, + mv_store: Option<&Rc>, +) -> Result { + if let Insn::SeekEnd { cursor_id } = *insn { + let mut cursor = state.get_cursor(cursor_id); + let cursor = cursor.as_btree_mut(); + return_if_io!(cursor.seek_end()); + } else { + unreachable!("unexpected Insn {:?}", insn) + } + state.pc += 1; + Ok(InsnFunctionStepResult::Step) +} + pub fn op_idx_le( program: &Program, state: &mut ProgramState, @@ -3706,6 +3724,73 @@ pub fn op_delete_async( Ok(InsnFunctionStepResult::Step) } +pub fn op_idx_insert_async( + program: &Program, + state: &mut ProgramState, + insn: &Insn, + pager: &Rc, + mv_store: Option<&Rc>, +) -> Result { + if let Insn::IdxInsertAsync { + cursor_id, + record_reg, + flags, + .. + } = *insn + { + let (_, cursor_type) = program.cursor_ref.get(cursor_id).unwrap(); + let CursorType::BTreeIndex(index_meta) = cursor_type else { + panic!("IdxInsert: not a BTree index cursor"); + }; + { + let mut cursor = state.get_cursor(cursor_id); + let cursor = cursor.as_btree_mut(); + let record = match &state.registers[record_reg] { + Register::Record(ref r) => r, + _ => return Err(LimboError::InternalError("expected record".into())), + }; + let moved_before = if index_meta.unique { + // check for uniqueness violation + match cursor.key_exists_in_index(record)? { + CursorResult::Ok(true) => { + return Err(LimboError::Constraint( + "UNIQUE constraint failed: duplicate key".into(), + )) + } + CursorResult::IO => return Ok(InsnFunctionStepResult::IO), + CursorResult::Ok(false) => {} + }; + false + } else { + flags.has(IdxInsertFlags::USE_SEEK) + }; + // insert record as key + return_if_io!(cursor.insert_index_key(record)); + } + state.pc += 1; + } + Ok(InsnFunctionStepResult::Step) +} + +pub fn op_idx_insert_await( + program: &Program, + state: &mut ProgramState, + insn: &Insn, + pager: &Rc, + mv_store: Option<&Rc>, +) -> Result { + if let Insn::IdxInsertAwait { cursor_id } = *insn { + { + let mut cursor = state.get_cursor(cursor_id); + let cursor = cursor.as_btree_mut(); + cursor.wait_for_completion()?; + } + // TODO: flag optimizations, update n_change if OPFLAG_NCHANGE + state.pc += 1; + } + Ok(InsnFunctionStepResult::Step) +} + pub fn op_delete_await( program: &Program, state: &mut ProgramState, @@ -3889,6 +3974,7 @@ pub fn op_open_write_async( let Insn::OpenWriteAsync { cursor_id, root_page, + .. } = insn else { unreachable!("unexpected Insn {:?}", insn) diff --git a/core/vdbe/explain.rs b/core/vdbe/explain.rs index 67333c334..f80a442f1 100644 --- a/core/vdbe/explain.rs +++ b/core/vdbe/explain.rs @@ -760,6 +760,39 @@ pub fn insn_to_str( 0, "".to_string(), ), + Insn::SeekEnd { cursor_id } => ( + "SeekEnd", + *cursor_id as i32, + 0, + 0, + OwnedValue::build_text(""), + 0, + "".to_string(), + ), + Insn::IdxInsertAsync { + cursor_id, + record_reg, + unpacked_start, + flags, + .. + } => ( + "IdxInsertAsync", + *cursor_id as i32, + *record_reg as i32, + unpacked_start.unwrap_or(0) as i32, + OwnedValue::build_text(""), + flags.0 as u16, + format!("key=r[{}]", record_reg), + ), + Insn::IdxInsertAwait { cursor_id } => ( + "IdxInsertAwait", + *cursor_id as i32, + 0, + 0, + OwnedValue::build_text(""), + 0, + "".to_string(), + ), Insn::IdxGT { cursor_id, start_reg, @@ -1097,6 +1130,7 @@ pub fn insn_to_str( Insn::OpenWriteAsync { cursor_id, root_page, + .. } => ( "OpenWriteAsync", *cursor_id as i32, diff --git a/core/vdbe/insn.rs b/core/vdbe/insn.rs index f45e7ce35..5ce68f14a 100644 --- a/core/vdbe/insn.rs +++ b/core/vdbe/insn.rs @@ -38,6 +38,44 @@ impl CmpInsFlags { } } +#[derive(Clone, Copy, Debug, Default)] +pub struct IdxInsertFlags(pub u8); +impl IdxInsertFlags { + pub const APPEND: u8 = 0x01; // Hint: insert likely at the end + pub const NCHANGE: u8 = 0x02; // Increment the change counter + pub const USE_SEEK: u8 = 0x04; // Skip seek if last one was same key + pub fn new() -> Self { + IdxInsertFlags(0) + } + pub fn has(&self, flag: u8) -> bool { + (self.0 & flag) != 0 + } + pub fn append(mut self, append: bool) -> Self { + if append { + self.0 |= IdxInsertFlags::APPEND; + } else { + self.0 &= !IdxInsertFlags::APPEND; + } + self + } + pub fn use_seek(mut self, seek: bool) -> Self { + if seek { + self.0 |= IdxInsertFlags::USE_SEEK; + } else { + self.0 &= !IdxInsertFlags::USE_SEEK; + } + self + } + pub fn nchange(mut self, change: bool) -> Self { + if change { + self.0 |= IdxInsertFlags::NCHANGE; + } else { + self.0 &= !IdxInsertFlags::NCHANGE; + } + self + } +} + #[derive(Description, Debug)] pub enum Insn { /// Initialize the program state and jump to the given PC. @@ -401,6 +439,9 @@ pub enum Insn { src_reg: usize, target_pc: BranchOffset, }, + SeekEnd { + cursor_id: CursorID, + }, /// P1 is an open index cursor and P3 is a cursor on the corresponding table. This opcode does a deferred seek of the P3 table cursor to the row that corresponds to the current row of P1. /// This is a deferred seek. Nothing actually happens until the cursor is used to read a record. That way, if no reads occur, no unnecessary I/O happens. @@ -431,8 +472,20 @@ pub enum Insn { target_pc: BranchOffset, }, - /// The P4 register values beginning with P3 form an unpacked index key that omits the PRIMARY KEY. Compare this key value against the index that P1 is currently pointing to, ignoring the PRIMARY KEY or ROWID fields at the end. - /// If the P1 index entry is greater or equal than the key value then jump to P2. Otherwise fall through to the next instruction. + // cursor_id is a cursor pointing to a B-Tree index that uses integer keys, this op writes the value obtained from MakeRecord into the index. + // P3 + P4 are for the original column values that make up that key in unpacked (pre-serialized) form. + // If P5 has the OPFLAG_APPEND bit set, that is a hint to the b-tree layer that this insert is likely to be an append. + // OPFLAG_NCHANGE bit set, then the change counter is incremented by this instruction. If the OPFLAG_NCHANGE bit is clear, then the change counter is unchanged + IdxInsertAsync { + cursor_id: CursorID, + record_reg: usize, // P2 the register containing the record to insert + unpacked_start: Option, // P3 the index of the first register for the unpacked key + unpacked_count: Option, // P4 # of unpacked values in the key in P2 + flags: IdxInsertFlags, // TODO: optimization + }, + IdxInsertAwait { + cursor_id: CursorID, + }, IdxGE { cursor_id: CursorID, start_reg: usize, @@ -588,6 +641,7 @@ pub enum Insn { OpenWriteAsync { cursor_id: CursorID, root_page: PageIdx, + is_new_idx: bool, }, OpenWriteAwait {}, @@ -1237,10 +1291,13 @@ impl Insn { Insn::DeferredSeek { .. } => execute::op_deferred_seek, Insn::SeekGE { .. } => execute::op_seek_ge, Insn::SeekGT { .. } => execute::op_seek_gt, + Insn::SeekEnd { .. } => execute::op_seek_end, Insn::IdxGE { .. } => execute::op_idx_ge, Insn::IdxGT { .. } => execute::op_idx_gt, Insn::IdxLE { .. } => execute::op_idx_le, Insn::IdxLT { .. } => execute::op_idx_lt, + Insn::IdxInsertAsync { .. } => execute::op_idx_insert_async, + Insn::IdxInsertAwait { .. } => execute::op_idx_insert_await, Insn::DecrJumpZero { .. } => execute::op_decr_jump_zero, Insn::AggStep { .. } => execute::op_agg_step, From e020ba3dfe97e49e85652d9f36320a5c347ba576 Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Sun, 30 Mar 2025 16:16:55 -0400 Subject: [PATCH 07/16] Add enum for interpreting a value as a register or literal for insns --- core/translate/index.rs | 8 +++----- core/translate/insert.rs | 7 +++---- core/translate/main_loop.rs | 18 ++++++------------ core/translate/schema.rs | 10 ++++------ core/vdbe/execute.rs | 17 ++++++++++++++--- core/vdbe/explain.rs | 7 +++++-- core/vdbe/insn.rs | 20 +++++++++++++++----- 7 files changed, 50 insertions(+), 37 deletions(-) diff --git a/core/translate/index.rs b/core/translate/index.rs index 6f1a02c3c..c9a474cab 100644 --- a/core/translate/index.rs +++ b/core/translate/index.rs @@ -6,7 +6,7 @@ use crate::{ util::normalize_ident, vdbe::{ builder::{CursorType, ProgramBuilder, QueryMode}, - insn::{IdxInsertFlags, Insn}, + insn::{IdxInsertFlags, Insn, RegisterOrLiteral}, }, OwnedValue, }; @@ -97,8 +97,7 @@ pub fn translate_create_index( // open the sqlite schema table for writing and create a new entry for the index program.emit_insn(Insn::OpenWriteAsync { cursor_id: sqlite_schema_cursor_id, - root_page: sqlite_table.root_page, - is_new_idx: false, + root_page: RegisterOrLiteral::Literal(sqlite_table.root_page), }); program.emit_insn(Insn::OpenWriteAwait {}); let sql = create_idx_stmt_to_sql(&tbl_name, &idx_name, unique_if_not_exists, &columns); @@ -197,8 +196,7 @@ pub fn translate_create_index( // newly sorted index records. program.emit_insn(Insn::OpenWriteAsync { cursor_id: btree_cursor_id, - root_page: root_page_reg, - is_new_idx: true, + root_page: RegisterOrLiteral::Register(root_page_reg), }); program.emit_insn(Insn::OpenWriteAwait {}); diff --git a/core/translate/insert.rs b/core/translate/insert.rs index d6faeb2c8..371d85be1 100644 --- a/core/translate/insert.rs +++ b/core/translate/insert.rs @@ -9,6 +9,7 @@ use crate::error::SQLITE_CONSTRAINT_PRIMARYKEY; use crate::schema::Table; use crate::util::normalize_ident; use crate::vdbe::builder::{ProgramBuilderOpts, QueryMode}; +use crate::vdbe::insn::RegisterOrLiteral; use crate::vdbe::BranchOffset; use crate::{ schema::{Column, Schema}, @@ -152,8 +153,7 @@ pub fn translate_insert( program.emit_insn(Insn::OpenWriteAsync { cursor_id, - root_page, - is_new_idx: false, + root_page: RegisterOrLiteral::Literal(root_page), }); program.emit_insn(Insn::OpenWriteAwait {}); @@ -169,8 +169,7 @@ pub fn translate_insert( // Single row - populate registers directly program.emit_insn(Insn::OpenWriteAsync { cursor_id, - root_page, - is_new_idx: false, + root_page: RegisterOrLiteral::Literal(root_page), }); program.emit_insn(Insn::OpenWriteAwait {}); diff --git a/core/translate/main_loop.rs b/core/translate/main_loop.rs index d9b7d6a36..e7235e8f0 100644 --- a/core/translate/main_loop.rs +++ b/core/translate/main_loop.rs @@ -102,8 +102,7 @@ pub fn init_loop( let root_page = btree.root_page; program.emit_insn(Insn::OpenWriteAsync { cursor_id, - root_page, - is_new_idx: false, + root_page: root_page.into(), }); program.emit_insn(Insn::OpenWriteAwait {}); } @@ -111,8 +110,7 @@ pub fn init_loop( let root_page = btree.root_page; program.emit_insn(Insn::OpenWriteAsync { cursor_id, - root_page, - is_new_idx: false, + root_page: root_page.into(), }); program.emit_insn(Insn::OpenWriteAwait {}); } @@ -146,16 +144,14 @@ pub fn init_loop( OperationMode::DELETE => { program.emit_insn(Insn::OpenWriteAsync { cursor_id: table_cursor_id, - root_page: table.table.get_root_page(), - is_new_idx: false, + root_page: table.table.get_root_page().into(), }); program.emit_insn(Insn::OpenWriteAwait {}); } OperationMode::UPDATE => { program.emit_insn(Insn::OpenWriteAsync { cursor_id: table_cursor_id, - root_page: table.table.get_root_page(), - is_new_idx: false, + root_page: table.table.get_root_page().into(), }); program.emit_insn(Insn::OpenWriteAwait {}); } @@ -181,16 +177,14 @@ pub fn init_loop( OperationMode::DELETE => { program.emit_insn(Insn::OpenWriteAsync { cursor_id: index_cursor_id, - root_page: index.root_page, - is_new_idx: false, + root_page: index.root_page.into(), }); program.emit_insn(Insn::OpenWriteAwait {}); } OperationMode::UPDATE => { program.emit_insn(Insn::OpenWriteAsync { cursor_id: index_cursor_id, - root_page: index.root_page, - is_new_idx: false, + root_page: index.root_page.into(), }); program.emit_insn(Insn::OpenWriteAwait {}); } diff --git a/core/translate/schema.rs b/core/translate/schema.rs index 6f87937c9..eea0868d0 100644 --- a/core/translate/schema.rs +++ b/core/translate/schema.rs @@ -7,6 +7,7 @@ use crate::translate::ProgramBuilderOpts; use crate::translate::QueryMode; use crate::util::PRIMARY_KEY_AUTOMATIC_INDEX_NAME_PREFIX; use crate::vdbe::builder::CursorType; +use crate::vdbe::insn::RegisterOrLiteral; use crate::vdbe::insn::{CmpInsFlags, Insn}; use crate::LimboError; use crate::{bail_parse_error, Result}; @@ -103,8 +104,7 @@ pub fn translate_create_table( ); program.emit_insn(Insn::OpenWriteAsync { cursor_id: sqlite_schema_cursor_id, - root_page: 1, - is_new_idx: false, + root_page: 1usize.into(), }); program.emit_insn(Insn::OpenWriteAwait {}); @@ -501,8 +501,7 @@ pub fn translate_create_virtual_table( ); program.emit_insn(Insn::OpenWriteAsync { cursor_id: sqlite_schema_cursor_id, - root_page: 1, - is_new_idx: false, + root_page: 1usize.into(), }); program.emit_insn(Insn::OpenWriteAwait {}); @@ -581,8 +580,7 @@ pub fn translate_drop_table( ); program.emit_insn(Insn::OpenWriteAsync { cursor_id: sqlite_schema_cursor_id, - root_page: 1, - is_new_idx: false, + root_page: 1usize.into(), }); program.emit_insn(Insn::OpenWriteAwait {}); diff --git a/core/vdbe/execute.rs b/core/vdbe/execute.rs index 89bd26859..968fea4eb 100644 --- a/core/vdbe/execute.rs +++ b/core/vdbe/execute.rs @@ -29,7 +29,7 @@ use crate::{info, MvCursor, RefValue, Row, StepResult, TransactionState}; use super::insn::{ exec_add, exec_and, exec_bit_and, exec_bit_not, exec_bit_or, exec_boolean_not, exec_concat, exec_divide, exec_multiply, exec_or, exec_remainder, exec_shift_left, exec_shift_right, - exec_subtract, Cookie, + exec_subtract, Cookie, RegisterOrLiteral, }; use super::HaltState; use rand::thread_rng; @@ -3979,12 +3979,23 @@ pub fn op_open_write_async( else { unreachable!("unexpected Insn {:?}", insn) }; + let root_page = match root_page { + RegisterOrLiteral::Literal(lit) => *lit as u64, + RegisterOrLiteral::Register(reg) => match &state.registers[*reg].get_owned_value() { + OwnedValue::Integer(val) => *val as u64, + _ => { + return Err(LimboError::InternalError( + "OpenWriteAsync: the value in root_page is not an integer".into(), + )); + } + }, + }; let (_, cursor_type) = program.cursor_ref.get(*cursor_id).unwrap(); let mut cursors = state.cursors.borrow_mut(); let is_index = cursor_type.is_index(); let mv_cursor = match state.mv_tx_id { Some(tx_id) => { - let table_id = *root_page as u64; + let table_id = root_page; let mv_store = mv_store.unwrap().clone(); let mv_cursor = Rc::new(RefCell::new( MvCursor::new(mv_store.clone(), tx_id, table_id).unwrap(), @@ -3993,7 +4004,7 @@ pub fn op_open_write_async( } None => None, }; - let cursor = BTreeCursor::new(mv_cursor, pager.clone(), *root_page); + let cursor = BTreeCursor::new(mv_cursor, pager.clone(), root_page as usize); if is_index { cursors .get_mut(*cursor_id) diff --git a/core/vdbe/explain.rs b/core/vdbe/explain.rs index f80a442f1..66c68d9c0 100644 --- a/core/vdbe/explain.rs +++ b/core/vdbe/explain.rs @@ -1,4 +1,4 @@ -use crate::vdbe::builder::CursorType; +use crate::vdbe::{builder::CursorType, insn::RegisterOrLiteral}; use super::{Insn, InsnReference, OwnedValue, Program}; use crate::function::{Func, ScalarFunc}; @@ -1134,7 +1134,10 @@ pub fn insn_to_str( } => ( "OpenWriteAsync", *cursor_id as i32, - *root_page as i32, + match root_page { + RegisterOrLiteral::Literal(i) => *i as _, + RegisterOrLiteral::Register(i) => *i as _, + }, 0, OwnedValue::build_text(""), 0, diff --git a/core/vdbe/insn.rs b/core/vdbe/insn.rs index 5ce68f14a..4a6bc1ea4 100644 --- a/core/vdbe/insn.rs +++ b/core/vdbe/insn.rs @@ -76,6 +76,18 @@ impl IdxInsertFlags { } } +#[derive(Clone, Copy, Debug)] +pub enum RegisterOrLiteral { + Register(usize), + Literal(T), +} + +impl From for RegisterOrLiteral { + fn from(value: PageIdx) -> Self { + RegisterOrLiteral::Literal(value) + } +} + #[derive(Description, Debug)] pub enum Insn { /// Initialize the program state and jump to the given PC. @@ -640,8 +652,7 @@ pub enum Insn { OpenWriteAsync { cursor_id: CursorID, - root_page: PageIdx, - is_new_idx: bool, + root_page: RegisterOrLiteral, }, OpenWriteAwait {}, @@ -1296,8 +1307,6 @@ impl Insn { Insn::IdxGT { .. } => execute::op_idx_gt, Insn::IdxLE { .. } => execute::op_idx_le, Insn::IdxLT { .. } => execute::op_idx_lt, - Insn::IdxInsertAsync { .. } => execute::op_idx_insert_async, - Insn::IdxInsertAwait { .. } => execute::op_idx_insert_await, Insn::DecrJumpZero { .. } => execute::op_decr_jump_zero, Insn::AggStep { .. } => execute::op_agg_step, @@ -1315,7 +1324,8 @@ impl Insn { Insn::Yield { .. } => execute::op_yield, Insn::InsertAsync { .. } => execute::op_insert_async, Insn::InsertAwait { .. } => execute::op_insert_await, - + Insn::IdxInsertAsync { .. } => execute::op_idx_insert_async, + Insn::IdxInsertAwait { .. } => execute::op_idx_insert_await, Insn::DeleteAsync { .. } => execute::op_delete_async, Insn::DeleteAwait { .. } => execute::op_delete_await, From 2c3fd509feb5142bf2ef77aec756f58c445f89d9 Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Sun, 30 Mar 2025 16:29:06 -0400 Subject: [PATCH 08/16] Remove unused imports and consolidate ordering comparison --- core/storage/btree.rs | 9 ++++----- core/translate/schema.rs | 1 - 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/core/storage/btree.rs b/core/storage/btree.rs index 0f094ac5b..002459955 100644 --- a/core/storage/btree.rs +++ b/core/storage/btree.rs @@ -2017,6 +2017,8 @@ impl BTreeCursor { { BTreeCell::IndexInteriorCell(IndexInteriorCell { payload, .. }) | BTreeCell::IndexLeafCell(IndexLeafCell { payload, .. }) => { + // TODO: implement efficient comparison of records + // e.g. https://github.com/sqlite/sqlite/blob/master/src/vdbeaux.c#L4719 read_record( payload, self.get_immutable_record_or_create().as_mut().unwrap(), @@ -2027,10 +2029,7 @@ impl BTreeCursor { self.get_immutable_record().as_ref().unwrap().get_values(), ); match order { - Ordering::Less => { - break; - } - Ordering::Equal => { + Ordering::Less | Ordering::Equal => { break; } Ordering::Greater => {} @@ -2044,7 +2043,7 @@ impl BTreeCursor { } pub fn seek_end(&mut self) -> Result> { - assert!(self.mv_cursor.is_none()); + assert!(self.mv_cursor.is_none()); // unsure about this -_- self.move_to_root(); loop { let mem_page = self.stack.top(); diff --git a/core/translate/schema.rs b/core/translate/schema.rs index eea0868d0..3d5aa79db 100644 --- a/core/translate/schema.rs +++ b/core/translate/schema.rs @@ -7,7 +7,6 @@ use crate::translate::ProgramBuilderOpts; use crate::translate::QueryMode; use crate::util::PRIMARY_KEY_AUTOMATIC_INDEX_NAME_PREFIX; use crate::vdbe::builder::CursorType; -use crate::vdbe::insn::RegisterOrLiteral; use crate::vdbe::insn::{CmpInsFlags, Insn}; use crate::LimboError; use crate::{bail_parse_error, Result}; From 007fbe8cc73c80af004d25f4e72f6c4e5682c65f Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Sun, 30 Mar 2025 19:52:54 -0400 Subject: [PATCH 09/16] Fix unique index issue and prealloc in sql string for schema --- core/storage/btree.rs | 34 ++++++++++++++++++++-------------- core/translate/index.rs | 14 +++++++------- 2 files changed, 27 insertions(+), 21 deletions(-) diff --git a/core/storage/btree.rs b/core/storage/btree.rs index 002459955..d4f58151a 100644 --- a/core/storage/btree.rs +++ b/core/storage/btree.rs @@ -2501,22 +2501,28 @@ impl BTreeCursor { /// Search for a key in an Index Btree. Looking up indexes that need to be unique, we cannot compare the rowid pub fn key_exists_in_index(&mut self, key: &ImmutableRecord) -> Result> { return_if_io!(self.do_seek(SeekKey::IndexKey(key), SeekOp::GE)); - if let Some(record) = self.record().as_ref() { - // get existing record, excluding the rowid - assert!(record.len() > 0); - let existing_key = &record.get_values()[..record.count() - 1]; - let inserted_key_vals = &key.get_values(); - if existing_key - .iter() - .zip(inserted_key_vals.iter()) - .all(|(a, b)| a == b) - { - return Ok(CursorResult::Ok(true)); // duplicate + + let record_opt = self.record(); + match record_opt.as_ref() { + Some(record) => { + // Existing record found — compare prefix + let existing_key = &record.get_values()[..record.count().saturating_sub(1)]; + let inserted_key_vals = &key.get_values(); + if existing_key + .iter() + .zip(inserted_key_vals.iter()) + .all(|(a, b)| a == b) + { + return Ok(CursorResult::Ok(true)); // duplicate + } + } + None => { + // Cursor not pointing at a record — table is empty or past last + return Ok(CursorResult::Ok(false)); } - } else { - return Err(LimboError::InvalidArgument("Expected Record key".into())); } - Ok(CursorResult::Ok(false)) // no matching key found + + Ok(CursorResult::Ok(false)) // not a duplicate } pub fn exists(&mut self, key: &OwnedValue) -> Result> { diff --git a/core/translate/index.rs b/core/translate/index.rs index c9a474cab..32d7cd2e9 100644 --- a/core/translate/index.rs +++ b/core/translate/index.rs @@ -62,13 +62,13 @@ pub fn translate_create_index( unique: unique_if_not_exists.0, }); - // Allocate the necessary cursors. + // Allocate the necessary cursors: // - // 1. sqlite_schema_cursor_id - for the sqlite_schema table - // 2. btree_cursor_id - for the index btree - // 3. table_cursor_id - for the table we are creating the index on - // 4. sorter_cursor_id - for the sorter - // 5. pseudo_cursor_id - for the pseudo table to store the sorted index values + // 1. sqlite_schema_cursor_id - sqlite_schema table + // 2. btree_cursor_id - new index btree + // 3. table_cursor_id - table we are creating the index on + // 4. sorter_cursor_id - sorter + // 5. pseudo_cursor_id - pseudo table to store the sorted index values let sqlite_table = schema.get_btree_table(SQLITE_TABLEID).unwrap(); let sqlite_schema_cursor_id = program.alloc_cursor_id( Some(SQLITE_TABLEID.to_owned()), @@ -292,7 +292,7 @@ fn create_idx_stmt_to_sql( unique_if_not_exists: (bool, bool), cols: &[((usize, &Column), SortOrder)], ) -> String { - let mut sql = String::new(); + let mut sql = String::with_capacity(128); sql.push_str("CREATE "); if unique_if_not_exists.0 { sql.push_str("UNIQUE "); From 068ab4ab2779b3e6deebad4b38f375585b8d7a31 Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Mon, 31 Mar 2025 20:48:02 -0400 Subject: [PATCH 10/16] Refactor btree to reuse existing insert and seek with idx keys --- core/storage/btree.rs | 259 ++++++++++++++++++++++-------------------- core/vdbe/execute.rs | 11 +- 2 files changed, 142 insertions(+), 128 deletions(-) diff --git a/core/storage/btree.rs b/core/storage/btree.rs index d4f58151a..11a80ef4b 100644 --- a/core/storage/btree.rs +++ b/core/storage/btree.rs @@ -171,6 +171,50 @@ enum ReadPayloadOverflow { }, } +#[derive(Clone, Debug)] +pub enum BTreeKey<'a> { + TableRowId((u64, Option<&'a ImmutableRecord>)), + IndexKey(&'a ImmutableRecord), +} + +impl<'a> BTreeKey<'_> { + pub fn new_table_rowid(rowid: u64, record: Option<&'a ImmutableRecord>) -> BTreeKey<'a> { + BTreeKey::TableRowId((rowid, record)) + } + pub fn new_index_key(record: &'a ImmutableRecord) -> BTreeKey<'a> { + BTreeKey::IndexKey(record) + } + fn get_record(&self) -> Option<&'_ ImmutableRecord> { + match self { + BTreeKey::TableRowId((_, record)) => *record, + BTreeKey::IndexKey(record) => Some(record), + } + } + + fn maybe_rowid(&self) -> Option { + match self { + BTreeKey::TableRowId((rowid, _)) => Some(*rowid), + BTreeKey::IndexKey(_) => None, + } + } + + /// Assert that the key is a rowid and return it. + fn to_rowid(&self) -> u64 { + match self { + BTreeKey::TableRowId((rowid, _)) => *rowid, + BTreeKey::IndexKey(_) => panic!("BTreeKey::assert_rowid() called on IndexKey"), + } + } + + /// Assert that the key is an index key and return it. + fn to_index_key_values(&self) -> &'_ Vec { + match self { + BTreeKey::TableRowId(_) => panic!("BTreeKey::assert_index_key() called on TableRowId"), + BTreeKey::IndexKey(key) => key.get_values(), + } + } +} + #[derive(Clone)] struct BalanceInfo { /// Old pages being balanced. @@ -1047,72 +1091,13 @@ impl BTreeCursor { } } - pub fn insert_index_key(&mut self, key: &ImmutableRecord) -> Result> { - if let CursorState::None = &self.state { - self.state = CursorState::Write(WriteInfo::new()); - } - - let ret = loop { - let write_state = self.state.mut_write_info().unwrap().state; - match write_state { - WriteState::Start => { - let page = self.stack.top(); - return_if_locked!(page); - page.set_dirty(); - self.pager.add_dirty(page.get().id); - let page = page.get().contents.as_mut().unwrap(); - - assert!(matches!(page.page_type(), PageType::IndexLeaf)); - let cell_idx = self.find_index_cell(page, key); - let mut cell_payload: Vec = Vec::new(); - fill_cell_payload( - page.page_type(), - None, - &mut cell_payload, - key, - self.usable_space() as u16, - self.pager.clone(), - ); - // insert - let overflow = { - debug!( - "insert_index_key(overflow, cell_count={})", - page.cell_count() - ); - insert_into_cell( - page, - cell_payload.as_slice(), - cell_idx, - self.usable_space() as u16, - )?; - page.overflow_cells.len() - }; - let write_info = self.state.mut_write_info().unwrap(); - write_info.state = if overflow > 0 { - WriteState::BalanceStart - } else { - WriteState::Finish - }; - } - WriteState::BalanceStart - | WriteState::BalanceNonRoot - | WriteState::BalanceNonRootWaitLoadPages => { - return_if_io!(self.balance()); - } - WriteState::Finish => break Ok(CursorResult::Ok(())), - } - }; - self.state = CursorState::None; - ret - } - /// Insert a record into the btree. /// If the insert operation overflows the page, it will be split and the btree will be balanced. - fn insert_into_page( - &mut self, - key: &OwnedValue, - record: &ImmutableRecord, - ) -> Result> { + fn insert_into_page(&mut self, bkey: &BTreeKey) -> Result> { + let record = bkey + .get_record() + .expect("expected record present on insert"); + if let CursorState::None = &self.state { self.state = CursorState::Write(WriteInfo::new()); } @@ -1128,10 +1113,6 @@ impl BTreeCursor { WriteState::Start => { let page = self.stack.top(); return_if_locked_maybe_load!(self.pager, page); - let int_key = match key { - OwnedValue::Integer(i) => *i as u64, - _ => unreachable!("btree tables are indexed by integers!"), - }; // get page and find cell let (cell_idx, page_type) = { @@ -1141,23 +1122,27 @@ impl BTreeCursor { self.pager.add_dirty(page.get().id); let page = page.get().contents.as_mut().unwrap(); - assert!(matches!(page.page_type(), PageType::TableLeaf)); + assert!(matches!( + page.page_type(), + PageType::TableLeaf | PageType::IndexLeaf + )); // find cell - (self.find_cell(page, int_key), page.page_type()) + (self.find_cell(page, bkey), page.page_type()) }; tracing::debug!("insert_into_page(cell_idx={})", cell_idx); // if the cell index is less than the total cells, check: if its an existing // rowid, we are going to update / overwrite the cell if cell_idx < page.get_contents().cell_count() { - if let BTreeCell::TableLeafCell(tbl_leaf) = page.get_contents().cell_get( + match page.get_contents().cell_get( cell_idx, payload_overflow_threshold_max(page_type, self.usable_space() as u16), payload_overflow_threshold_min(page_type, self.usable_space() as u16), self.usable_space(), )? { - if tbl_leaf._rowid == int_key { + BTreeCell::TableLeafCell(tbl_leaf) => { + if tbl_leaf._rowid == bkey.to_rowid() { tracing::debug!("insert_into_page: found exact match with cell_idx={cell_idx}, overwriting"); self.overwrite_cell(page.clone(), cell_idx, record)?; self.state @@ -1167,12 +1152,37 @@ impl BTreeCursor { continue; } } + BTreeCell::IndexLeafCell(idx_leaf) => { + read_record( + idx_leaf.payload, + self.get_immutable_record_or_create().as_mut().unwrap(), + ) + .expect("failed to read record"); + if compare_immutable( + record.get_values(), + self.get_immutable_record() + .as_ref() + .unwrap() + .get_values() + ) == Ordering::Equal { + + tracing::debug!("insert_into_page: found exact match with cell_idx={cell_idx}, overwriting"); + self.overwrite_cell(page.clone(), cell_idx, record)?; + self.state + .mut_write_info() + .expect("expected write info") + .state = WriteState::Finish; + continue; + } + } + other => panic!("unexpected cell type, expected TableLeaf or IndexLeaf, found: {:?}", other), + } } // insert cell let mut cell_payload: Vec = Vec::with_capacity(record.len() + 4); fill_cell_payload( page_type, - Some(int_key), + bkey.maybe_rowid(), &mut cell_payload, record, self.usable_space() as u16, @@ -1971,8 +1981,7 @@ impl BTreeCursor { } /// Find the index of the cell in the page that contains the given rowid. - /// BTree tables only. - fn find_cell(&self, page: &PageContent, int_key: u64) -> usize { + fn find_cell(&self, page: &PageContent, key: &BTreeKey) -> usize { let mut cell_idx = 0; let cell_count = page.cell_count(); while cell_idx < cell_count { @@ -1986,35 +1995,15 @@ impl BTreeCursor { .unwrap() { BTreeCell::TableLeafCell(cell) => { - if int_key <= cell._rowid { + if key.to_rowid() <= cell._rowid { break; } } BTreeCell::TableInteriorCell(cell) => { - if int_key <= cell._rowid { + if key.to_rowid() <= cell._rowid { break; } } - _ => todo!(), - } - cell_idx += 1; - } - cell_idx - } - - fn find_index_cell(&self, page: &PageContent, key: &ImmutableRecord) -> usize { - let mut cell_idx = 0; - let cell_count = page.cell_count(); - while cell_idx < cell_count { - match page - .cell_get( - cell_idx, - payload_overflow_threshold_max(page.page_type(), self.usable_space() as u16), - payload_overflow_threshold_min(page.page_type(), self.usable_space() as u16), - self.usable_space(), - ) - .unwrap() - { BTreeCell::IndexInteriorCell(IndexInteriorCell { payload, .. }) | BTreeCell::IndexLeafCell(IndexLeafCell { payload, .. }) => { // TODO: implement efficient comparison of records @@ -2025,7 +2014,7 @@ impl BTreeCursor { ) .expect("failed to read record"); let order = compare_immutable( - key.get_values(), + key.to_index_key_values(), self.get_immutable_record().as_ref().unwrap().get_values(), ); match order { @@ -2035,7 +2024,6 @@ impl BTreeCursor { Ordering::Greater => {} } } - _ => unreachable!("Expected Index cell types"), } cell_idx += 1; } @@ -2158,28 +2146,34 @@ impl BTreeCursor { pub fn insert( &mut self, - key: &OwnedValue, - record: &ImmutableRecord, + key: &BTreeKey, moved_before: bool, /* Indicate whether it's necessary to traverse to find the leaf page */ ) -> Result> { - let int_key = match key { - OwnedValue::Integer(i) => i, - _ => unreachable!("btree tables are indexed by integers!"), - }; match &self.mv_cursor { - Some(mv_cursor) => { - let row_id = - crate::mvcc::database::RowID::new(self.table_id() as u64, *int_key as u64); - let record_buf = record.get_payload().to_vec(); - let row = crate::mvcc::database::Row::new(row_id, record_buf); - mv_cursor.borrow_mut().insert(row).unwrap(); - } + Some(mv_cursor) => match key.maybe_rowid() { + Some(rowid) => { + let row_id = crate::mvcc::database::RowID::new(self.table_id() as u64, rowid); + let record_buf = key.get_record().unwrap().get_payload().to_vec(); + let row = crate::mvcc::database::Row::new(row_id, record_buf); + mv_cursor.borrow_mut().insert(row).unwrap(); + } + None => todo!("Support mvcc inserts with index btrees"), + }, None => { if !moved_before { - return_if_io!(self.move_to(SeekKey::TableRowId(*int_key as u64), SeekOp::EQ)); + return_if_io!(self.move_to( + match key { + BTreeKey::IndexKey(_) => SeekKey::IndexKey(key.get_record().unwrap()), + BTreeKey::TableRowId(_) => SeekKey::TableRowId(key.to_rowid()), + }, + SeekOp::EQ + )); + } + return_if_io!(self.insert_into_page(key)); + if key.maybe_rowid().is_some() { + let int_key = key.to_rowid(); + self.rowid.replace(Some(int_key)); } - return_if_io!(self.insert_into_page(key, record)); - self.rowid.replace(Some(*int_key as u64)); } }; Ok(CursorResult::Ok(())) @@ -2543,7 +2537,7 @@ impl BTreeCursor { OwnedValue::Integer(i) => *i as u64, _ => unreachable!("btree tables are indexed by integers!"), }; - let cell_idx = self.find_cell(contents, int_key); + let cell_idx = self.find_cell(contents, &BTreeKey::new_table_rowid(int_key, None)); if cell_idx >= contents.cell_count() { Ok(CursorResult::Ok(false)) } else { @@ -4199,12 +4193,15 @@ mod tests { pager.deref(), ) .unwrap(); - let key = OwnedValue::Integer(*key); let value = ImmutableRecord::from_registers(&[Register::OwnedValue( OwnedValue::Blob(vec![0; *size]), )]); tracing::info!("insert key:{}", key); - run_until_done(|| cursor.insert(&key, &value, true), pager.deref()).unwrap(); + run_until_done( + || cursor.insert(&BTreeKey::new_table_rowid(*key, Some(&value)), true), + pager.deref(), + ) + .unwrap(); tracing::info!( "=========== btree ===========\n{}\n\n", format_btree(pager.clone(), root_page, 0) @@ -4279,12 +4276,14 @@ mod tests { pager.deref(), ) .unwrap(); - - let key = OwnedValue::Integer(key); let value = ImmutableRecord::from_registers(&[Register::OwnedValue( OwnedValue::Blob(vec![0; size]), )]); - run_until_done(|| cursor.insert(&key, &value, true), pager.deref()).unwrap(); + run_until_done( + || cursor.insert(&BTreeKey::new_table_rowid(key as u64, Some(&value)), true), + pager.deref(), + ) + .unwrap(); if matches!(validate_btree(pager.clone(), root_page), (_, false)) { panic!("invalid btree"); } @@ -5170,7 +5169,11 @@ mod tests { pager.deref(), ) .unwrap(); - run_until_done(|| cursor.insert(&key, &value, true), pager.deref()).unwrap(); + run_until_done( + || cursor.insert(&BTreeKey::new_table_rowid(i as u64, Some(&value)), true), + pager.deref(), + ) + .unwrap(); keys.push(i); } if matches!(validate_btree(pager.clone(), root_page), (_, false)) { @@ -5248,7 +5251,11 @@ mod tests { ) .unwrap(); - run_until_done(|| cursor.insert(&key, &value, true), pager.deref()).unwrap(); + run_until_done( + || cursor.insert(&BTreeKey::new_table_rowid(i as u64, Some(&value)), true), + pager.deref(), + ) + .unwrap(); } match validate_btree(pager.clone(), root_page) { @@ -5326,7 +5333,11 @@ mod tests { pager.deref(), ) .unwrap(); - run_until_done(|| cursor.insert(&key, &value, true), pager.deref()).unwrap(); + run_until_done( + || cursor.insert(&BTreeKey::new_table_rowid(i as u64, Some(&value)), true), + pager.deref(), + ) + .unwrap(); tracing::debug!( "=========== btree after ===========\n{}\n\n", format_btree(pager.clone(), root_page, 0) diff --git a/core/vdbe/execute.rs b/core/vdbe/execute.rs index 968fea4eb..f446da777 100644 --- a/core/vdbe/execute.rs +++ b/core/vdbe/execute.rs @@ -11,7 +11,7 @@ use std::{borrow::BorrowMut, rc::Rc}; use crate::pseudo::PseudoCursor; use crate::result::LimboResult; use crate::schema::{affinity, Affinity}; -use crate::storage::btree::BTreeCursor; +use crate::storage::btree::{BTreeCursor, BTreeKey}; use crate::storage::wal::CheckpointResult; use crate::types::{ AggContext, Cursor, CursorResult, ExternalAggState, OwnedValue, SeekKey, SeekOp, @@ -3666,11 +3666,14 @@ pub fn op_insert_async( Register::Record(r) => r, _ => unreachable!("Not a record! Cannot insert a non record value."), }; - let key = &state.registers[*key_reg]; + let key = match &state.registers[*key_reg].get_owned_value() { + OwnedValue::Integer(i) => *i, + _ => unreachable!("expected integer key"), + }; // NOTE(pere): Sending moved_before == true is okay because we moved before but // if we were to set to false after starting a balance procedure, it might // leave undefined state. - return_if_io!(cursor.insert(key.get_owned_value(), record, true)); + return_if_io!(cursor.insert(&BTreeKey::new_table_rowid(key as u64, Some(record)), true)); } state.pc += 1; Ok(InsnFunctionStepResult::Step) @@ -3765,7 +3768,7 @@ pub fn op_idx_insert_async( flags.has(IdxInsertFlags::USE_SEEK) }; // insert record as key - return_if_io!(cursor.insert_index_key(record)); + return_if_io!(cursor.insert(&BTreeKey::new_index_key(record), moved_before)); } state.pc += 1; } From a2b9d8d371c50b931e4bce18f0ea3f065e14df45 Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Mon, 31 Mar 2025 20:55:37 -0400 Subject: [PATCH 11/16] Use Correct flag on idx insert to prevent seeking --- core/vdbe/mod.rs | 44 +++++++++++++++++++++----------------------- 1 file changed, 21 insertions(+), 23 deletions(-) diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index 673b836a4..8794b208a 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -394,7 +394,7 @@ impl Program { } mv_transactions.clear(); } - return Ok(StepResult::Done); + Ok(StepResult::Done) } else { let connection = self .connection @@ -408,30 +408,28 @@ impl Program { ); if program_state.halt_state.is_some() { self.step_end_write_txn(&pager, &mut program_state.halt_state, connection.deref()) - } else { - if auto_commit { - let current_state = connection.transaction_state.borrow().clone(); - match current_state { - TransactionState::Write => self.step_end_write_txn( - &pager, - &mut program_state.halt_state, - connection.deref(), - ), - TransactionState::Read => { - connection.transaction_state.replace(TransactionState::None); - pager.end_read_tx()?; - Ok(StepResult::Done) - } - TransactionState::None => Ok(StepResult::Done), + } else if auto_commit { + let current_state = connection.transaction_state.borrow().clone(); + match current_state { + TransactionState::Write => self.step_end_write_txn( + &pager, + &mut program_state.halt_state, + connection.deref(), + ), + TransactionState::Read => { + connection.transaction_state.replace(TransactionState::None); + pager.end_read_tx()?; + Ok(StepResult::Done) } - } else { - if self.change_cnt_on { - if let Some(conn) = self.connection.upgrade() { - conn.set_changes(self.n_change.get()); - } - } - Ok(StepResult::Done) + TransactionState::None => Ok(StepResult::Done), } + } else { + if self.change_cnt_on { + if let Some(conn) = self.connection.upgrade() { + conn.set_changes(self.n_change.get()); + } + } + Ok(StepResult::Done) } } } From 6b42808f1a86ef297de57a39a2ad2f1fbcf7a09a Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Mon, 31 Mar 2025 21:57:56 -0400 Subject: [PATCH 12/16] Dont re-seek if we are inserting a new unique index --- core/storage/btree.rs | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/core/storage/btree.rs b/core/storage/btree.rs index 11a80ef4b..7dfb82005 100644 --- a/core/storage/btree.rs +++ b/core/storage/btree.rs @@ -2161,13 +2161,15 @@ impl BTreeCursor { }, None => { if !moved_before { - return_if_io!(self.move_to( - match key { - BTreeKey::IndexKey(_) => SeekKey::IndexKey(key.get_record().unwrap()), - BTreeKey::TableRowId(_) => SeekKey::TableRowId(key.to_rowid()), - }, - SeekOp::EQ - )); + match key { + BTreeKey::IndexKey(_) => { + return_if_io!(self + .move_to(SeekKey::IndexKey(key.get_record().unwrap()), SeekOp::GE)) + } + BTreeKey::TableRowId(_) => return_if_io!( + self.move_to(SeekKey::TableRowId(key.to_rowid()), SeekOp::EQ) + ), + } } return_if_io!(self.insert_into_page(key)); if key.maybe_rowid().is_some() { From abc97c877482033b30db263c9c924c2147be738e Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Tue, 1 Apr 2025 10:27:30 -0400 Subject: [PATCH 13/16] Add doc comments to new btree key enum and remove unused lifetimes --- core/storage/btree.rs | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/core/storage/btree.rs b/core/storage/btree.rs index 7dfb82005..ae710859b 100644 --- a/core/storage/btree.rs +++ b/core/storage/btree.rs @@ -177,13 +177,19 @@ pub enum BTreeKey<'a> { IndexKey(&'a ImmutableRecord), } -impl<'a> BTreeKey<'_> { - pub fn new_table_rowid(rowid: u64, record: Option<&'a ImmutableRecord>) -> BTreeKey<'a> { +impl BTreeKey<'_> { + /// Create a new table rowid key from a rowid and an optional immutable record. + /// The record is optional because it may not be available when the key is created. + pub fn new_table_rowid(rowid: u64, record: Option<&ImmutableRecord>) -> BTreeKey<'_> { BTreeKey::TableRowId((rowid, record)) } - pub fn new_index_key(record: &'a ImmutableRecord) -> BTreeKey<'a> { + + /// Create a new index key from an immutable record. + pub fn new_index_key(record: &ImmutableRecord) -> BTreeKey<'_> { BTreeKey::IndexKey(record) } + + /// Get the record, if present. Index will always be present, fn get_record(&self) -> Option<&'_ ImmutableRecord> { match self { BTreeKey::TableRowId((_, record)) => *record, @@ -191,6 +197,7 @@ impl<'a> BTreeKey<'_> { } } + /// Get the rowid, if present. Index will never be present. fn maybe_rowid(&self) -> Option { match self { BTreeKey::TableRowId((rowid, _)) => Some(*rowid), @@ -198,18 +205,18 @@ impl<'a> BTreeKey<'_> { } } - /// Assert that the key is a rowid and return it. + /// Assert that the key is an integer rowid and return it. fn to_rowid(&self) -> u64 { match self { BTreeKey::TableRowId((rowid, _)) => *rowid, - BTreeKey::IndexKey(_) => panic!("BTreeKey::assert_rowid() called on IndexKey"), + BTreeKey::IndexKey(_) => panic!("BTreeKey::to_rowid called on IndexKey"), } } /// Assert that the key is an index key and return it. fn to_index_key_values(&self) -> &'_ Vec { match self { - BTreeKey::TableRowId(_) => panic!("BTreeKey::assert_index_key() called on TableRowId"), + BTreeKey::TableRowId(_) => panic!("BTreeKey::to_index_key called on TableRowId"), BTreeKey::IndexKey(key) => key.get_values(), } } From bd04b10f172d1edc3969b53453351a43395b3b73 Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Tue, 1 Apr 2025 22:24:30 -0400 Subject: [PATCH 14/16] Fix btree tests to adapt to new type for BTreeKey --- core/storage/btree.rs | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/core/storage/btree.rs b/core/storage/btree.rs index ae710859b..789f9cb39 100644 --- a/core/storage/btree.rs +++ b/core/storage/btree.rs @@ -4196,7 +4196,7 @@ mod tests { for (key, size) in sequence.iter() { run_until_done( || { - let key = SeekKey::TableRowId(*key as u64); + let key = SeekKey::TableRowId(*key); cursor.move_to(key, SeekOp::EQ) }, pager.deref(), @@ -4217,7 +4217,7 @@ mod tests { ); } for (key, _) in sequence.iter() { - let seek_key = SeekKey::TableRowId(*key as u64); + let seek_key = SeekKey::TableRowId(*key); assert!( matches!( cursor.seek(seek_key, SeekOp::EQ).unwrap(), @@ -5166,7 +5166,6 @@ mod tests { for i in 0..10000 { let mut cursor = BTreeCursor::new(None, pager.clone(), root_page); tracing::info!("INSERT INTO t VALUES ({});", i,); - let key = OwnedValue::Integer(i); let value = ImmutableRecord::from_registers(&[Register::OwnedValue(OwnedValue::Integer(i))]); tracing::trace!("before insert {}", i); @@ -5246,7 +5245,6 @@ mod tests { // Insert 10,000 records in to the BTree. for i in 1..=10000 { let mut cursor = BTreeCursor::new(None, pager.clone(), root_page); - let key = OwnedValue::Integer(i); let value = ImmutableRecord::from_registers(&[Register::OwnedValue(OwnedValue::Text( Text::new("hello world"), ))]); @@ -5323,7 +5321,6 @@ mod tests { for i in 0..iterations { let mut cursor = BTreeCursor::new(None, pager.clone(), root_page); tracing::info!("INSERT INTO t VALUES ({});", i,); - let key = OwnedValue::Integer(i as i64); let value = ImmutableRecord::from_registers(&[Register::OwnedValue(OwnedValue::Text(Text { value: huge_texts[i].as_bytes().to_vec(), From 399994bf6644acb1878cde58ed25538078ea8b49 Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Wed, 2 Apr 2025 21:29:08 -0400 Subject: [PATCH 15/16] Fix ext tests start with no default schema --- testing/cli_tests/extensions.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/testing/cli_tests/extensions.py b/testing/cli_tests/extensions.py index d898908f9..bab8cb74f 100755 --- a/testing/cli_tests/extensions.py +++ b/testing/cli_tests/extensions.py @@ -342,6 +342,7 @@ def test_kv(): # first, create a normal table to ensure no issues limbo.execute_dot("CREATE TABLE other (a,b,c);") limbo.execute_dot("INSERT INTO other values (23,32,23);") + limbo = TestLimboShell() limbo.run_test_fn( "create virtual table t using kv_store;", lambda res: "Virtual table module not found: kv_store" in res, @@ -350,6 +351,7 @@ def test_kv(): limbo.debug_print( "create virtual table t using kv_store;", ) + limbo.run_test_fn(".schema", lambda res: "CREATE VIRTUAL TABLE t" in res) limbo.run_test_fn( "insert into t values ('hello', 'world');", null, @@ -496,12 +498,6 @@ def test_vfs(): "Tested large write to testfs", ) print("Tested large write to testfs") - # Pere: I commented this out because it added an extra row that made the test test_sqlite_vfs_compat fail - # it didn't segfault from my side so maybe this is necessary? - # # open regular db file to ensure we don't segfault when vfs file is dropped - # limbo.execute_dot(".open testing/vfs.db") - # limbo.execute_dot("create table test (id integer primary key, value float);") - # limbo.execute_dot("insert into test (value) values (1.0);") limbo.quit() @@ -548,10 +544,10 @@ if __name__ == "__main__": test_aggregates() test_crypto() test_series() - test_kv() test_ipaddr() test_vfs() test_sqlite_vfs_compat() + test_kv() except Exception as e: print(f"Test FAILED: {e}") cleanup() From 83af71f1409141192c3408591d3a1cf9c4c8a2a2 Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Sat, 5 Apr 2025 11:30:57 -0400 Subject: [PATCH 16/16] Return accidentally deleted comment on SeekGE insn from merge conflict --- core/vdbe/insn.rs | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/core/vdbe/insn.rs b/core/vdbe/insn.rs index 4a6bc1ea4..8d3a9afca 100644 --- a/core/vdbe/insn.rs +++ b/core/vdbe/insn.rs @@ -484,10 +484,10 @@ pub enum Insn { target_pc: BranchOffset, }, - // cursor_id is a cursor pointing to a B-Tree index that uses integer keys, this op writes the value obtained from MakeRecord into the index. - // P3 + P4 are for the original column values that make up that key in unpacked (pre-serialized) form. - // If P5 has the OPFLAG_APPEND bit set, that is a hint to the b-tree layer that this insert is likely to be an append. - // OPFLAG_NCHANGE bit set, then the change counter is incremented by this instruction. If the OPFLAG_NCHANGE bit is clear, then the change counter is unchanged + /// cursor_id is a cursor pointing to a B-Tree index that uses integer keys, this op writes the value obtained from MakeRecord into the index. + /// P3 + P4 are for the original column values that make up that key in unpacked (pre-serialized) form. + /// If P5 has the OPFLAG_APPEND bit set, that is a hint to the b-tree layer that this insert is likely to be an append. + /// OPFLAG_NCHANGE bit set, then the change counter is incremented by this instruction. If the OPFLAG_NCHANGE bit is clear, then the change counter is unchanged IdxInsertAsync { cursor_id: CursorID, record_reg: usize, // P2 the register containing the record to insert @@ -498,6 +498,9 @@ pub enum Insn { IdxInsertAwait { cursor_id: CursorID, }, + + /// The P4 register values beginning with P3 form an unpacked index key that omits the PRIMARY KEY. Compare this key value against the index that P1 is currently pointing to, ignoring the PRIMARY KEY or ROWID fields at the end. + /// If the P1 index entry is greater or equal than the key value then jump to P2. Otherwise fall through to the next instruction. IdxGE { cursor_id: CursorID, start_reg: usize,