mirrors/limbo
1
0
Fork 0
mirror of https://github.com/tursodatabase/limbo.git synced 2025-07-07 12:35:00 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 9

index scan wip foo doesnt work yet

Browse source
This commit is contained in:
jussisaurio 2024-09-29 19:13:50 +03:00
parent d3015ad854
commit f02da18acd
14 changed files with 792 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

8
core/pseudo.rs
View file

@ -52,6 +52,14 @@ impl Cursor for PseudoCursor {
unimplemented!();
}
fn seek_ge(&mut self, _: &OwnedRecord) -> Result<CursorResult<bool>> {
unimplemented!();
}
fn seek_gt(&mut self, _: &OwnedRecord) -> Result<CursorResult<bool>> {
unimplemented!();
}
fn seek_to_last(&mut self) -> Result<CursorResult<()>> {
unimplemented!();
}

14
core/schema.rs
View file

@ -46,6 +46,7 @@ impl Schema {
#[derive(Clone, Debug)]
pub enum Table {
BTree(Rc<BTreeTable>),
Index(Rc<Index>),
Pseudo(Rc<PseudoTable>),
}
@ -57,6 +58,7 @@ impl Table {
pub fn get_rowid_alias_column(&self) -> Option<(usize, &Column)> {
match self {
Table::BTree(table) => table.get_rowid_alias_column(),
Table::Index(_) => None,
Table::Pseudo(_) => None,
}
}
@ -64,6 +66,7 @@ impl Table {
pub fn column_is_rowid_alias(&self, col: &Column) -> bool {
match self {
Table::BTree(table) => table.column_is_rowid_alias(col),
Table::Index(_) => false,
Table::Pseudo(_) => false,
}
}
@ -71,6 +74,7 @@ impl Table {
pub fn get_name(&self) -> &str {
match self {
Table::BTree(table) => &table.name,
Table::Index(index) => &index.name,
Table::Pseudo(_) => "",
}
}
@ -81,6 +85,10 @@ impl Table {
Some(column) => Some(&column.name),
None => None,
},
Table::Index(i) => match i.columns.get(index) {
Some(column) => Some(&column.name),
None => None,
},
Table::Pseudo(table) => match table.columns.get(index) {
Some(column) => Some(&column.name),
None => None,
@ -91,6 +99,7 @@ impl Table {
pub fn get_column(&self, name: &str) -> Option<(usize, &Column)> {
match self {
Table::BTree(table) => table.get_column(name),
Table::Index(index) => unimplemented!(),
Table::Pseudo(table) => table.get_column(name),
}
}
@ -98,6 +107,7 @@ impl Table {
pub fn get_column_at(&self, index: usize) -> &Column {
match self {
Table::BTree(table) => table.columns.get(index).unwrap(),
Table::Index(index) => unimplemented!(),
Table::Pseudo(table) => table.columns.get(index).unwrap(),
}
}
@ -105,6 +115,7 @@ impl Table {
pub fn columns(&self) -> &Vec<Column> {
match self {
Table::BTree(table) => &table.columns,
Table::Index(index) => unimplemented!(),
Table::Pseudo(table) => &table.columns,
}
}
@ -112,7 +123,8 @@ impl Table {
pub fn has_rowid(&self) -> bool {
match self {
Table::BTree(table) => table.has_rowid,
Table::Pseudo(_) => todo!(),
Table::Index(_) => unimplemented!(),
Table::Pseudo(_) => unimplemented!(),
}
}
}

176
core/storage/btree.rs
View file

@ -11,7 +11,7 @@ use crate::Result;
use std::cell::{Ref, RefCell};
use std::rc::Rc;
use super::sqlite3_ondisk::{write_varint_to_vec, OverflowCell};
use super::sqlite3_ondisk::{write_varint_to_vec, IndexInteriorCell, IndexLeafCell, OverflowCell};
/*
These are offsets of fields in the header of a b-tree page.
@ -23,6 +23,12 @@ const BTREE_HEADER_OFFSET_CELL_CONTENT: usize = 5; /* pointer to first byte of c
const BTREE_HEADER_OFFSET_FRAGMENTED: usize = 7; /* number of fragmented bytes -> u8 */
const BTREE_HEADER_OFFSET_RIGHTMOST: usize = 8; /* if internalnode, pointer right most pointer (saved separately from cells) -> u32 */
#[derive(Clone)]
pub enum IndexSeekOp {
GT,
GE,
}
pub struct MemPage {
parent: Option<Rc<MemPage>>,
page_idx: usize,
@ -145,16 +151,77 @@ impl BTreeCursor {
let record = crate::storage::sqlite3_ondisk::read_record(_payload)?;
return Ok(CursorResult::Ok((Some(*_rowid), Some(record))));
}
BTreeCell::IndexInteriorCell(_) => {
unimplemented!();
BTreeCell::IndexInteriorCell(IndexInteriorCell {
left_child_page, ..
}) => {
mem_page.advance();
let mem_page =
MemPage::new(Some(mem_page.clone()), *left_child_page as usize, 0);
self.page.replace(Some(Rc::new(mem_page)));
continue;
}
BTreeCell::IndexLeafCell(_) => {
unimplemented!();
BTreeCell::IndexLeafCell(IndexLeafCell { payload, .. }) => {
mem_page.advance();
let record = crate::storage::sqlite3_ondisk::read_record(payload)?;
let rowid = match record.values[1] {
OwnedValue::Integer(rowid) => rowid as u64,
_ => unreachable!("index cells should have an integer rowid"),
};
return Ok(CursorResult::Ok((Some(rowid), Some(record))));
}
}
}
}
fn btree_index_seek(
&mut self,
key: &OwnedRecord,
op: IndexSeekOp,
) -> Result<CursorResult<(Option<u64>, Option<OwnedRecord>)>> {
self.move_to_index_leaf(key, op.clone())?;
let mem_page = self.get_mem_page();
let page_idx = mem_page.page_idx;
let page = self.pager.read_page(page_idx)?;
let page = RefCell::borrow(&page);
if page.is_locked() {
return Ok(CursorResult::IO);
}
let page = page.contents.read().unwrap();
let page = page.as_ref().unwrap();
for cell_idx in 0..page.cell_count() {
match &page.cell_get(
cell_idx,
self.pager.clone(),
self.max_local(page.page_type()),
self.min_local(page.page_type()),
self.usable_space(),
)? {
BTreeCell::IndexLeafCell(IndexLeafCell { payload, .. }) => {
mem_page.advance();
let record = crate::storage::sqlite3_ondisk::read_record(payload)?;
let comparison = match op {
IndexSeekOp::GT => record > *key,
IndexSeekOp::GE => record >= *key,
};
if comparison {
let rowid = match record.values.get(1) {
Some(OwnedValue::Integer(rowid)) => *rowid as u64,
_ => unreachable!("index cells should have an integer rowid"),
};
return Ok(CursorResult::Ok((Some(rowid), Some(record))));
}
}
cell_type => {
unreachable!("unexpected cell type: {:?}", cell_type);
}
}
}
Ok(CursorResult::Ok((None, None)))
}
fn btree_seek_rowid(
&mut self,
rowid: u64,
@ -337,6 +404,81 @@ impl BTreeCursor {
}
}
fn move_to_index_leaf(
&mut self,
key: &OwnedRecord,
cmp: IndexSeekOp,
) -> Result<CursorResult<()>> {
self.move_to_root();
loop {
let mem_page = self.get_mem_page();
let page_idx = mem_page.page_idx;
let page = self.pager.read_page(page_idx)?;
let page = RefCell::borrow(&page);
if page.is_locked() {
return Ok(CursorResult::IO);
}
let page = page.contents.read().unwrap();
let page = page.as_ref().unwrap();
if page.is_leaf() {
return Ok(CursorResult::Ok(()));
}
let mut found_cell = false;
for cell_idx in 0..page.cell_count() {
match &page.cell_get(
cell_idx,
self.pager.clone(),
self.max_local(page.page_type()),
self.min_local(page.page_type()),
self.usable_space(),
)? {
BTreeCell::IndexInteriorCell(IndexInteriorCell {
left_child_page,
payload,
..
}) => {
// get the logic for this from btree_index_seek
let record = crate::storage::sqlite3_ondisk::read_record(payload)?;
let comparison = match cmp {
IndexSeekOp::GT => record > *key,
IndexSeekOp::GE => record >= *key,
};
if comparison {
mem_page.advance();
let mem_page =
MemPage::new(Some(mem_page.clone()), *left_child_page as usize, 0);
self.page.replace(Some(Rc::new(mem_page)));
found_cell = true;
break;
}
}
_ => {
unreachable!(
"we don't iterate leaf cells while trying to move to a leaf cell"
);
}
}
}
if !found_cell {
let parent = mem_page.clone();
match page.rightmost_pointer() {
Some(right_most_pointer) => {
let mem_page = MemPage::new(Some(parent), right_most_pointer as usize, 0);
self.page.replace(Some(Rc::new(mem_page)));
continue;
}
None => {
unreachable!("we shall not go back up! The only way is down the slope");
}
}
}
}
}
fn insert_to_page(
&mut self,
key: &OwnedValue,
@ -1296,6 +1438,30 @@ impl Cursor for BTreeCursor {
}
}
fn seek_ge(&mut self, key: &OwnedRecord) -> Result<CursorResult<bool>> {
match self.btree_index_seek(key, IndexSeekOp::GE)? {
CursorResult::Ok((rowid, record)) => {
self.rowid.replace(rowid);
println!("seek_ge: {:?}", record);
self.record.replace(record);
Ok(CursorResult::Ok(rowid.is_some()))
}
CursorResult::IO => Ok(CursorResult::IO),
}
}
fn seek_gt(&mut self, key: &OwnedRecord) -> Result<CursorResult<bool>> {
match self.btree_index_seek(key, IndexSeekOp::GT)? {
CursorResult::Ok((rowid, record)) => {
self.rowid.replace(rowid);
println!("seek_gt: {:?}", record);
self.record.replace(record);
Ok(CursorResult::Ok(rowid.is_some()))
}
CursorResult::IO => Ok(CursorResult::IO),
}
}
fn record(&self) -> Result<Ref<Option<OwnedRecord>>> {
Ok(self.record.borrow())
}

199
core/translate/emitter.rs
View file

@ -2,6 +2,8 @@ use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;
use sqlite3_parser::ast;
use crate::schema::{BTreeTable, Column, PseudoTable, Table};
use crate::storage::sqlite3_ondisk::DatabaseHeader;
use crate::translate::expr::resolve_ident_pseudo_table;
@ -251,6 +253,165 @@ impl Emitter for Operator {
_ => Ok(OpStepResult::Done),
}
}
Operator::IndexScan {
table,
table_identifier,
index,
seek_cmp,
seek_expr,
predicates,
step,
id,
..
} => {
*step += 1;
const INDEX_SCAN_OPEN_AND_SEEK: usize = 1;
const INDEX_SCAN_NEXT: usize = 2;
match *step {
INDEX_SCAN_OPEN_AND_SEEK => {
let table_cursor_id = program.alloc_cursor_id(
Some(table_identifier.clone()),
Some(Table::BTree(table.clone())),
);
let index_cursor_id = program.alloc_cursor_id(
Some(index.name.clone()),
Some(Table::Index(index.clone())),
);
let next_row_label = program.allocate_label();
m.next_row_labels.insert(*id, next_row_label);
let rewind_label = program.allocate_label();
m.rewind_labels.push(rewind_label);
program.emit_insn(Insn::OpenReadAsync {
cursor_id: table_cursor_id,
root_page: table.root_page,
});
program.emit_insn(Insn::OpenReadAwait);
program.emit_insn(Insn::OpenReadAsync {
cursor_id: index_cursor_id,
root_page: index.root_page,
});
program.emit_insn(Insn::OpenReadAwait);
let cmp_reg = program.alloc_register();
// TODO this only handles ascending indexes
match seek_cmp {
ast::Operator::Equals
| ast::Operator::Greater
| ast::Operator::GreaterEquals => {
translate_expr(
program,
Some(referenced_tables),
seek_expr,
cmp_reg,
None,
None,
)?;
}
ast::Operator::Less | ast::Operator::LessEquals => {
program.emit_insn(Insn::Null {
dest: cmp_reg,
dest_end: None,
});
}
_ => unreachable!(),
}
program.emit_insn_with_label_dependency(
match seek_cmp {
ast::Operator::Equals | ast::Operator::GreaterEquals => {
Insn::SeekGE {
cursor_id: index_cursor_id,
start_reg: cmp_reg,
num_regs: 1,
target_pc: *m.termination_label_stack.last().unwrap(),
}
}
ast::Operator::Greater
| ast::Operator::Less
| ast::Operator::LessEquals => Insn::SeekGT {
cursor_id: index_cursor_id,
start_reg: cmp_reg,
num_regs: 1,
target_pc: *m.termination_label_stack.last().unwrap(),
},
_ => unreachable!(),
},
*m.termination_label_stack.last().unwrap(),
);
if *seek_cmp == ast::Operator::Less
|| *seek_cmp == ast::Operator::LessEquals
{
translate_expr(
program,
Some(referenced_tables),
seek_expr,
cmp_reg,
None,
None,
)?;
}
program.defer_label_resolution(rewind_label, program.offset() as usize);
// We are currently only handling ascending indexes.
// For conditions like index_key > 10, we have already seeked to the first key greater than 10, and can just scan forward.
// For conditions like index_key < 10, we are at the beginning of the index, and will scan forward and emit IdxGE(10) with a conditional jump to the end.
// For conditions like index_key = 10, we have already seeked to the first key greater than or equal to 10, and can just scan forward and emit IdxGT(10) with a conditional jump to the end.
// For conditions like index_key >= 10, we have already seeked to the first key greater than or equal to 10, and can just scan forward.
// For conditions like index_key <= 10, we are at the beginning of the index, and will scan forward and emit IdxGT(10) with a conditional jump to the end.
// For conditions like index_key != 10, TODO. probably the optimal way is not to use an index at all.
let abort_jump_target = *m.termination_label_stack.last().unwrap();
match seek_cmp {
ast::Operator::Equals | ast::Operator::LessEquals => {
program.emit_insn_with_label_dependency(
Insn::IdxGT {
cursor_id: index_cursor_id,
start_reg: cmp_reg,
num_regs: 1,
target_pc: abort_jump_target,
},
abort_jump_target,
);
}
ast::Operator::Less => {
program.emit_insn_with_label_dependency(
Insn::IdxGE {
cursor_id: index_cursor_id,
start_reg: cmp_reg,
num_regs: 1,
target_pc: abort_jump_target,
},
abort_jump_target,
);
}
_ => {}
}
program.emit_insn(Insn::DeferredSeek {
index_cursor_id,
table_cursor_id,
});
Ok(OpStepResult::ReadyToEmit)
}
INDEX_SCAN_NEXT => {
let cursor_id = program.resolve_cursor_id(&index.name, None);
program
.resolve_label(*m.next_row_labels.get(id).unwrap(), program.offset());
program.emit_insn(Insn::NextAsync { cursor_id });
let jump_label = m.rewind_labels.pop().unwrap();
program.emit_insn_with_label_dependency(
Insn::NextAwait {
cursor_id,
pc_if_next: jump_label,
},
jump_label,
);
Ok(OpStepResult::Done)
}
_ => Ok(OpStepResult::Done),
}
}
Operator::SeekRowid {
table,
table_identifier,
@ -1196,6 +1357,23 @@ impl Emitter for Operator {
Ok(start_reg)
}
Operator::IndexScan {
table,
table_identifier,
..
} => {
let start_reg = program.alloc_registers(col_count);
let table = cursor_override
.map(|c| c.pseudo_table.clone())
.unwrap_or_else(|| Table::BTree(table.clone()));
let cursor_id = cursor_override
.map(|c| c.cursor_id)
.unwrap_or_else(|| program.resolve_cursor_id(table_identifier, None));
let start_column_offset = cursor_override.map(|c| c.sort_key_len).unwrap_or(0);
translate_table_columns(program, cursor_id, &table, start_column_offset, start_reg);
Ok(start_reg)
}
Operator::Join { left, right, .. } => {
let left_start_reg =
left.result_columns(program, referenced_tables, m, cursor_override)?;
@ -1415,13 +1593,7 @@ impl Emitter for Operator {
fn prologue(
cache: ExpressionResultCache,
) -> Result<(
ProgramBuilder,
Metadata,
BranchOffset,
BranchOffset,
BranchOffset,
)> {
) -> Result<(ProgramBuilder, Metadata, BranchOffset, BranchOffset)> {
let mut program = ProgramBuilder::new();
let init_label = program.allocate_label();
let halt_label = program.allocate_label();
@ -1446,16 +1618,19 @@ fn prologue(
sorts: HashMap::new(),
};
Ok((program, metadata, init_label, halt_label, start_offset))
Ok((program, metadata, init_label, start_offset))
}
fn epilogue(
program: &mut ProgramBuilder,
metadata: &mut Metadata,
init_label: BranchOffset,
halt_label: BranchOffset,
start_offset: BranchOffset,
) -> Result<()> {
program.resolve_label(halt_label, program.offset());
program.resolve_label(
metadata.termination_label_stack.pop().unwrap(),
program.offset(),
);
program.emit_insn(Insn::Halt {
err_code: 0,
description: String::new(),
@ -1479,7 +1654,7 @@ pub fn emit_program(
mut plan: Plan,
cache: ExpressionResultCache,
) -> Result<Program> {
let (mut program, mut metadata, init_label, halt_label, start_offset) = prologue(cache)?;
let (mut program, mut metadata, init_label, start_offset) = prologue(cache)?;
loop {
match plan
.root_operator
@ -1495,7 +1670,7 @@ pub fn emit_program(
)?;
}
OpStepResult::Done => {
epilogue(&mut program, init_label, halt_label, start_offset)?;
epilogue(&mut program, &mut metadata, init_label, start_offset)?;
return Ok(program.build(database_header));
}
}

1
core/translate/insert.rs
View file

@ -50,6 +50,7 @@ pub fn translate_insert(
);
let root_page = match table.as_ref() {
Table::BTree(btree) => btree.root_page,
Table::Index(index) => index.root_page,
Table::Pseudo(_) => todo!(),
};

116
core/translate/optimizer.rs
View file

@ -2,7 +2,12 @@ use std::{collections::HashMap, rc::Rc};
use sqlite3_parser::ast;
use crate::{schema::BTreeTable, util::normalize_ident, Result};
use crate::{
schema::{BTreeTable, Index},
types::OwnedValue,
util::normalize_ident,
Result,
};
use super::plan::{
get_table_ref_bitmask_for_ast_expr, get_table_ref_bitmask_for_operator, Operator, Plan,
@ -25,6 +30,7 @@ pub fn optimize_plan(mut select_plan: Plan) -> Result<(Plan, ExpressionResultCac
Plan {
root_operator: Operator::Nothing,
referenced_tables: vec![],
available_indexes: vec![],
},
expr_result_cache,
));
@ -32,6 +38,7 @@ pub fn optimize_plan(mut select_plan: Plan) -> Result<(Plan, ExpressionResultCac
use_indexes(
&mut select_plan.root_operator,
&select_plan.referenced_tables,
&select_plan.available_indexes,
)?;
find_shared_expressions_in_child_operators_and_mark_them_so_that_the_parent_operator_doesnt_recompute_them(&select_plan.root_operator, &mut expr_result_cache);
Ok((select_plan, expr_result_cache))
@ -43,8 +50,10 @@ pub fn optimize_plan(mut select_plan: Plan) -> Result<(Plan, ExpressionResultCac
fn use_indexes(
operator: &mut Operator,
referenced_tables: &[(Rc<BTreeTable>, String)],
available_indexes: &[Rc<Index>],
) -> Result<()> {
match operator {
Operator::IndexScan { .. } => Ok(()),
Operator::Scan {
table,
predicates: filter,
@ -90,35 +99,72 @@ fn use_indexes(
predicates: predicates_owned,
id: *id,
step: 0,
};
return Ok(());
}
let mut maybe_index_predicate = None;
let mut maybe_index_idx = None;
let fs = filter.as_mut().unwrap();
for i in 0..fs.len() {
let mut f = fs[i].take_ownership();
let index_idx = f.check_index_scan(available_indexes)?;
if index_idx.is_some() {
maybe_index_predicate = Some(f);
maybe_index_idx = index_idx;
fs.remove(i);
break;
}
}
if let Some(index_idx) = maybe_index_idx {
let index_predicate = maybe_index_predicate.unwrap();
match index_predicate {
ast::Expr::Binary(lhs, op, rhs) => {
*operator = Operator::IndexScan {
table: table.clone(),
index: available_indexes[index_idx].clone(),
index_predicate: ast::Expr::Binary(lhs, op, rhs.clone()),
predicates: Some(std::mem::take(fs)),
seek_cmp: op,
seek_expr: *rhs,
table_identifier: table_identifier.clone(),
id: *id,
step: 0,
};
}
_ => {
crate::bail_parse_error!("Unsupported index predicate");
}
}
}
Ok(())
}
Operator::Aggregate { source, .. } => {
use_indexes(source, referenced_tables)?;
use_indexes(source, referenced_tables, available_indexes)?;
Ok(())
}
Operator::Filter { source, .. } => {
use_indexes(source, referenced_tables)?;
use_indexes(source, referenced_tables, available_indexes)?;
Ok(())
}
Operator::SeekRowid { .. } => Ok(()),
Operator::Limit { source, .. } => {
use_indexes(source, referenced_tables)?;
use_indexes(source, referenced_tables, available_indexes)?;
Ok(())
}
Operator::Join { left, right, .. } => {
use_indexes(left, referenced_tables)?;
use_indexes(right, referenced_tables)?;
use_indexes(left, referenced_tables, available_indexes)?;
use_indexes(right, referenced_tables, available_indexes)?;
Ok(())
}
Operator::Order { source, .. } => {
use_indexes(source, referenced_tables)?;
use_indexes(source, referenced_tables, available_indexes)?;
Ok(())
}
Operator::Projection { source, .. } => {
use_indexes(source, referenced_tables)?;
use_indexes(source, referenced_tables, available_indexes)?;
Ok(())
}
Operator::Nothing => Ok(()),
@ -279,6 +325,7 @@ fn eliminate_constants(operator: &mut Operator) -> Result<ConstantConditionElimi
}
Ok(ConstantConditionEliminationResult::Continue)
}
Operator::IndexScan { .. } => Ok(ConstantConditionEliminationResult::Continue),
Operator::Nothing => Ok(ConstantConditionEliminationResult::Continue),
}
}
@ -379,6 +426,7 @@ fn push_predicates(
Ok(())
}
Operator::Scan { .. } => Ok(()),
Operator::IndexScan { .. } => Ok(()),
Operator::Nothing => Ok(()),
}
}
@ -424,6 +472,7 @@ fn push_predicate(
Ok(None)
}
Operator::IndexScan { .. } => Ok(Some(predicate)),
Operator::Filter {
source,
predicates: ps,
@ -684,6 +733,7 @@ fn find_indexes_of_all_result_columns_in_operator_that_match_expr_either_fully_o
mask
}
Operator::Scan { .. } => 0,
Operator::IndexScan { .. } => 0,
Operator::Nothing => 0,
};
@ -883,6 +933,7 @@ fn find_shared_expressions_in_child_operators_and_mark_them_so_that_the_parent_o
find_shared_expressions_in_child_operators_and_mark_them_so_that_the_parent_operator_doesnt_recompute_them(source, expr_result_cache)
}
Operator::Scan { .. } => {}
Operator::IndexScan { .. } => {}
Operator::Nothing => {}
}
}
@ -915,6 +966,7 @@ pub trait Optimizable {
&self,
referenced_tables: &[(Rc<BTreeTable>, String)],
) -> Result<Option<usize>>;
fn check_index_scan(&mut self, available_indexes: &[Rc<Index>]) -> Result<Option<usize>>;
}
impl Optimizable for ast::Expr {
@ -967,6 +1019,54 @@ impl Optimizable for ast::Expr {
_ => Ok(None),
}
}
fn check_index_scan(&mut self, available_indexes: &[Rc<Index>]) -> Result<Option<usize>> {
match self {
ast::Expr::Id(ident) => {
let ident = normalize_ident(&ident.0);
let indexes = available_indexes
.iter()
.enumerate()
.filter(|(_, i)| i.columns.iter().any(|c| c.name == ident))
.collect::<Vec<_>>();
if indexes.is_empty() {
return Ok(None);
}
if indexes.len() > 1 {
crate::bail_parse_error!("ambiguous column name {}", ident)
}
Ok(Some(indexes.first().unwrap().0))
}
ast::Expr::Qualified(tbl, ident) => {
let tbl = normalize_ident(&tbl.0);
let ident = normalize_ident(&ident.0);
let index = available_indexes.iter().enumerate().find(|(_, i)| {
let normalized_tbl = normalize_ident(&i.table_name);
normalized_tbl == tbl
&& i.columns.iter().any(|c| normalize_ident(&c.name) == ident)
});
if index.is_none() {
return Ok(None);
}
Ok(Some(index.unwrap().0))
}
ast::Expr::Binary(lhs, op, rhs) => {
let lhs_index = lhs.check_index_scan(available_indexes)?;
if lhs_index.is_some() {
return Ok(lhs_index);
}
let rhs_index = rhs.check_index_scan(available_indexes)?;
if rhs_index.is_some() {
// swap lhs and rhs
let lhs_new = rhs.take_ownership();
let rhs_new = lhs.take_ownership();
*self = ast::Expr::Binary(Box::new(lhs_new), *op, Box::new(rhs_new));
return Ok(rhs_index);
}
Ok(None)
}
_ => Ok(None),
}
}
fn check_constant(&self) -> Result<Option<ConstantPredicate>> {
match self {
ast::Expr::Literal(lit) => match lit {

35
core/translate/plan.rs
View file

@ -6,12 +6,18 @@ use std::{
use sqlite3_parser::ast;
use crate::{function::AggFunc, schema::BTreeTable, util::normalize_ident, Result};
use crate::{
function::AggFunc,
schema::{BTreeTable, Index},
util::normalize_ident,
Result,
};
#[derive(Debug)]
pub struct Plan {
pub root_operator: Operator,
pub referenced_tables: Vec<(Rc<BTreeTable>, String)>,
pub available_indexes: Vec<Rc<Index>>,
}
impl Display for Plan {
@ -123,6 +129,17 @@ pub enum Operator {
predicates: Option<Vec<ast::Expr>>,
step: usize,
},
IndexScan {
id: usize,
index: Rc<Index>,
seek_cmp: ast::Operator,
seek_expr: ast::Expr,
index_predicate: ast::Expr,
table: Rc<BTreeTable>,
table_identifier: String,
predicates: Option<Vec<ast::Expr>>,
step: usize,
},
// Nothing operator
// This operator is used to represent an empty query.
// e.g. SELECT * from foo WHERE 0 will eventually be optimized to Nothing.
@ -172,6 +189,7 @@ impl Operator {
.map(|e| e.column_count(referenced_tables))
.sum(),
Operator::Scan { table, .. } => table.columns.len(),
Operator::IndexScan { table, .. } => table.columns.len(),
Operator::Nothing => 0,
}
}
@ -226,6 +244,9 @@ impl Operator {
})
.collect(),
Operator::Scan { table, .. } => table.columns.iter().map(|c| c.name.clone()).collect(),
Operator::IndexScan { table, .. } => {
table.columns.iter().map(|c| c.name.clone()).collect()
}
Operator::Nothing => vec![],
}
}
@ -240,6 +261,7 @@ impl Operator {
Operator::Order { id, .. } => *id,
Operator::Projection { id, .. } => *id,
Operator::Scan { id, .. } => *id,
Operator::IndexScan { id, .. } => *id,
Operator::Nothing => unreachable!(),
}
}
@ -429,6 +451,10 @@ impl Display for Operator {
}?;
Ok(())
}
Operator::IndexScan { table, .. } => {
writeln!(f, "{}INDEX SCAN {}", indent, table.name)?;
Ok(())
}
Operator::Nothing => Ok(()),
}
}
@ -489,6 +515,13 @@ pub fn get_table_ref_bitmask_for_operator<'a>(
.position(|(t, _)| Rc::ptr_eq(t, table))
.unwrap();
}
Operator::IndexScan { table, .. } => {
table_refs_mask |= 1
<< tables
.iter()
.position(|(t, _)| Rc::ptr_eq(t, table))
.unwrap();
}
Operator::Nothing => {}
}
Ok(table_refs_mask)

