Skip to content

Commit

Permalink
refactor: simplify index cover implementation
Browse files Browse the repository at this point in the history
  • Loading branch information
@KKould
KKould committed Jan 17, 2025
1 parent 8683df7 commit 919adbf
Show file tree
Hide file tree
Showing 8 changed files with 167 additions and 145 deletions.
2 changes: 0 additions & 2 deletions core/translate/delete.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@ use crate::translate::planner::{parse_limit, parse_where};
use crate::{schema::Schema, storage::sqlite3_ondisk::DatabaseHeader, vdbe::Program};
use crate::{Connection, Result, SymbolTable};
use sqlite3_parser::ast::{Expr, Limit, QualifiedName};
use std::collections::BTreeSet;
use std::rc::Weak;
use std::{cell::RefCell, rc::Rc};

Expand Down Expand Up @@ -59,7 +58,6 @@ pub fn prepare_delete_plan(
iter_dir: None,
},
result_columns: vec![],
related_columns: BTreeSet::new(),
where_clause: resolved_where_clauses,
order_by: None,
limit: resolved_limit,
Expand Down
192 changes: 81 additions & 111 deletions core/translate/emitter.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,18 +2,18 @@
// It handles translating high-level SQL operations into low-level bytecode that can be executed by the virtual machine.

use std::cell::RefCell;
use std::collections::{BTreeSet, HashMap};
use std::collections::HashMap;
use std::rc::{Rc, Weak};

use sqlite3_parser::ast::{self};

use crate::schema::{Column, PseudoTable, Table};
use crate::storage::sqlite3_ondisk::DatabaseHeader;
use crate::translate::plan::{ColumnBinding, DeletePlan, IterationDirection, Plan, Search};
use crate::translate::plan::{DeletePlan, IterationDirection, Plan, Search};
use crate::types::{OwnedRecord, OwnedValue};
use crate::util::exprs_are_equivalent;
use crate::vdbe::builder::ProgramBuilder;
use crate::vdbe::{insn::Insn, BranchOffset, CursorID, Program};
use crate::vdbe::{insn::Insn, BranchOffset, Program};
use crate::{Connection, Result, SymbolTable};

