Support UPDATE for virtual tables

This commit is contained in:
PThorpe92
2025-03-24 22:14:48 -04:00
parent 2d7a27fbfa
commit b685086cad
3 changed files with 210 additions and 48 deletions

View File

@@ -6,6 +6,7 @@ use std::rc::Rc;
use limbo_sqlite3_parser::ast::{self};
use crate::function::Func;
use crate::schema::Table;
use crate::translate::plan::{DeletePlan, Plan, Search};
use crate::util::exprs_are_equivalent;
use crate::vdbe::builder::ProgramBuilder;
@@ -600,20 +601,41 @@ fn emit_update_insns(
if table_column.primary_key {
program.emit_null(dest, None);
} else {
program.emit_insn(Insn::Column {
cursor_id: *index
.as_ref()
.and_then(|(_, id)| {
if column_idx_in_index.is_some() {
Some(id)
} else {
None
}
})
.unwrap_or(&cursor_id),
column: column_idx_in_index.unwrap_or(idx),
dest,
});
match &table_ref.table {
Table::BTree(_) => {
program.emit_insn(Insn::Column {
cursor_id: *index
.as_ref()
.and_then(|(_, id)| {
if column_idx_in_index.is_some() {
Some(id)
} else {
None
}
})
.unwrap_or(&cursor_id),
column: column_idx_in_index.unwrap_or(idx),
dest,
});
}
Table::Virtual(_) => {
program.emit_insn(Insn::VColumn {
cursor_id: *index
.as_ref()
.and_then(|(_, id)| {
if column_idx_in_index.is_some() {
Some(id)
} else {
None
}
})
.unwrap_or(&cursor_id),
column: column_idx_in_index.unwrap_or(idx),
dest,
});
}
typ => unreachable!("query plan generated on unexpected table type {:?}", typ),
}
}
}
}
@@ -633,13 +655,34 @@ fn emit_update_insns(
count: table_ref.columns().len(),
dest_reg: record_reg,
});
program.emit_insn(Insn::InsertAsync {
cursor: cursor_id,
key_reg: rowid_reg,
record_reg,
flag: 0,
});
program.emit_insn(Insn::InsertAwait { cursor_id });
match &table_ref.table {
Table::BTree(_) => {
program.emit_insn(Insn::InsertAsync {
cursor: cursor_id,
key_reg: rowid_reg,
record_reg,
flag: 0,
});
program.emit_insn(Insn::InsertAwait { cursor_id });
}
Table::Virtual(vtab) => {
let new_rowid = program.alloc_register();
program.emit_insn(Insn::Copy {
src_reg: rowid_reg,
dst_reg: new_rowid,
amount: 0,
});
let arg_count = table_ref.columns().len() + 2;
program.emit_insn(Insn::VUpdate {
cursor_id,
arg_count,
start_reg: record_reg,
vtab_ptr: vtab.implementation.as_ref().ctx as usize,
conflict_action: 0u16,
});
}
_ => unreachable!("unexpected table type"),
}
if let Some(limit_reg) = t_ctx.reg_limit {
program.emit_insn(Insn::DecrJumpZero {

View File

@@ -109,7 +109,7 @@ pub fn init_loop(
program.emit_insn(Insn::OpenReadAwait {});
}
}
(OperationMode::DELETE, Table::BTree(btree)) => {
(OperationMode::DELETE | OperationMode::UPDATE, Table::BTree(btree)) => {
let root_page = btree.root_page;
program.emit_insn(Insn::OpenWriteAsync {
cursor_id,
@@ -131,11 +131,7 @@ pub fn init_loop(
}
program.emit_insn(Insn::OpenWriteAwait {});
}
(OperationMode::SELECT, Table::Virtual(_)) => {
program.emit_insn(Insn::VOpenAsync { cursor_id });
program.emit_insn(Insn::VOpenAwait {});
}
(OperationMode::DELETE, Table::Virtual(_)) => {
(_, Table::Virtual(_)) => {
program.emit_insn(Insn::VOpenAsync { cursor_id });
program.emit_insn(Insn::VOpenAwait {});
}
@@ -158,14 +154,7 @@ pub fn init_loop(
});
program.emit_insn(Insn::OpenReadAwait {});
}
OperationMode::DELETE => {
program.emit_insn(Insn::OpenWriteAsync {
cursor_id: table_cursor_id,
root_page: table.table.get_root_page().into(),
});
program.emit_insn(Insn::OpenWriteAwait {});
}
OperationMode::UPDATE => {
OperationMode::DELETE | OperationMode::UPDATE => {
program.emit_insn(Insn::OpenWriteAsync {
cursor_id: table_cursor_id,
root_page: table.table.get_root_page().into(),
@@ -191,17 +180,10 @@ pub fn init_loop(
});
program.emit_insn(Insn::OpenReadAwait);
}
OperationMode::DELETE => {
OperationMode::UPDATE | OperationMode::DELETE => {
program.emit_insn(Insn::OpenWriteAsync {
cursor_id: index_cursor_id,
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.into(),
root_page: index.root_page,
});
program.emit_insn(Insn::OpenWriteAwait {});
}

View File

@@ -1,4 +1,7 @@
use std::sync::Arc;
use crate::translate::plan::Operation;
use crate::vdbe::BranchOffset;
use crate::{
bail_parse_error,
schema::{Schema, Table},
@@ -8,7 +11,7 @@ use crate::{
};
use limbo_sqlite3_parser::ast::{self, Expr, ResultColumn, SortOrder, Update};
use super::emitter::emit_program;
use super::emitter::{emit_program, Resolver};
use super::optimizer::optimize_plan;
use super::plan::{
Direction, IterationDirection, Plan, ResultSetColumn, TableReference, UpdatePlan,
@@ -53,6 +56,7 @@ pub fn translate_update(
) -> crate::Result<ProgramBuilder> {
let mut plan = prepare_update_plan(schema, body)?;
optimize_plan(&mut plan, schema)?;
let resolver = Resolver::new(syms);
// TODO: freestyling these numbers
let mut program = ProgramBuilder::new(ProgramBuilderOpts {
query_mode,
@@ -65,6 +69,12 @@ pub fn translate_update(
}
pub fn prepare_update_plan(schema: &Schema, body: &mut Update) -> crate::Result<Plan> {
if body.with.is_some() {
bail_parse_error!("WITH clause is not supported");
}
if body.or_conflict.is_some() {
bail_parse_error!("ON CONFLICT clause is not supported");
}
let table_name = &body.tbl_name.name;
let table = match schema.get_table(table_name.0.as_str()) {
Some(table) => table,
@@ -86,7 +96,11 @@ pub fn prepare_update_plan(schema: &Schema, body: &mut Update) -> crate::Result<
})
.unwrap_or(IterationDirection::Forwards);
let table_references = vec![TableReference {
table: Table::BTree(btree_table.clone()),
table: match table.as_ref() {
Table::Virtual(vtab) => Table::Virtual(vtab.clone()),
Table::BTree(btree_table) => Table::BTree(btree_table.clone()),
_ => unreachable!(),
},
identifier: table_name.0.clone(),
op: Operation::Scan {
iter_dir,
@@ -99,8 +113,8 @@ pub fn prepare_update_plan(schema: &Schema, body: &mut Update) -> crate::Result<
.iter_mut()
.map(|set| {
let ident = normalize_ident(set.col_names[0].0.as_str());
let col_index = btree_table
.columns
let col_index = table
.columns()
.iter()
.enumerate()
.find_map(|(i, col)| {
@@ -185,3 +199,126 @@ pub fn prepare_update_plan(schema: &Schema, body: &mut Update) -> crate::Result<
contains_constant_false_condition: false,
}))
}
// fn translate_vtab_update(
// mut program: ProgramBuilder,
// body: &mut Update,
// table: Arc<Table>,
// resolver: &Resolver,
// ) -> crate::Result<ProgramBuilder> {
// let start_label = program.allocate_label();
// program.emit_insn(Insn::Init {
// target_pc: start_label,
// });
// let start_offset = program.offset();
// let vtab = table.virtual_table().unwrap();
// let cursor_id = program.alloc_cursor_id(
// Some(table.get_name().to_string()),
// CursorType::VirtualTable(vtab.clone()),
// );
// let referenced_tables = vec![TableReference {
// table: Table::Virtual(table.virtual_table().unwrap().clone()),
// identifier: table.get_name().to_string(),
// op: Operation::Scan { iter_dir: None },
// join_info: None,
// }];
// program.emit_insn(Insn::VOpenAsync { cursor_id });
// program.emit_insn(Insn::VOpenAwait {});
//
// let argv_start = program.alloc_registers(0);
// let end_label = program.allocate_label();
// let skip_label = program.allocate_label();
// program.emit_insn(Insn::VFilter {
// cursor_id,
// pc_if_empty: end_label,
// args_reg: argv_start,
// arg_count: 0,
// });
//
// let loop_start = program.offset();
// let start_reg = program.alloc_registers(2 + table.columns().len());
// let old_rowid = start_reg;
// let new_rowid = start_reg + 1;
// let column_regs = start_reg + 2;
//
// program.emit_insn(Insn::RowId {
// cursor_id,
// dest: old_rowid,
// });
// program.emit_insn(Insn::RowId {
// cursor_id,
// dest: new_rowid,
// });
//
// for (i, _) in table.columns().iter().enumerate() {
// let dest = column_regs + i;
// program.emit_insn(Insn::VColumn {
// cursor_id,
// column: i,
// dest,
// });
// }
//
// if let Some(ref mut where_clause) = body.where_clause {
// bind_column_references(where_clause, &referenced_tables, None)?;
// translate_condition_expr(
// &mut program,
// &referenced_tables,
// where_clause,
// ConditionMetadata {
// jump_if_condition_is_true: false,
// jump_target_when_true: BranchOffset::Placeholder,
// jump_target_when_false: skip_label,
// },
// resolver,
// )?;
// }
// // prepare updated columns in place
// for expr in body.sets.iter() {
// let Some(col_index) = table.columns().iter().position(|t| {
// t.name
// .as_ref()
// .unwrap()
// .eq_ignore_ascii_case(&expr.col_names[0].0)
// }) else {
// bail_parse_error!("column {} not found", expr.col_names[0].0);
// };
// translate_expr(
// &mut program,
// Some(&referenced_tables),
// &expr.expr,
// column_regs + col_index,
// resolver,
// )?;
// }
//
// let arg_count = 2 + table.columns().len();
// program.emit_insn(Insn::VUpdate {
// cursor_id,
// arg_count,
// start_reg: old_rowid,
// vtab_ptr: vtab.implementation.ctx as usize,
// conflict_action: 0,
// });
//
// program.resolve_label(skip_label, program.offset());
// program.emit_insn(Insn::VNext {
// cursor_id,
// pc_if_next: loop_start,
// });
//
// program.resolve_label(end_label, program.offset());
// program.emit_insn(Insn::Halt {
// err_code: 0,
// description: String::new(),
// });
// program.resolve_label(start_label, program.offset());
// program.emit_insn(Insn::Transaction { write: true });
//
// program.emit_constant_insns();
// program.emit_insn(Insn::Goto {
// target_pc: start_offset,
// });
// program.table_references = referenced_tables.clone();
// Ok(program)
// }