Remove TableReferenceType enum to clean up planner

This commit is contained in:
PThorpe92
2025-02-05 21:30:55 -05:00
parent cd83ac6146
commit ae88d51e6f
8 changed files with 76 additions and 102 deletions

View File

@@ -136,6 +136,7 @@ impl Database {
name: name.to_string(), name: name.to_string(),
implementation: vtab_module, implementation: vtab_module,
columns, columns,
args: None,
}; };
self.syms.borrow_mut().vtabs.insert(name.to_string(), vtab); self.syms.borrow_mut().vtabs.insert(name.to_string(), vtab);
ResultCode::OK ResultCode::OK

View File

@@ -514,6 +514,7 @@ pub type StepResult = vdbe::StepResult;
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct VirtualTable { pub struct VirtualTable {
name: String, name: String,
args: Option<Vec<ast::Expr>>,
pub implementation: Rc<VTabModuleImpl>, pub implementation: Rc<VTabModuleImpl>,
columns: Vec<Column>, columns: Vec<Column>,
} }
@@ -537,7 +538,7 @@ impl VirtualTable {
OwnedValue::Null => Ok(ExtValue::null()), OwnedValue::Null => Ok(ExtValue::null()),
OwnedValue::Integer(i) => Ok(ExtValue::from_integer(*i)), OwnedValue::Integer(i) => Ok(ExtValue::from_integer(*i)),
OwnedValue::Float(f) => Ok(ExtValue::from_float(*f)), OwnedValue::Float(f) => Ok(ExtValue::from_float(*f)),
OwnedValue::Text(t) => Ok(ExtValue::from_text((*t.value).clone())), OwnedValue::Text(t) => Ok(ExtValue::from_text(t.as_str().to_string())),
OwnedValue::Blob(b) => Ok(ExtValue::from_blob((**b).clone())), OwnedValue::Blob(b) => Ok(ExtValue::from_blob((**b).clone())),
other => Err(LimboError::ExtensionError(format!( other => Err(LimboError::ExtensionError(format!(
"Unsupported value type: {:?}", "Unsupported value type: {:?}",

View File

@@ -7,7 +7,7 @@ use crate::vdbe::builder::{ProgramBuilder, ProgramBuilderOpts, QueryMode};
use crate::{schema::Schema, Result, SymbolTable}; use crate::{schema::Schema, Result, SymbolTable};
use sqlite3_parser::ast::{Expr, Limit, QualifiedName}; use sqlite3_parser::ast::{Expr, Limit, QualifiedName};
use super::plan::{TableReference, TableReferenceType}; use super::plan::TableReference;
pub fn translate_delete( pub fn translate_delete(
query_mode: QueryMode, query_mode: QueryMode,
@@ -48,7 +48,6 @@ pub fn prepare_delete_plan(
identifier: table.name.clone(), identifier: table.name.clone(),
op: Operation::Scan { iter_dir: None }, op: Operation::Scan { iter_dir: None },
join_info: None, join_info: None,
reference_type: TableReferenceType::BTreeTable,
}]; }];
let mut where_predicates = vec![]; let mut where_predicates = vec![];

View File

@@ -3,7 +3,7 @@ use sqlite3_parser::ast::{self, UnaryOperator};
#[cfg(feature = "json")] #[cfg(feature = "json")]
use crate::function::JsonFunc; use crate::function::JsonFunc;
use crate::function::{Func, FuncCtx, MathFuncArity, ScalarFunc, VectorFunc}; use crate::function::{Func, FuncCtx, MathFuncArity, ScalarFunc, VectorFunc};
use crate::schema::Type; use crate::schema::{Table, Type};
use crate::util::normalize_ident; use crate::util::normalize_ident;
use crate::vdbe::{ use crate::vdbe::{
builder::ProgramBuilder, builder::ProgramBuilder,
@@ -13,7 +13,7 @@ use crate::vdbe::{
use crate::Result; use crate::Result;
use super::emitter::Resolver; use super::emitter::Resolver;
use super::plan::{Operation, TableReference, TableReferenceType}; use super::plan::{Operation, TableReference};
#[derive(Debug, Clone, Copy)] #[derive(Debug, Clone, Copy)]
pub struct ConditionMetadata { pub struct ConditionMetadata {
@@ -1823,49 +1823,38 @@ pub fn translate_expr(
match table_reference.op { match table_reference.op {
// If we are reading a column from a table, we find the cursor that corresponds to // If we are reading a column from a table, we find the cursor that corresponds to
// the table and read the column from the cursor. // the table and read the column from the cursor.
Operation::Scan { .. } | Operation::Search(_) => { Operation::Scan { .. } | Operation::Search(_) => match &table_reference.table {
match &table_reference.reference_type { Table::BTree(_) => {
TableReferenceType::BTreeTable => { let cursor_id = program.resolve_cursor_id(&table_reference.identifier);
let cursor_id = program.resolve_cursor_id(&table_reference.identifier); if *is_rowid_alias {
if *is_rowid_alias { program.emit_insn(Insn::RowId {
program.emit_insn(Insn::RowId { cursor_id,
cursor_id, dest: target_register,
dest: target_register, });
}); } else {
} else { program.emit_insn(Insn::Column {
program.emit_insn(Insn::Column {
cursor_id,
column: *column,
dest: target_register,
});
}
let Some(column) = table_reference.table.get_column_at(*column) else {
crate::bail_parse_error!("column index out of bounds");
};
maybe_apply_affinity(column.ty, target_register, program);
Ok(target_register)
}
TableReferenceType::VirtualTable { .. } => {
let cursor_id = program.resolve_cursor_id(&table_reference.identifier);
program.emit_insn(Insn::VColumn {
cursor_id, cursor_id,
column: *column, column: *column,
dest: target_register, dest: target_register,
}); });
Ok(target_register)
}
TableReferenceType::Subquery {
result_columns_start_reg,
} => {
program.emit_insn(Insn::Copy {
src_reg: result_columns_start_reg + *column,
dst_reg: target_register,
amount: 0,
});
Ok(target_register)
} }
let Some(column) = table_reference.table.get_column_at(*column) else {
crate::bail_parse_error!("column index out of bounds");
};
maybe_apply_affinity(column.ty, target_register, program);
Ok(target_register)
} }
} Table::Virtual(_) => {
let cursor_id = program.resolve_cursor_id(&table_reference.identifier);
program.emit_insn(Insn::VColumn {
cursor_id,
column: *column,
dest: target_register,
});
Ok(target_register)
}
_ => unreachable!(),
},
// If we are reading a column from a subquery, we instead copy the column from the // If we are reading a column from a subquery, we instead copy the column from the
// subquery's result registers. // subquery's result registers.
Operation::Subquery { Operation::Subquery {

View File

@@ -1,6 +1,7 @@
use sqlite3_parser::ast; use sqlite3_parser::ast;
use crate::{ use crate::{
schema::Table,
translate::result_row::emit_select_result, translate::result_row::emit_select_result,
vdbe::{ vdbe::{
builder::{CursorType, ProgramBuilder}, builder::{CursorType, ProgramBuilder},
@@ -17,7 +18,7 @@ use super::{
order_by::{order_by_sorter_insert, sorter_insert}, order_by::{order_by_sorter_insert, sorter_insert},
plan::{ plan::{
IterationDirection, Operation, Search, SelectPlan, SelectQueryType, TableReference, IterationDirection, Operation, Search, SelectPlan, SelectQueryType, TableReference,
TableReferenceType, WhereTerm, WhereTerm,
}, },
}; };
@@ -78,21 +79,18 @@ pub fn init_loop(
} }
match &table.op { match &table.op {
Operation::Scan { .. } => { Operation::Scan { .. } => {
let ref_type = &table.reference_type;
let cursor_id = program.alloc_cursor_id( let cursor_id = program.alloc_cursor_id(
Some(table.identifier.clone()), Some(table.identifier.clone()),
match ref_type { match &table.table {
TableReferenceType::BTreeTable => { Table::BTree(_) => CursorType::BTreeTable(table.btree().unwrap().clone()),
CursorType::BTreeTable(table.btree().unwrap().clone()) Table::Virtual(_) => {
}
TableReferenceType::VirtualTable { .. } => {
CursorType::VirtualTable(table.virtual_table().unwrap().clone()) CursorType::VirtualTable(table.virtual_table().unwrap().clone())
} }
other => panic!("Invalid table reference type in Scan: {:?}", other), other => panic!("Invalid table reference type in Scan: {:?}", other),
}, },
); );
match (mode, ref_type) { match (mode, &table.table) {
(OperationMode::SELECT, TableReferenceType::BTreeTable) => { (OperationMode::SELECT, Table::BTree(_)) => {
let root_page = table.btree().unwrap().root_page; let root_page = table.btree().unwrap().root_page;
program.emit_insn(Insn::OpenReadAsync { program.emit_insn(Insn::OpenReadAsync {
cursor_id, cursor_id,
@@ -100,7 +98,7 @@ pub fn init_loop(
}); });
program.emit_insn(Insn::OpenReadAwait {}); program.emit_insn(Insn::OpenReadAwait {});
} }
(OperationMode::DELETE, TableReferenceType::BTreeTable) => { (OperationMode::DELETE, Table::BTree(_)) => {
let root_page = table.btree().unwrap().root_page; let root_page = table.btree().unwrap().root_page;
program.emit_insn(Insn::OpenWriteAsync { program.emit_insn(Insn::OpenWriteAsync {
cursor_id, cursor_id,
@@ -108,7 +106,7 @@ pub fn init_loop(
}); });
program.emit_insn(Insn::OpenWriteAwait {}); program.emit_insn(Insn::OpenWriteAwait {});
} }
(OperationMode::SELECT, TableReferenceType::VirtualTable { .. }) => { (OperationMode::SELECT, Table::Virtual(_)) => {
program.emit_insn(Insn::VOpenAsync { cursor_id }); program.emit_insn(Insn::VOpenAsync { cursor_id });
program.emit_insn(Insn::VOpenAwait {}); program.emit_insn(Insn::VOpenAwait {});
} }
@@ -258,10 +256,9 @@ pub fn open_loop(
} }
} }
Operation::Scan { iter_dir } => { Operation::Scan { iter_dir } => {
let ref_type = &table.reference_type;
let cursor_id = program.resolve_cursor_id(&table.identifier); let cursor_id = program.resolve_cursor_id(&table.identifier);
if !matches!(ref_type, TableReferenceType::VirtualTable { .. }) { if !matches!(&table.table, Table::Virtual(_)) {
if iter_dir if iter_dir
.as_ref() .as_ref()
.is_some_and(|dir| *dir == IterationDirection::Backwards) .is_some_and(|dir| *dir == IterationDirection::Backwards)
@@ -271,8 +268,8 @@ pub fn open_loop(
program.emit_insn(Insn::RewindAsync { cursor_id }); program.emit_insn(Insn::RewindAsync { cursor_id });
} }
} }
match ref_type { match &table.table {
TableReferenceType::BTreeTable => program.emit_insn( Table::BTree(_) => program.emit_insn(
if iter_dir if iter_dir
.as_ref() .as_ref()
.is_some_and(|dir| *dir == IterationDirection::Backwards) .is_some_and(|dir| *dir == IterationDirection::Backwards)
@@ -288,13 +285,18 @@ pub fn open_loop(
} }
}, },
), ),
TableReferenceType::VirtualTable { args, .. } => { Table::Virtual(ref table) => {
let args = if let Some(args) = table.args.as_ref() {
args
} else {
&vec![]
};
let start_reg = program.alloc_registers(args.len()); let start_reg = program.alloc_registers(args.len());
let mut cur_reg = start_reg; let mut cur_reg = start_reg;
for arg in args { for arg in args {
let reg = cur_reg; let reg = cur_reg;
cur_reg += 1; cur_reg += 1;
translate_expr(program, Some(tables), arg, reg, &t_ctx.resolver)?; translate_expr(program, Some(tables), &arg, reg, &t_ctx.resolver)?;
} }
program.emit_insn(Insn::VFilter { program.emit_insn(Insn::VFilter {
cursor_id, cursor_id,
@@ -722,11 +724,10 @@ pub fn close_loop(
}); });
} }
Operation::Scan { iter_dir, .. } => { Operation::Scan { iter_dir, .. } => {
let ref_type = &table.reference_type;
program.resolve_label(loop_labels.next, program.offset()); program.resolve_label(loop_labels.next, program.offset());
let cursor_id = program.resolve_cursor_id(&table.identifier); let cursor_id = program.resolve_cursor_id(&table.identifier);
match ref_type { match &table.table {
TableReferenceType::BTreeTable { .. } => { Table::BTree(_) => {
if iter_dir if iter_dir
.as_ref() .as_ref()
.is_some_and(|dir| *dir == IterationDirection::Backwards) .is_some_and(|dir| *dir == IterationDirection::Backwards)
@@ -750,7 +751,7 @@ pub fn close_loop(
}); });
} }
} }
TableReferenceType::VirtualTable { .. } => { Table::Virtual(_) => {
program.emit_insn(Insn::VNext { program.emit_insn(Insn::VNext {
cursor_id, cursor_id,
pc_if_next: loop_labels.loop_start, pc_if_next: loop_labels.loop_start,

View File

@@ -198,7 +198,6 @@ pub struct TableReference {
pub identifier: String, pub identifier: String,
/// The join info for this table reference, if it is the right side of a join (which all except the first table reference have) /// The join info for this table reference, if it is the right side of a join (which all except the first table reference have)
pub join_info: Option<JoinInfo>, pub join_info: Option<JoinInfo>,
pub reference_type: TableReferenceType,
} }
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
@@ -225,35 +224,17 @@ pub enum Operation {
}, },
} }
/// The type of the table reference, either BTreeTable or Subquery
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum TableReferenceType {
/// A BTreeTable is a table that is stored on disk in a B-tree index.
BTreeTable,
/// A subquery.
Subquery {
/// The index of the first register in the query plan that contains the result columns of the subquery.
result_columns_start_reg: usize,
},
/// A virtual table.
VirtualTable {
/// Arguments to pass e.g. generate_series(1, 10, 2)
args: Vec<ast::Expr>,
},
}
impl TableReference { impl TableReference {
/// Returns the btree table for this table reference, if it is a BTreeTable. /// Returns the btree table for this table reference, if it is a BTreeTable.
pub fn btree(&self) -> Option<Rc<BTreeTable>> { pub fn btree(&self) -> Option<Rc<BTreeTable>> {
match &self.reference_type { match &self.table {
TableReferenceType::BTreeTable => self.table.btree(), Table::BTree(_) => self.table.btree(),
TableReferenceType::Subquery { .. } => None, _ => None,
TableReferenceType::VirtualTable { .. } => None,
} }
} }
pub fn virtual_table(&self) -> Option<Rc<VirtualTable>> { pub fn virtual_table(&self) -> Option<Rc<VirtualTable>> {
match &self.reference_type { match &self.table {
TableReferenceType::VirtualTable { .. } => self.table.virtual_table(), Table::Virtual(_) => self.table.virtual_table(),
_ => None, _ => None,
} }
} }
@@ -280,9 +261,6 @@ impl TableReference {
result_columns_start_reg: 0, // Will be set in the bytecode emission phase result_columns_start_reg: 0, // Will be set in the bytecode emission phase
}, },
table, table,
reference_type: TableReferenceType::Subquery {
result_columns_start_reg: 0, // Will be set in the bytecode emission phase
},
identifier: identifier.clone(), identifier: identifier.clone(),
join_info, join_info,
} }

View File

@@ -1,7 +1,7 @@
use super::{ use super::{
plan::{ plan::{
Aggregate, JoinInfo, Operation, Plan, ResultSetColumn, SelectQueryType, TableReference, Aggregate, JoinInfo, Operation, Plan, ResultSetColumn, SelectQueryType, TableReference,
TableReferenceType, WhereTerm, WhereTerm,
}, },
select::prepare_select_plan, select::prepare_select_plan,
SymbolTable, SymbolTable,
@@ -11,7 +11,7 @@ use crate::{
schema::{Schema, Table}, schema::{Schema, Table},
util::{exprs_are_equivalent, normalize_ident}, util::{exprs_are_equivalent, normalize_ident},
vdbe::BranchOffset, vdbe::BranchOffset,
Result, Result, VirtualTable,
}; };
use sqlite3_parser::ast::{self, Expr, FromClause, JoinType, Limit, UnaryOperator}; use sqlite3_parser::ast::{self, Expr, FromClause, JoinType, Limit, UnaryOperator};
@@ -301,7 +301,6 @@ fn parse_from_clause_table(
table: Table::BTree(table.clone()), table: Table::BTree(table.clone()),
identifier: alias.unwrap_or(normalized_qualified_name), identifier: alias.unwrap_or(normalized_qualified_name),
join_info: None, join_info: None,
reference_type: TableReferenceType::BTreeTable,
}) })
} }
ast::SelectTable::Select(subselect, maybe_alias) => { ast::SelectTable::Select(subselect, maybe_alias) => {
@@ -320,9 +319,9 @@ fn parse_from_clause_table(
.unwrap_or(format!("subquery_{}", cur_table_index)); .unwrap_or(format!("subquery_{}", cur_table_index));
Ok(TableReference::new_subquery(identifier, subplan, None)) Ok(TableReference::new_subquery(identifier, subplan, None))
} }
ast::SelectTable::TableCall(qualified_name, mut maybe_args, maybe_alias) => { ast::SelectTable::TableCall(qualified_name, maybe_args, maybe_alias) => {
let normalized_name = normalize_ident(qualified_name.name.0.as_str()); let normalized_name = &normalize_ident(qualified_name.name.0.as_str());
let Some(vtab) = syms.vtabs.get(&normalized_name) else { let Some(vtab) = syms.vtabs.get(normalized_name) else {
crate::bail_parse_error!("Virtual table {} not found", normalized_name); crate::bail_parse_error!("Virtual table {} not found", normalized_name);
}; };
let alias = maybe_alias let alias = maybe_alias
@@ -331,16 +330,22 @@ fn parse_from_clause_table(
ast::As::As(id) => id.0.clone(), ast::As::As(id) => id.0.clone(),
ast::As::Elided(id) => id.0.clone(), ast::As::Elided(id) => id.0.clone(),
}) })
.unwrap_or(normalized_name); .unwrap_or(normalized_name.to_string());
Ok(TableReference { Ok(TableReference {
op: Operation::Scan { iter_dir: None }, op: Operation::Scan { iter_dir: None },
join_info: None, join_info: None,
table: Table::Virtual(vtab.clone().into()), table: Table::Virtual(
VirtualTable {
name: normalized_name.clone(),
args: maybe_args,
implementation: vtab.implementation.clone(),
columns: vtab.columns.clone(),
}
.into(),
)
.into(),
identifier: alias.clone(), identifier: alias.clone(),
reference_type: TableReferenceType::VirtualTable {
args: maybe_args.take().unwrap_or_default(),
},
}) })
} }
_ => todo!(), _ => todo!(),

View File

@@ -307,7 +307,7 @@ def test_crypto(pipe):
run_test( run_test(
pipe, pipe,
"SELECT crypto_blake('a');", "SELECT crypto_blake('a');",
lambda res: "Error" in res, lambda res: "Parse error" in res,
"crypto_blake3 returns null when ext not loaded", "crypto_blake3 returns null when ext not loaded",
) )
run_test( run_test(