use super::expr::{
Expand Down Expand Up @@ -378,7 +378,6 @@ fn emit_query(
open_loop(
program,
&mut plan.source,
&plan.related_columns,
&plan.referenced_tables,
metadata,
syms,
Expand Down Expand Up @@ -471,7 +470,6 @@ fn emit_program_for_delete(
open_loop(
&mut program,
&mut plan.source,
&plan.related_columns,
&plan.referenced_tables,
&mut metadata,
syms,
Expand Down Expand Up @@ -666,34 +664,26 @@ fn init_source(
search,
..
} => {
let table_cursor_id = program.alloc_cursor_id(
Some(table_reference.table_identifier.clone()),
Some(table_reference.table.clone()),
);

match mode {
OperationMode::SELECT => {
program.emit_insn(Insn::OpenReadAsync {
cursor_id: table_cursor_id,
root_page: table_reference.table.get_root_page(),
});
program.emit_insn(Insn::OpenReadAwait {});
}
OperationMode::DELETE => {
program.emit_insn(Insn::OpenWriteAsync {
cursor_id: table_cursor_id,
root_page: table_reference.table.get_root_page(),
});
program.emit_insn(Insn::OpenWriteAwait {});
}
_ => {
unimplemented!()
}
}
let mut is_cover = false;

if let Search::IndexSearch { index, .. } = search {
let index_cursor_id = program
.alloc_cursor_id(Some(index.name.clone()), Some(Table::Index(index.clone())));
if let Search::IndexSearch {
index,
cover_mapping,
..
} = search
{
let index_cursor_id = if cover_mapping.is_some() {
is_cover = true;
program.alloc_cursor_id(
Some(table_reference.table_identifier.clone()),
Some(Table::Index(index.clone())),
)
} else {
program.alloc_cursor_id(
Some(index.name.clone()),
Some(Table::Index(index.clone())),
)
};

match mode {
OperationMode::SELECT => {
Expand All @@ -716,6 +706,33 @@ fn init_source(
}
}

if !is_cover {
let table_cursor_id = program.alloc_cursor_id(
Some(table_reference.table_identifier.clone()),
Some(table_reference.table.clone()),
);

match mode {
OperationMode::SELECT => {
program.emit_insn(Insn::OpenReadAsync {
cursor_id: table_cursor_id,
root_page: table_reference.table.get_root_page(),
});
program.emit_insn(Insn::OpenReadAwait {});
}
OperationMode::DELETE => {
program.emit_insn(Insn::OpenWriteAsync {
cursor_id: table_cursor_id,
root_page: table_reference.table.get_root_page(),
});
program.emit_insn(Insn::OpenWriteAwait {});
}
_ => {
unimplemented!()
}
}
}

Ok(())
}
SourceOperator::Nothing { .. } => Ok(()),
Expand All @@ -728,7 +745,6 @@ fn init_source(
fn open_loop(
program: &mut ProgramBuilder,
source: &mut SourceOperator,
related_columns: &BTreeSet<ColumnBinding>,
referenced_tables: &[TableReference],
metadata: &mut Metadata,
syms: &SymbolTable,
Expand Down Expand Up @@ -803,14 +819,7 @@ fn open_loop(
outer,
..
} => {
open_loop(
program,
left,
related_columns,
referenced_tables,
metadata,
syms,
)?;
open_loop(program, left, referenced_tables, metadata, syms)?;

let loop_labels = metadata
.loop_labels
Expand All @@ -828,14 +837,7 @@ fn open_loop(
jump_target_when_false = lj_meta.check_match_flag_label;
}

open_loop(
program,
right,
related_columns,
referenced_tables,
metadata,
syms,
)?;
open_loop(program, right, referenced_tables, metadata, syms)?;

if let Some(predicates) = predicates {
let jump_target_when_true = program.allocate_label();
Expand Down Expand Up @@ -938,6 +940,7 @@ fn open_loop(
predicates,
..
} => {
let mut is_index_cover = false;
let table_cursor_id = program.resolve_cursor_id(&table_reference.table_identifier);
let loop_labels = metadata
.loop_labels
Expand All @@ -946,8 +949,21 @@ fn open_loop(
// Open the loop for the index search.
// Rowid equality point lookups are handled with a SeekRowid instruction which does not loop, since it is a single row lookup.
if !matches!(search, Search::RowidEq { .. }) {
let index_cursor_id = if let Search::IndexSearch { index, .. } = search {
Some(program.resolve_cursor_id(&index.name))
let index_cursor_id = if let Search::IndexSearch {
index,
cover_mapping,
..
} = search
{
if let Some(cover_mapping) = &cover_mapping {
program
.table_cursors
.insert(table_cursor_id, cover_mapping.clone());
is_index_cover = true;
Some(table_cursor_id)
} else {
Some(program.resolve_cursor_id(&index.name))
}
} else {
None
};
Expand Down Expand Up @@ -1086,19 +1102,11 @@ fn open_loop(
_ => {}
}

if let Some(index_cursor_id) = index_cursor_id {
if !try_index_cover(
program,
related_columns,
table_reference,
if let (Some(index_cursor_id), false) = (index_cursor_id, is_index_cover) {
program.emit_insn(Insn::DeferredSeek {
index_cursor_id,
table_cursor_id,
) {
program.emit_insn(Insn::DeferredSeek {
index_cursor_id,
table_cursor_id,
});
}
});
}
}

Expand Down Expand Up @@ -1147,54 +1155,6 @@ fn open_loop(
}
}

fn try_index_cover(
program: &mut ProgramBuilder,
related_columns: &BTreeSet<ColumnBinding>,
table_reference: &TableReference,
index_cursor_id: CursorID,
table_cursor_id: CursorID,
) -> bool {
let TableReference {
table, table_index, ..
} = table_reference;

if let (_, Some(Table::Index(index))) = &program.cursor_ref[index_cursor_id] {
let mut index_column_map = HashMap::with_capacity(index.columns.len());
let index_clone = Rc::clone(index);
for (i, index_column) in index_clone.columns.iter().enumerate() {
// SAFETY: the column in the index must exist in the table
let (pos, _) = table.get_column(&index_column.name).unwrap();
index_column_map.insert(pos, i);
}
if let Some(pos) = table
.columns()
.iter()
.position(|column| column.is_rowid_alias)
{
index_column_map.insert(pos, index_column_map.len());
}

let lower = ColumnBinding {
table: *table_index,
column: 0,
};
let upper = ColumnBinding {
table: table_index + 1,
column: 0,
};
let is_cover = related_columns
.range(lower..upper)
.all(|binding| index_column_map.contains_key(&binding.column));
if is_cover {
program
.index_cover_cursors
.insert(table_cursor_id, (index_cursor_id, index_column_map));
return true;
}
}
false
}

/// SQLite (and so Limbo) processes joins as a nested loop.
/// The inner loop may emit rows to various destinations depending on the query:
/// - a GROUP BY sorter (grouping is done by sorting based on the GROUP BY keys and aggregating while the GROUP BY keys match)
Expand Down Expand Up @@ -1532,7 +1492,17 @@ fn close_loop(
return Ok(());
}
let cursor_id = match search {
Search::IndexSearch { index, .. } => program.resolve_cursor_id(&index.name),
Search::IndexSearch {
index,
cover_mapping,
..
} => {
if cover_mapping.is_some() {
program.resolve_cursor_id(&table_reference.table_identifier)
} else {
program.resolve_cursor_id(&index.name)
}
}
Search::RowidSearch { .. } => {
program.resolve_cursor_id(&table_reference.table_identifier)
}
Expand Down
6 changes: 2 additions & 4 deletions core/translate/expr.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1967,12 +1967,10 @@ pub fn translate_expr(
// the table and read the column from the cursor.
TableReferenceType::BTreeTable => {
let cursor_id = program.resolve_cursor_id(&tbl_ref.table_identifier);
if let Some((index_cursor_id, columns_mapping)) =
program.index_cover_cursors.get(&cursor_id)
{
if let Some(columns_mapping) = program.table_cursors.get(&cursor_id) {
// FIXME: row_id is displayed as `column` in explain
program.emit_insn(Insn::Column {
cursor_id: *index_cursor_id,
cursor_id,
column: columns_mapping[column],
dest: target_register,
});
Expand Down
Loading

0 comments on commit 919adbf

Please sign in to comment.