9
core/translate/planner.rs
View file

@ -277,10 +277,19 @@ pub fn prepare_select_plan<'a>(schema: &Schema, select: ast::Select) -> Result<P
}
}
println!("{:?}", schema.indexes);
// Return the unoptimized query plan
Ok(Plan {
root_operator: operator,
referenced_tables,
available_indexes: schema
.indexes
.clone()
.into_iter()
.map(|(_, v)| v)
.flatten()
.collect(),
})
}
_ => todo!(),

2
core/types.rs
View file

@ -416,6 +416,8 @@ pub trait Cursor {
fn wait_for_completion(&mut self) -> Result<()>;
fn rowid(&self) -> Result<Option<u64>>;
fn seek_rowid(&mut self, rowid: u64) -> Result<CursorResult<bool>>;
fn seek_ge(&mut self, key: &OwnedRecord) -> Result<CursorResult<bool>>;
fn seek_gt(&mut self, key: &OwnedRecord) -> Result<CursorResult<bool>>;
fn seek_to_last(&mut self) -> Result<CursorResult<()>>;
fn record(&self) -> Result<Ref<Option<OwnedRecord>>>;
fn insert(

16
core/vdbe/builder.rs
View file

@ -299,6 +299,22 @@ impl ProgramBuilder {
assert!(*target_pc_eq < 0);
*target_pc_eq = to_offset;
}
Insn::SeekGE { target_pc, .. } => {
assert!(*target_pc < 0);
*target_pc = to_offset;
}
Insn::SeekGT { target_pc, .. } => {
assert!(*target_pc < 0);
*target_pc = to_offset;
}
Insn::IdxGE { target_pc, .. } => {
assert!(*target_pc < 0);
*target_pc = to_offset;
}
Insn::IdxGT { target_pc, .. } => {
assert!(*target_pc < 0);
*target_pc = to_offset;
}
_ => {
todo!("missing resolve_label for {:?}", insn);
}

68
core/vdbe/explain.rs
View file

@ -521,6 +521,74 @@ pub fn insn_to_str(
target_pc
),
),
Insn::DeferredSeek {
index_cursor_id,
table_cursor_id,
} => (
"DeferredSeek",
*index_cursor_id as i32,
*table_cursor_id as i32,
0,
OwnedValue::Text(Rc::new("".to_string())),
0,
"".to_string(),
),
Insn::SeekGT {
cursor_id,
start_reg,
num_regs,
target_pc,
} => (
"SeekGT",
*cursor_id as i32,
*target_pc as i32,
*start_reg as i32,
OwnedValue::Text(Rc::new("".to_string())),
0,
"".to_string(),
),
Insn::SeekGE {
cursor_id,
start_reg,
num_regs,
target_pc,
} => (
"SeekGE",
*cursor_id as i32,
*target_pc as i32,
*start_reg as i32,
OwnedValue::Text(Rc::new("".to_string())),
0,
"".to_string(),
),
Insn::IdxGT {
cursor_id,
start_reg,
num_regs,
target_pc,
} => (
"IdxGT",
*cursor_id as i32,
*target_pc as i32,
*start_reg as i32,
OwnedValue::Text(Rc::new("".to_string())),
0,
"".to_string(),
),
Insn::IdxGE {
cursor_id,
start_reg,
num_regs,
target_pc,
} => (
"IdxGE",
*cursor_id as i32,
*target_pc as i32,
*start_reg as i32,
OwnedValue::Text(Rc::new("".to_string())),
0,
"".to_string(),
),
Insn::DecrJumpZero { reg, target_pc } => (
"DecrJumpZero",
*reg as i32,

168
core/vdbe/mod.rs
View file

@ -298,6 +298,47 @@ pub enum Insn {
target_pc: BranchOffset,
},
// 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.
DeferredSeek {
index_cursor_id: CursorID,
table_cursor_id: CursorID,
},
// Seek to the first index entry that is greater than or equal to the given key. If not found, jump to the given PC. Otherwise, continue to the next instruction.
SeekGE {
cursor_id: CursorID,
start_reg: usize,
num_regs: usize,
target_pc: BranchOffset,
},
// Seek to the first index entry that is greater than the given key. If not found, jump to the given PC. Otherwise, continue to the next instruction.
SeekGT {
cursor_id: CursorID,
start_reg: usize,
num_regs: usize,
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.
IdxGE {
cursor_id: CursorID,
start_reg: usize,
num_regs: usize,
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 than the key value then jump to P2. Otherwise fall through to the next instruction.
IdxGT {
cursor_id: CursorID,
start_reg: usize,
num_regs: usize,
target_pc: BranchOffset,
},
// Decrement the given register and jump to the given PC if the result is zero.
DecrJumpZero {
reg: usize,
@ -444,6 +485,7 @@ pub struct ProgramState {
cursors: RefCell<BTreeMap<CursorID, Box<dyn Cursor>>>,
registers: Vec<OwnedValue>,
last_compare: Option<std::cmp::Ordering>,
deferred_seek: Option<(CursorID, CursorID)>,
ended_coroutine: bool, // flag to notify yield coroutine finished
regex_cache: RegexCache,
}
@ -458,6 +500,7 @@ impl ProgramState {
cursors,
registers,
last_compare: None,
deferred_seek: None,
ended_coroutine: false,
regex_cache: RegexCache::new(),
}
@ -958,6 +1001,13 @@ impl Program {
column,
dest,
} => {
if let Some((index_cursor_id, table_cursor_id)) = state.deferred_seek.take() {
let index_cursor = cursors.get_mut(&index_cursor_id).unwrap();
let rowid = index_cursor.rowid()?;
let table_cursor = cursors.get_mut(&table_cursor_id).unwrap();
table_cursor.seek_rowid(rowid.unwrap())?;
}
let cursor = cursors.get_mut(cursor_id).unwrap();
if let Some(ref record) = *cursor.record()? {
let null_flag = cursor.get_null_flag();
@ -1085,6 +1135,13 @@ impl Program {
state.pc += 1;
}
Insn::RowId { cursor_id, dest } => {
if let Some((index_cursor_id, table_cursor_id)) = state.deferred_seek.take() {
let index_cursor = cursors.get_mut(&index_cursor_id).unwrap();
let rowid = index_cursor.rowid()?;
let table_cursor = cursors.get_mut(&table_cursor_id).unwrap();
table_cursor.seek_rowid(rowid.unwrap())?;
}
let cursor = cursors.get_mut(cursor_id).unwrap();
if let Some(ref rowid) = cursor.rowid()? {
state.registers[*dest] = OwnedValue::Integer(*rowid as i64);
@ -1121,6 +1178,104 @@ impl Program {
}
}
}
Insn::DeferredSeek {
index_cursor_id,
table_cursor_id,
} => {
state.deferred_seek = Some((*index_cursor_id, *table_cursor_id));
state.pc += 1;
}
Insn::SeekGE {
cursor_id,
start_reg,
num_regs,
target_pc,
} => {
let cursor = cursors.get_mut(cursor_id).unwrap();
let record_from_regs: OwnedRecord =
make_owned_record(&state.registers, start_reg, num_regs);
match cursor.seek_ge(&record_from_regs)? {
CursorResult::Ok(found) => {
if !found {
state.pc = *target_pc;
} else {
state.pc += 1;
}
}
CursorResult::IO => {
// If there is I/O, the instruction is restarted.
return Ok(StepResult::IO);
}
}
}
Insn::SeekGT {
cursor_id,
start_reg,
num_regs,
target_pc,
} => {
let cursor = cursors.get_mut(cursor_id).unwrap();
let record_from_regs: OwnedRecord =
make_owned_record(&state.registers, start_reg, num_regs);
match cursor.seek_gt(&record_from_regs)? {
CursorResult::Ok(found) => {
if !found {
state.pc = *target_pc;
} else {
state.pc += 1;
}
}
CursorResult::IO => {
// If there is I/O, the instruction is restarted.
return Ok(StepResult::IO);
}
}
}
Insn::IdxGE {
cursor_id,
start_reg,
num_regs,
target_pc,
} => {
assert!(*target_pc >= 0);
let cursor = cursors.get_mut(cursor_id).unwrap();
let record_from_regs: OwnedRecord =
make_owned_record(&state.registers, start_reg, num_regs);
if let Some(ref idx_record) = *cursor.record()? {
// omit the rowid from the idx_record, which is the last value
if idx_record.values[..idx_record.values.len() - 1]
>= *record_from_regs.values
{
state.pc = *target_pc;
} else {
state.pc += 1;
}
} else {
state.pc = *target_pc;
}
}
Insn::IdxGT {
cursor_id,
start_reg,
num_regs,
target_pc,
} => {
let cursor = cursors.get_mut(cursor_id).unwrap();
let record_from_regs: OwnedRecord =
make_owned_record(&state.registers, start_reg, num_regs);
if let Some(ref idx_record) = *cursor.record()? {
// omit the rowid from the idx_record, which is the last value
if idx_record.values[..idx_record.values.len() - 1]
> *record_from_regs.values
{
state.pc = *target_pc;
} else {
state.pc += 1;
}
} else {
state.pc = *target_pc;
}
}
Insn::DecrJumpZero { reg, target_pc } => {
assert!(*target_pc >= 0);
match state.registers[*reg] {
@ -1863,7 +2018,10 @@ fn print_insn(program: &Program, addr: InsnReference, insn: &Insn, indent: Strin
fn get_indent_count(indent_count: usize, curr_insn: &Insn, prev_insn: Option<&Insn>) -> usize {
let indent_count = if let Some(insn) = prev_insn {
match insn {
Insn::RewindAwait { .. } | Insn::SorterSort { .. } => indent_count + 1,
Insn::RewindAwait { .. }
| Insn::SorterSort { .. }
| Insn::SeekGE { .. }
| Insn::SeekGT { .. } => indent_count + 1,
_ => indent_count,
}
} else {
@ -2369,6 +2527,14 @@ mod tests {
self.seek_rowid(rowid)
}
fn seek_ge(&mut self, key: &OwnedRecord) -> Result<CursorResult<bool>> {
unimplemented!();
}
fn seek_gt(&mut self, key: &OwnedRecord) -> Result<CursorResult<bool>> {
unimplemented!();
}
fn rewind(&mut self) -> Result<CursorResult<()>> {
unimplemented!()
}

8
core/vdbe/sorter.rs
View file

@ -79,6 +79,14 @@ impl Cursor for Sorter {
unimplemented!();
}
fn seek_ge(&mut self, key: &OwnedRecord) -> Result<CursorResult<bool>> {
unimplemented!();
}
fn seek_gt(&mut self, key: &OwnedRecord) -> Result<CursorResult<bool>> {
unimplemented!();
}
fn seek_to_last(&mut self) -> Result<CursorResult<()>> {
unimplemented!();
}

BIN
testing/testing.db
View file

Binary file not shown.