aljaz/turso
1
0
Fork 0
mirror of https://github.com/aljazceru/turso.git synced 2026-02-02 14:54:23 +01:00
Code Issues Packages Projects Releases Wiki Activity

Merge '(core): Primary key index scans and single-column secondary index scans' from Jussi Saurio

Browse Source 
This PR adds an index on `users.age` to `testing.db`, and support for
indexed lookups. Only single-column ascending indexes are currently
supported.
This PR also gets rid of `Operator::Seekrowid` in favor of
`Operator::Search` which handles all non-full-table-scan searches: 1.
integer primary key (rowid) point queries 2. integer primary key index
scans, and 3. secondary index scans.
examples:
```
limbo> select first_name, age from users where age > 90 limit 10;
Miranda|90
Sarah|90
Justin|90
Justin|90
John|90
Jeremy|90
Stephanie|90
Joshua|90
Jenny|90
Jennifer|90

limbo> explain query plan select first_name, age from users where age > 90 limit 10;
QUERY PLAN
`--TAKE 10
   `--PROJECT first_name, age
   |  `--SEARCH users USING INDEX age_idx

limbo> explain select first_name, age from users where age > 90 limit 10;
addr  opcode             p1    p2    p3    p4             p5  comment
----  -----------------  ----  ----  ----  -------------  --  -------
0     Init               0     15    0                    0   Start at 15
1     OpenReadAsync      0     2     0                    0   table=users, root=2
2     OpenReadAwait      0     0     0                    0
3     OpenReadAsync      1     274   0                    0   table=age_idx, root=274
4     OpenReadAwait      0     0     0                    0
5     Integer            90    1     0                    0   r[1]=90
6     SeekGT             1     14    1                    0
7       DeferredSeek     1     0     0                    0
8       Column           0     1     2                    0   r[2]=users.first_name
9       Column           0     9     3                    0   r[3]=users.age
10      ResultRow        2     2     0                    0   output=r[2..3]
11      DecrJumpZero     4     14    0                    0   if (--r[4]==0) goto 14
12    NextAsync          1     0     0                    0
13    NextAwait          1     7     0                    0
14    Halt               0     0     0                    0
15    Transaction        0     0     0                    0
16    Integer            10    4     0                    0   r[4]=10
17    Goto               0     1     0                    0
```
Sqlite version:
```
sqlite> explain select first_name, age from users where age > 90 limit 10;
addr  opcode         p1    p2    p3    p4             p5  comment
----  -------------  ----  ----  ----  -------------  --  -------------
0     Init           0     13    0                    0   Start at 13
1     Integer        10    1     0                    0   r[1]=10; LIMIT counter
2     OpenRead       0     2     0     10             0   root=2 iDb=0; users
3     OpenRead       1     274   0     k(2,,)         0   root=274 iDb=0; age_idx
4     Integer        90    2     0                    0   r[2]=90
5     SeekGT         1     12    2     1              0   key=r[2]
6       DeferredSeek   1     0     0                    0   Move 0 to 1.rowid if needed
7       Column         0     1     3                    0   r[3]= cursor 0 column 1
8       Column         1     0     4                    0   r[4]= cursor 1 column 0
9       ResultRow      3     2     0                    0   output=r[3..4]
10      DecrJumpZero   1     12    0                    0   if (--r[1])==0 goto 12
11    Next           1     6     0                    0
12    Halt           0     0     0                    0
13    Transaction    0     0     3     0              1   usesStmtJournal=0
14    Goto           0     1     0                    0
```
---
´Seek` instructions are also now supported for primary key rowid
searches:
```
limbo> select id, first_name from users where id > 9995;
9996|Donald
9997|Ruth
9998|Dorothy
9999|Gina
10000|Nicole

limbo> explain query plan select id, first_name from users where id > 9995;
QUERY PLAN
`--PROJECT id, first_name
   `--SEARCH users USING INTEGER PRIMARY KEY (rowid=?)

limbo> explain select id, first_name from users where id > 9995;
addr  opcode             p1    p2    p3    p4             p5  comment
----  -----------------  ----  ----  ----  -------------  --  -------
0     Init               0     11    0                    0   Start at 11
1     OpenReadAsync      0     2     0                    0   table=users, root=2
2     OpenReadAwait      0     0     0                    0
3     Integer            9995  1     0                    0   r[1]=9995
4     SeekGT             0     10    1                    0
5       RowId            0     2     0                    0   r[2]=users.rowid
6       Column           0     1     3                    0   r[3]=users.first_name
7       ResultRow        2     2     0                    0   output=r[2..3]
8     NextAsync          0     0     0                    0
9     NextAwait          0     5     0                    0
10    Halt               0     0     0                    0
11    Transaction        0     0     0                    0
12    Goto               0     1     0                    0
```
sqlite:
```
sqlite> explain select id, first_name from users where id > 9995;
addr  opcode         p1    p2    p3    p4             p5  comment
----  -------------  ----  ----  ----  -------------  --  -------------
0     Init           0     8     0                    0   Start at 8
1     OpenRead       0     2     0     2              0   root=2 iDb=0; users
2     SeekGT         0     7     1                    0   key=r[1]; pk
3       Rowid          0     2     0                    0   r[2]=users.rowid
4       Column         0     1     3                    0   r[3]= cursor 0 column 1
5       ResultRow      2     2     0                    0   output=r[2..3]
6     Next           0     3     0                    0
7     Halt           0     0     0                    0
8     Transaction    0     0     3     0              1   usesStmtJournal=0
9     Integer        9995  1     0                    0   r[1]=9995
10    Goto           0     1     0                    0
```
---
More complex example with a join that uses both a rowid lookup and a
secondary index scan:
```
limbo> explain query plan select u.first_name, p.name from users u join products p on u.id = p.id and u.age > 70;
QUERY PLAN
`--PROJECT u.first_name, p.name
   `--JOIN
   |  |--SEARCH u USING INDEX age_idx
   |  `--SEARCH p USING INTEGER PRIMARY KEY (rowid=?)

limbo> explain select u.first_name, p.name from users u join products p on u.id = p.id and u.age > 70;
addr  opcode             p1    p2    p3    p4             p5  comment
----  -----------------  ----  ----  ----  -------------  --  -------
0     Init               0     18    0                    0   Start at 18
1     OpenReadAsync      0     2     0                    0   table=u, root=2
2     OpenReadAwait      0     0     0                    0
3     OpenReadAsync      1     274   0                    0   table=age_idx, root=274
4     OpenReadAwait      0     0     0                    0
5     OpenReadAsync      2     3     0                    0   table=p, root=3
6     OpenReadAwait      0     0     0                    0
7     Integer            70    1     0                    0   r[1]=70
8     SeekGT             1     17    1                    0
9       DeferredSeek     1     0     0                    0
10      RowId            0     2     0                    0   r[2]=u.rowid
11      SeekRowid        2     2     15                   0   if (r[2]!=p.rowid) goto 15
12      Column           0     1     3                    0   r[3]=u.first_name
13      Column           2     1     4                    0   r[4]=p.name
14      ResultRow        3     2     0                    0   output=r[3..4]
15    NextAsync          1     0     0                    0
16    NextAwait          1     9     0                    0
17    Halt               0     0     0                    0
18    Transaction        0     0     0                    0
19    Goto               0     1     0                    0
```
sqlite:
```
sqlite> explain select u.first_name, p.name from users u join products p on u.id = p.id and u.age > 70;
addr  opcode         p1    p2    p3    p4             p5  comment
----  -------------  ----  ----  ----  -------------  --  -------------
0     Init           0     14    0                    0   Start at 14
1     OpenRead       0     2     0     10             0   root=2 iDb=0; users
2     OpenRead       2     274   0     k(2,,)         0   root=274 iDb=0; age_idx
3     OpenRead       1     3     0     2              0   root=3 iDb=0; products
4     Integer        70    1     0                    0   r[1]=70
5     SeekGT         2     13    1     1              0   key=r[1]
6       DeferredSeek   2     0     0                    0   Move 0 to 2.rowid if needed
7       IdxRowid       2     2     0                    0   r[2]=rowid; users.rowid
8       SeekRowid      1     12    2                    0   intkey=r[2]
9       Column         0     1     3                    0   r[3]= cursor 0 column 1
10      Column         1     1     4                    0   r[4]= cursor 1 column 1
11      ResultRow      3     2     0                    0   output=r[3..4]
12    Next           2     6     0                    0
13    Halt           0     0     0                    0
14    Transaction    0     0     3     0              1   usesStmtJournal=0
15    Goto           0     1     0                    0
```

Closes #350
This commit is contained in:
Pekka Enberg
2024-10-13 10:11:23 +03:00
parent c58b7b42c0 556f4b73c9
commit fdf3b0d16b
20 changed files with 1217 additions and 307 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
cli/main.rs
View File

@@ -201,7 +201,7 @@ fn display_schema(
) -> anyhow::Result<()> {
let sql = match table {
Some(table_name) => format!(
"SELECT sql FROM sqlite_schema WHERE type='table' AND name = '{}' AND name NOT LIKE 'sqlite_%'",
"SELECT sql FROM sqlite_schema WHERE type IN ('table', 'index') AND tbl_name = '{}' AND name NOT LIKE 'sqlite_%'",
table_name
),
None => String::from(

7
core/pseudo.rs
View File

@@ -1,4 +1,7 @@
use crate::Result;
use crate::{
types::{SeekKey, SeekOp},
Result,
};
use std::cell::{Ref, RefCell};
use crate::types::{Cursor, CursorResult, OwnedRecord, OwnedValue};
@@ -48,7 +51,7 @@ impl Cursor for PseudoCursor {
Ok(x)
}
fn seek_rowid(&mut self, _: u64) -> Result<CursorResult<bool>> {
fn seek(&mut self, _: SeekKey<'_>, _: SeekOp) -> Result<CursorResult<bool>> {
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!(),
}
}
}

221
core/storage/btree.rs
View File

@@ -5,13 +5,13 @@ use crate::storage::sqlite3_ondisk::{
read_btree_cell, read_varint, write_varint, BTreeCell, DatabaseHeader, PageContent, PageType,
TableInteriorCell, TableLeafCell,
};
use crate::types::{Cursor, CursorResult, OwnedRecord, OwnedValue};
use crate::types::{Cursor, CursorResult, OwnedRecord, OwnedValue, SeekKey, SeekOp};
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,7 @@ 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(Debug)]
pub struct MemPage {
parent: Option<Rc<MemPage>>,
page_idx: usize,
@@ -56,6 +57,7 @@ pub struct BTreeCursor {
record: RefCell<Option<OwnedRecord>>,
null_flag: bool,
database_header: Rc<RefCell<DatabaseHeader>>,
going_upwards: bool,
}
impl BTreeCursor {
@@ -72,6 +74,7 @@ impl BTreeCursor {
record: RefCell::new(None),
null_flag: false,
database_header,
going_upwards: false,
}
}
@@ -87,7 +90,10 @@ impl BTreeCursor {
Ok(CursorResult::Ok(page.cell_count() == 0))
}
fn get_next_record(&mut self) -> Result<CursorResult<(Option<u64>, Option<OwnedRecord>)>> {
fn get_next_record<'a>(
&mut self,
predicate: Option<(SeekKey<'a>, SeekOp)>,
) -> Result<CursorResult<(Option<u64>, Option<OwnedRecord>)>> {
loop {
let mem_page = self.get_mem_page();
let page_idx = mem_page.page_idx;
@@ -109,6 +115,7 @@ impl BTreeCursor {
}
None => match parent {
Some(ref parent) => {
self.going_upwards = true;
self.page.replace(Some(parent.clone()));
continue;
}
@@ -130,6 +137,7 @@ impl BTreeCursor {
_left_child_page,
_rowid,
}) => {
assert!(predicate.is_none());
mem_page.advance();
let mem_page =
MemPage::new(Some(mem_page.clone()), *_left_child_page as usize, 0);
@@ -141,61 +149,181 @@ impl BTreeCursor {
_payload,
first_overflow_page: _,
}) => {
assert!(predicate.is_none());
mem_page.advance();
let record = crate::storage::sqlite3_ondisk::read_record(_payload)?;
return Ok(CursorResult::Ok((Some(*_rowid), Some(record))));
}
BTreeCell::IndexInteriorCell(_) => {
unimplemented!();
BTreeCell::IndexInteriorCell(IndexInteriorCell {
payload,
left_child_page,
..
}) => {
if !self.going_upwards {
let mem_page =
MemPage::new(Some(mem_page.clone()), *left_child_page as usize, 0);
self.page.replace(Some(Rc::new(mem_page)));
continue;
}
self.going_upwards = false;
mem_page.advance();
let record = crate::storage::sqlite3_ondisk::read_record(payload)?;
if predicate.is_none() {
let rowid = match record.values.last() {
Some(OwnedValue::Integer(rowid)) => *rowid as u64,
_ => unreachable!("index cells should have an integer rowid"),
};
return Ok(CursorResult::Ok((Some(rowid), Some(record))));
}
let (key, op) = predicate.as_ref().unwrap();
let SeekKey::IndexKey(index_key) = key else {
unreachable!("index seek key should be a record");
};
let found = match op {
SeekOp::GT => &record > *index_key,
SeekOp::GE => &record >= *index_key,
SeekOp::EQ => &record == *index_key,
};
if found {
let rowid = match record.values.last() {
Some(OwnedValue::Integer(rowid)) => *rowid as u64,
_ => unreachable!("index cells should have an integer rowid"),
};
return Ok(CursorResult::Ok((Some(rowid), Some(record))));
} else {
continue;
}
}
BTreeCell::IndexLeafCell(_) => {
unimplemented!();
BTreeCell::IndexLeafCell(IndexLeafCell { payload, .. }) => {
mem_page.advance();
let record = crate::storage::sqlite3_ondisk::read_record(payload)?;
if predicate.is_none() {
let rowid = match record.values.last() {
Some(OwnedValue::Integer(rowid)) => *rowid as u64,
_ => unreachable!("index cells should have an integer rowid"),
};
return Ok(CursorResult::Ok((Some(rowid), Some(record))));
}
let (key, op) = predicate.as_ref().unwrap();
let SeekKey::IndexKey(index_key) = key else {
unreachable!("index seek key should be a record");
};
let found = match op {
SeekOp::GT => &record > *index_key,
SeekOp::GE => &record >= *index_key,
SeekOp::EQ => &record == *index_key,
};
if found {
let rowid = match record.values.last() {
Some(OwnedValue::Integer(rowid)) => *rowid as u64,
_ => unreachable!("index cells should have an integer rowid"),
};
return Ok(CursorResult::Ok((Some(rowid), Some(record))));
} else {
continue;
}
}
}
}
}
fn btree_seek_rowid(
fn seek<'a>(
&mut self,
rowid: u64,
key: SeekKey<'a>,
op: SeekOp,
) -> Result<CursorResult<(Option<u64>, Option<OwnedRecord>)>> {
self.move_to(rowid)?;
match self.move_to(key.clone(), op.clone())? {
CursorResult::Ok(_) => {}
CursorResult::IO => return Ok(CursorResult::IO),
};
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(
let cell = page.cell_get(
cell_idx,
self.pager.clone(),
self.max_local(page.page_type()),
self.min_local(page.page_type()),
self.usable_space(),
)? {
)?;
match &cell {
BTreeCell::TableLeafCell(TableLeafCell {
_rowid: cell_rowid,
_payload: p,
_payload: payload,
first_overflow_page: _,
}) => {
if *cell_rowid == rowid {
let record = crate::storage::sqlite3_ondisk::read_record(p)?;
let SeekKey::TableRowId(rowid_key) = key else {
unreachable!("table seek key should be a rowid");
};
mem_page.advance();
let found = match op {
SeekOp::GT => *cell_rowid > rowid_key,
SeekOp::GE => *cell_rowid >= rowid_key,
SeekOp::EQ => *cell_rowid == rowid_key,
};
if found {
let record = crate::storage::sqlite3_ondisk::read_record(payload)?;
return Ok(CursorResult::Ok((Some(*cell_rowid), Some(record))));
}
}
BTreeCell::IndexLeafCell(IndexLeafCell { payload, .. }) => {
let SeekKey::IndexKey(index_key) = key else {
unreachable!("index seek key should be a record");
};
mem_page.advance();
let record = crate::storage::sqlite3_ondisk::read_record(payload)?;
let found = match op {
SeekOp::GT => record > *index_key,
SeekOp::GE => record >= *index_key,
SeekOp::EQ => record == *index_key,
};
if found {
let rowid = match record.values.last() {
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);
}
}
}
// We have now iterated over all cells in the leaf page and found no match.
let is_index = matches!(key, SeekKey::IndexKey(_));
if is_index {
// Unlike tables, indexes store payloads in interior cells as well. self.move_to() always moves to a leaf page, so there are cases where we need to
// move back up to the parent interior cell and get the next record from there to perform a correct seek.
// an example of how this can occur:
//
// we do an index seek for key K with cmp = SeekOp::GT, meaning we want to seek to the first key that is greater than K.
// in self.move_to(), we encounter an interior cell with key K' = K+2, and move the left child page, which is a leaf page.
// the reason we move to the left child page is that we know that in an index, all keys in the left child page are less than K' i.e. less than K+2,
// meaning that the left subtree may contain a key greater than K, e.g. K+1. however, it is possible that it doesn't, in which case the correct
// next key is K+2, which is in the parent interior cell.
//
// In the seek() method, once we have landed in the leaf page and find that there is no cell with a key greater than K,
// if we were to return Ok(CursorResult::Ok((None, None))), self.record would be None, which is incorrect, because we already know
// that there is a record with a key greater than K (K' = K+2) in the parent interior cell. Hence, we need to move back up the tree
// and get the next matching record from there.
return self.get_next_record(Some((key, op)));
}
Ok(CursorResult::Ok((None, None)))
}
@@ -240,7 +368,7 @@ impl BTreeCursor {
}
}
pub fn move_to(&mut self, key: u64) -> Result<CursorResult<()>> {
pub fn move_to<'a>(&mut self, key: SeekKey<'a>, cmp: SeekOp) -> Result<CursorResult<()>> {
// For a table with N rows, we can find any row by row id in O(log(N)) time by starting at the root page and following the B-tree pointers.
// B-trees consist of interior pages and leaf pages. Interior pages contain pointers to other pages, while leaf pages contain the actual row data.
//
@@ -294,8 +422,16 @@ impl BTreeCursor {
_left_child_page,
_rowid,
}) => {
if key < *_rowid {
mem_page.advance();
let SeekKey::TableRowId(rowid_key) = key else {
unreachable!("table seek key should be a rowid");
};
mem_page.advance();
let target_leaf_page_is_in_left_subtree = match cmp {
SeekOp::GT => rowid_key < *_rowid,
SeekOp::GE => rowid_key <= *_rowid,
SeekOp::EQ => rowid_key <= *_rowid,
};
if target_leaf_page_is_in_left_subtree {
let mem_page =
MemPage::new(Some(mem_page.clone()), *_left_child_page as usize, 0);
self.page.replace(Some(Rc::new(mem_page)));
@@ -312,20 +448,43 @@ impl BTreeCursor {
"we don't iterate leaf cells while trying to move to a leaf cell"
);
}
BTreeCell::IndexInteriorCell(_) => {
unimplemented!();
BTreeCell::IndexInteriorCell(IndexInteriorCell {
left_child_page,
payload,
..
}) => {
let SeekKey::IndexKey(index_key) = key else {
unreachable!("index seek key should be a record");
};
let record = crate::storage::sqlite3_ondisk::read_record(payload)?;
let target_leaf_page_is_in_the_left_subtree = match cmp {
SeekOp::GT => index_key < &record,
SeekOp::GE => index_key <= &record,
SeekOp::EQ => index_key <= &record,
};
if target_leaf_page_is_in_the_left_subtree {
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;
} else {
mem_page.advance();
}
}
BTreeCell::IndexLeafCell(_) => {
unimplemented!();
unreachable!(
"we don't iterate leaf cells while trying to move to a leaf cell"
);
}
}
}
if !found_cell {
let parent = mem_page.clone();
let parent = mem_page.parent.clone();
match page.rightmost_pointer() {
Some(right_most_pointer) => {
let mem_page = MemPage::new(Some(parent), right_most_pointer as usize, 0);
let mem_page = MemPage::new(parent, right_most_pointer as usize, 0);
self.page.replace(Some(Rc::new(mem_page)));
continue;
}
@@ -1230,7 +1389,7 @@ fn find_free_cell(page_ref: &PageContent, db_header: Ref<DatabaseHeader>, amount
impl Cursor for BTreeCursor {
fn seek_to_last(&mut self) -> Result<CursorResult<()>> {
self.move_to_rightmost()?;
match self.get_next_record()? {
match self.get_next_record(None)? {
CursorResult::Ok((rowid, next)) => {
if rowid.is_none() {
match self.is_empty_table()? {
@@ -1255,7 +1414,7 @@ impl Cursor for BTreeCursor {
fn rewind(&mut self) -> Result<CursorResult<()>> {
let mem_page = MemPage::new(None, self.root_page, 0);
self.page.replace(Some(Rc::new(mem_page)));
match self.get_next_record()? {
match self.get_next_record(None)? {
CursorResult::Ok((rowid, next)) => {
self.rowid.replace(rowid);
self.record.replace(next);
@@ -1266,7 +1425,7 @@ impl Cursor for BTreeCursor {
}
fn next(&mut self) -> Result<CursorResult<()>> {
match self.get_next_record()? {
match self.get_next_record(None)? {
CursorResult::Ok((rowid, next)) => {
self.rowid.replace(rowid);
self.record.replace(next);
@@ -1285,8 +1444,8 @@ impl Cursor for BTreeCursor {
Ok(*self.rowid.borrow())
}
fn seek_rowid(&mut self, rowid: u64) -> Result<CursorResult<bool>> {
match self.btree_seek_rowid(rowid)? {
fn seek(&mut self, key: SeekKey<'_>, op: SeekOp) -> Result<CursorResult<bool>> {
match self.seek(key, op)? {
CursorResult::Ok((rowid, record)) => {
self.rowid.replace(rowid);
self.record.replace(record);
@@ -1311,7 +1470,7 @@ impl Cursor for BTreeCursor {
_ => unreachable!("btree tables are indexed by integers!"),
};
if !moved_before {
match self.move_to(*int_key as u64)? {
match self.move_to(SeekKey::TableRowId(*int_key as u64), SeekOp::EQ)? {
CursorResult::Ok(_) => {}
CursorResult::IO => return Ok(CursorResult::IO),
};
@@ -1336,7 +1495,7 @@ impl Cursor for BTreeCursor {
OwnedValue::Integer(i) => i,
_ => unreachable!("btree tables are indexed by integers!"),
};
match self.move_to(*int_key as u64)? {
match self.move_to(SeekKey::TableRowId(*int_key as u64), SeekOp::EQ)? {
CursorResult::Ok(_) => {}
CursorResult::IO => return Ok(CursorResult::IO),
};

342
core/translate/emitter.rs
View File

@@ -2,9 +2,12 @@ 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;
use crate::translate::plan::Search;
use crate::types::{OwnedRecord, OwnedValue};
use crate::vdbe::builder::ProgramBuilder;
use crate::vdbe::{BranchOffset, Insn, Program};
@@ -174,7 +177,7 @@ impl Emitter for Operator {
} => {
*step += 1;
const SCAN_OPEN_READ: usize = 1;
const SCAN_REWIND_AND_CONDITIONS: usize = 2;
const SCAN_BODY: usize = 2;
const SCAN_NEXT: usize = 3;
match *step {
SCAN_OPEN_READ => {
@@ -193,7 +196,7 @@ impl Emitter for Operator {
Ok(OpStepResult::Continue)
}
SCAN_REWIND_AND_CONDITIONS => {
SCAN_BODY => {
let cursor_id = program.resolve_cursor_id(table_identifier, None);
program.emit_insn(Insn::RewindAsync { cursor_id });
let scan_loop_body_label = program.allocate_label();
@@ -251,56 +254,242 @@ impl Emitter for Operator {
_ => Ok(OpStepResult::Done),
}
}
Operator::SeekRowid {
Operator::Search {
table,
table_identifier,
rowid_predicate,
search,
predicates,
step,
id,
..
} => {
*step += 1;
const SEEKROWID_OPEN_READ: usize = 1;
const SEEKROWID_SEEK_AND_CONDITIONS: usize = 2;
const SEARCH_OPEN_READ: usize = 1;
const SEARCH_BODY: usize = 2;
const SEARCH_NEXT: usize = 3;
match *step {
SEEKROWID_OPEN_READ => {
let cursor_id = program.alloc_cursor_id(
SEARCH_OPEN_READ => {
let table_cursor_id = program.alloc_cursor_id(
Some(table_identifier.clone()),
Some(Table::BTree(table.clone())),
);
let root_page = table.root_page;
let next_row_label = program.allocate_label();
if !matches!(search, Search::PrimaryKeyEq { .. }) {
// Primary key equality search is handled with a SeekRowid instruction which does not loop, since it is a single row lookup.
m.next_row_labels.insert(*id, next_row_label);
}
let scan_loop_body_label = program.allocate_label();
m.scan_loop_body_labels.push(scan_loop_body_label);
program.emit_insn(Insn::OpenReadAsync {
cursor_id,
root_page,
cursor_id: table_cursor_id,
root_page: table.root_page,
});
program.emit_insn(Insn::OpenReadAwait);
if let Search::IndexSearch { index, .. } = search {
let index_cursor_id = program.alloc_cursor_id(
Some(index.name.clone()),
Some(Table::Index(index.clone())),
);
program.emit_insn(Insn::OpenReadAsync {
cursor_id: index_cursor_id,
root_page: index.root_page,
});
program.emit_insn(Insn::OpenReadAwait);
}
Ok(OpStepResult::Continue)
}
SEEKROWID_SEEK_AND_CONDITIONS => {
let cursor_id = program.resolve_cursor_id(table_identifier, None);
let rowid_reg = program.alloc_register();
translate_expr(
program,
Some(referenced_tables),
rowid_predicate,
rowid_reg,
None,
None,
)?;
SEARCH_BODY => {
let table_cursor_id = program.resolve_cursor_id(table_identifier, None);
// Open the loop for the index search.
// Primary key equality search is handled with a SeekRowid instruction which does not loop, since it is a single row lookup.
if !matches!(search, Search::PrimaryKeyEq { .. }) {
let index_cursor_id = if let Search::IndexSearch { index, .. } = search
{
Some(program.resolve_cursor_id(&index.name, None))
} else {
None
};
let scan_loop_body_label = *m.scan_loop_body_labels.last().unwrap();
let cmp_reg = program.alloc_register();
let (cmp_expr, cmp_op) = match search {
Search::IndexSearch {
cmp_expr, cmp_op, ..
} => (cmp_expr, cmp_op),
Search::PrimaryKeySearch { cmp_expr, cmp_op } => (cmp_expr, cmp_op),
Search::PrimaryKeyEq { .. } => unreachable!(),
};
// TODO this only handles ascending indexes
match cmp_op {
ast::Operator::Equals
| ast::Operator::Greater
| ast::Operator::GreaterEquals => {
translate_expr(
program,
Some(referenced_tables),
cmp_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 cmp_op {
ast::Operator::Equals | ast::Operator::GreaterEquals => {
Insn::SeekGE {
is_index: index_cursor_id.is_some(),
cursor_id: index_cursor_id.unwrap_or(table_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 {
is_index: index_cursor_id.is_some(),
cursor_id: index_cursor_id.unwrap_or(table_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 *cmp_op == ast::Operator::Less
|| *cmp_op == ast::Operator::LessEquals
{
translate_expr(
program,
Some(referenced_tables),
cmp_expr,
cmp_reg,
None,
None,
)?;
}
program.defer_label_resolution(
scan_loop_body_label,
program.offset() as usize,
);
// TODO: 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.
//
// For primary key searches we emit RowId and then compare it to the seek value.
let abort_jump_target = *m
.next_row_labels
.get(id)
.unwrap_or(m.termination_label_stack.last().unwrap());
match cmp_op {
ast::Operator::Equals | ast::Operator::LessEquals => {
if index_cursor_id.is_some() {
program.emit_insn_with_label_dependency(
Insn::IdxGT {
cursor_id: index_cursor_id.unwrap(),
start_reg: cmp_reg,
num_regs: 1,
target_pc: abort_jump_target,
},
abort_jump_target,
);
} else {
let rowid_reg = program.alloc_register();
program.emit_insn(Insn::RowId {
cursor_id: table_cursor_id,
dest: rowid_reg,
});
program.emit_insn_with_label_dependency(
Insn::Gt {
lhs: rowid_reg,
rhs: cmp_reg,
target_pc: abort_jump_target,
},
abort_jump_target,
);
}
}
ast::Operator::Less => {
if index_cursor_id.is_some() {
program.emit_insn_with_label_dependency(
Insn::IdxGE {
cursor_id: index_cursor_id.unwrap(),
start_reg: cmp_reg,
num_regs: 1,
target_pc: abort_jump_target,
},
abort_jump_target,
);
} else {
let rowid_reg = program.alloc_register();
program.emit_insn(Insn::RowId {
cursor_id: table_cursor_id,
dest: rowid_reg,
});
program.emit_insn_with_label_dependency(
Insn::Ge {
lhs: rowid_reg,
rhs: cmp_reg,
target_pc: abort_jump_target,
},
abort_jump_target,
);
}
}
_ => {}
}
if index_cursor_id.is_some() {
program.emit_insn(Insn::DeferredSeek {
index_cursor_id: index_cursor_id.unwrap(),
table_cursor_id,
});
}
}
let jump_label = m
.next_row_labels
.get(id)
.unwrap_or(m.termination_label_stack.last().unwrap());
program.emit_insn_with_label_dependency(
Insn::SeekRowid {
cursor_id,
src_reg: rowid_reg,
target_pc: *jump_label,
},
*jump_label,
);
if let Search::PrimaryKeyEq { cmp_expr } = search {
let src_reg = program.alloc_register();
translate_expr(
program,
Some(referenced_tables),
cmp_expr,
src_reg,
None,
None,
)?;
program.emit_insn_with_label_dependency(
Insn::SeekRowid {
cursor_id: table_cursor_id,
src_reg,
target_pc: *jump_label,
},
*jump_label,
);
}
if let Some(predicates) = predicates {
for predicate in predicates.iter() {
let jump_target_when_true = program.allocate_label();
@@ -322,6 +511,33 @@ impl Emitter for Operator {
Ok(OpStepResult::ReadyToEmit)
}
SEARCH_NEXT => {
if matches!(search, Search::PrimaryKeyEq { .. }) {
// Primary key equality search is handled with a SeekRowid instruction which does not loop, so there is no need to emit a NextAsync instruction.
return Ok(OpStepResult::Done);
}
let cursor_id = match search {
Search::IndexSearch { index, .. } => {
program.resolve_cursor_id(&index.name, None)
}
Search::PrimaryKeySearch { .. } => {
program.resolve_cursor_id(table_identifier, None)
}
Search::PrimaryKeyEq { .. } => unreachable!(),
};
program
.resolve_label(*m.next_row_labels.get(id).unwrap(), program.offset());
program.emit_insn(Insn::NextAsync { cursor_id });
let jump_label = m.scan_loop_body_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),
}
}
@@ -428,7 +644,7 @@ impl Emitter for Operator {
Operator::Scan {
table_identifier, ..
} => program.resolve_cursor_id(table_identifier, None),
Operator::SeekRowid {
Operator::Search {
table_identifier, ..
} => program.resolve_cursor_id(table_identifier, None),
_ => unreachable!(),
@@ -443,9 +659,14 @@ impl Emitter for Operator {
},
lj_meta.set_match_flag_true_label,
);
// This points to the NextAsync instruction of the left table
program.resolve_label(lj_meta.on_match_jump_to_label, program.offset());
}
let next_row_label = if *outer {
m.left_joins.get(id).unwrap().on_match_jump_to_label
} else {
*m.next_row_labels.get(&right.id()).unwrap()
};
// This points to the NextAsync instruction of the left table
program.resolve_label(next_row_label, program.offset());
left.step(program, m, referenced_tables)?;
Ok(OpStepResult::Done)
@@ -1196,6 +1417,23 @@ impl Emitter for Operator {
Ok(start_reg)
}
Operator::Search {
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)?;
@@ -1256,23 +1494,6 @@ impl Emitter for Operator {
}
}
Operator::Filter { .. } => unreachable!("predicates have been pushed down"),
Operator::SeekRowid {
table_identifier,
table,
..
} => {
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::Limit { .. } => {
unimplemented!()
}
@@ -1415,13 +1636,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 +1661,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 +1697,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 +1713,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!(),
};

320
core/translate/optimizer.rs
View File

@@ -2,11 +2,15 @@ use std::{collections::HashMap, rc::Rc};
use sqlite3_parser::ast;
use crate::{schema::BTreeTable, util::normalize_ident, Result};
use crate::{
schema::{BTreeTable, Index},
util::normalize_ident,
Result,
};
use super::plan::{
get_table_ref_bitmask_for_ast_expr, get_table_ref_bitmask_for_operator, Operator, Plan,
ProjectionColumn,
ProjectionColumn, Search,
};
/**
@@ -25,6 +29,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 +37,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 +49,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::Search { .. } => Ok(()),
Operator::Scan {
table,
predicates: filter,
@@ -57,68 +65,57 @@ fn use_indexes(
}
let fs = filter.as_mut().unwrap();
let mut i = 0;
let mut maybe_rowid_predicate = None;
while i < fs.len() {
for i in 0..fs.len() {
let f = fs[i].take_ownership();
let table_index = referenced_tables
let table = referenced_tables
.iter()
.position(|(t, t_id)| Rc::ptr_eq(t, table) && t_id == table_identifier)
.find(|(t, t_id)| Rc::ptr_eq(t, table) && t_id == table_identifier)
.unwrap();
let (can_use, expr) =
try_extract_rowid_comparison_expression(f, table_index, referenced_tables)?;
if can_use {
maybe_rowid_predicate = Some(expr);
fs.remove(i);
break;
} else {
fs[i] = expr;
i += 1;
}
}
match try_extract_index_search_expression(f, table, available_indexes)? {
Either::Left(non_index_using_expr) => {
fs[i] = non_index_using_expr;
}
Either::Right(index_search) => {
fs.remove(i);
*operator = Operator::Search {
id: *id,
table: table.0.clone(),
table_identifier: table.1.clone(),
predicates: Some(fs.clone()),
search: index_search,
step: 0,
};
if let Some(rowid_predicate) = maybe_rowid_predicate {
let predicates_owned = if fs.is_empty() {
None
} else {
Some(std::mem::take(fs))
};
*operator = Operator::SeekRowid {
table: table.clone(),
table_identifier: table_identifier.clone(),
rowid_predicate,
predicates: predicates_owned,
id: *id,
step: 0,
return Ok(());
}
}
}
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(()),
@@ -206,31 +203,6 @@ fn eliminate_constants(operator: &mut Operator) -> Result<ConstantConditionElimi
// Aggregation operator can return a row even if the source is empty e.g. count(1) from users where 0
Ok(ConstantConditionEliminationResult::Continue)
}
Operator::SeekRowid {
rowid_predicate,
predicates,
..
} => {
if let Some(predicates) = predicates {
let mut i = 0;
while i < predicates.len() {
let predicate = &predicates[i];
if predicate.is_always_true()? {
predicates.remove(i);
} else if predicate.is_always_false()? {
return Ok(ConstantConditionEliminationResult::ImpossibleCondition);
} else {
i += 1;
}
}
}
if rowid_predicate.is_always_false()? {
return Ok(ConstantConditionEliminationResult::ImpossibleCondition);
}
Ok(ConstantConditionEliminationResult::Continue)
}
Operator::Limit { source, .. } => {
let constant_elimination_result = eliminate_constants(source)?;
if constant_elimination_result
@@ -279,6 +251,23 @@ fn eliminate_constants(operator: &mut Operator) -> Result<ConstantConditionElimi
}
Ok(ConstantConditionEliminationResult::Continue)
}
Operator::Search { predicates, .. } => {
if let Some(predicates) = predicates {
let mut i = 0;
while i < predicates.len() {
let predicate = &predicates[i];
if predicate.is_always_true()? {
predicates.remove(i);
} else if predicate.is_always_false()? {
return Ok(ConstantConditionEliminationResult::ImpossibleCondition);
} else {
i += 1;
}
}
}
Ok(ConstantConditionEliminationResult::Continue)
}
Operator::Nothing => Ok(ConstantConditionEliminationResult::Continue),
}
}
@@ -365,7 +354,6 @@ fn push_predicates(
Ok(())
}
Operator::SeekRowid { .. } => Ok(()),
Operator::Limit { source, .. } => {
push_predicates(source, referenced_tables)?;
Ok(())
@@ -379,6 +367,7 @@ fn push_predicates(
Ok(())
}
Operator::Scan { .. } => Ok(()),
Operator::Search { .. } => Ok(()),
Operator::Nothing => Ok(()),
}
}
@@ -424,6 +413,7 @@ fn push_predicate(
Ok(None)
}
Operator::Search { .. } => Ok(Some(predicate)),
Operator::Filter {
source,
predicates: ps,
@@ -486,7 +476,6 @@ fn push_predicate(
Ok(Some(push_result.unwrap()))
}
Operator::SeekRowid { .. } => Ok(Some(predicate)),
Operator::Limit { source, .. } => {
let push_result = push_predicate(source, predicate, referenced_tables)?;
if push_result.is_none() {
@@ -663,7 +652,6 @@ fn find_indexes_of_all_result_columns_in_operator_that_match_expr_either_fully_o
mask
}
Operator::Filter { .. } => 0,
Operator::SeekRowid { .. } => 0,
Operator::Limit { .. } => 0,
Operator::Join { .. } => 0,
Operator::Order { .. } => 0,
@@ -684,6 +672,7 @@ fn find_indexes_of_all_result_columns_in_operator_that_match_expr_either_fully_o
mask
}
Operator::Scan { .. } => 0,
Operator::Search { .. } => 0,
Operator::Nothing => 0,
};
@@ -847,7 +836,6 @@ fn find_shared_expressions_in_child_operators_and_mark_them_so_that_the_parent_o
)
}
Operator::Filter { .. } => unreachable!(),
Operator::SeekRowid { .. } => {}
Operator::Limit { source, .. } => {
find_shared_expressions_in_child_operators_and_mark_them_so_that_the_parent_operator_doesnt_recompute_them(source, expr_result_cache)
}
@@ -883,6 +871,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::Search { .. } => {}
Operator::Nothing => {}
}
}
@@ -910,59 +899,87 @@ pub trait Optimizable {
.check_constant()?
.map_or(false, |c| c == ConstantPredicate::AlwaysFalse))
}
// if the expression is the primary key of a table, returns the index of the table
fn check_primary_key(
&self,
referenced_tables: &[(Rc<BTreeTable>, String)],
fn is_primary_key_of(&self, table: &(Rc<BTreeTable>, String)) -> bool;
fn check_index_scan(
&mut self,
table: &(Rc<BTreeTable>, String),
available_indexes: &[Rc<Index>],
) -> Result<Option<usize>>;
}
impl Optimizable for ast::Expr {
fn check_primary_key(
&self,
referenced_tables: &[(Rc<BTreeTable>, String)],
) -> Result<Option<usize>> {
fn is_primary_key_of(&self, table: &(Rc<BTreeTable>, String)) -> bool {
match self {
ast::Expr::Id(ident) => {
let ident = normalize_ident(&ident.0);
let tables = referenced_tables
.iter()
.enumerate()
.filter_map(|(i, (t, _))| {
if t.get_column(&ident).map_or(false, |(_, c)| c.primary_key) {
Some(i)
} else {
None
}
});
let mut matches = 0;
let mut matching_tbl = None;
for tbl in tables {
matching_tbl = Some(tbl);
matches += 1;
if matches > 1 {
crate::bail_parse_error!("ambiguous column name {}", ident)
}
}
Ok(matching_tbl)
table
.0
.get_column(&ident)
.map_or(false, |(_, c)| c.primary_key)
}
ast::Expr::Qualified(tbl, ident) => {
let tbl = normalize_ident(&tbl.0);
let ident = normalize_ident(&ident.0);
let table = referenced_tables.iter().enumerate().find(|(_, (t, t_id))| {
*t_id == tbl && t.get_column(&ident).map_or(false, |(_, c)| c.primary_key)
});
if table.is_none() {
tbl == table.1
&& table
.0
.get_column(&ident)
.map_or(false, |(_, c)| c.primary_key)
}
_ => false,
}
}
fn check_index_scan(
&mut self,
table: &(Rc<BTreeTable>, String),
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.table_name == table.1 && i.columns.iter().any(|c| c.name == ident)
})
.collect::<Vec<_>>();
if indexes.is_empty() {
return Ok(None);
}
let table = table.unwrap();
Ok(Some(table.0))
if indexes.len() > 1 {
crate::bail_parse_error!("ambiguous column name {}", ident)
}
Ok(Some(indexes.first().unwrap().0))
}
ast::Expr::Qualified(_, ident) => {
let ident = normalize_ident(&ident.0);
let index = available_indexes.iter().enumerate().find(|(_, i)| {
if i.table_name != table.0.name {
return false;
}
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(table, available_indexes)?;
if lhs_index.is_some() {
return Ok(lhs_index);
}
let rhs_index = rhs.check_index_scan(table, 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),
}
@@ -1086,28 +1103,91 @@ impl Optimizable for ast::Expr {
}
}
pub fn try_extract_rowid_comparison_expression(
pub enum Either<T, U> {
Left(T),
Right(U),
}
pub fn try_extract_index_search_expression(
expr: ast::Expr,
table_index: usize,
referenced_tables: &[(Rc<BTreeTable>, String)],
) -> Result<(bool, ast::Expr)> {
table: &(Rc<BTreeTable>, String),
available_indexes: &[Rc<Index>],
) -> Result<Either<ast::Expr, Search>> {
match expr {
ast::Expr::Binary(lhs, ast::Operator::Equals, rhs) => {
if let Some(lhs_table_index) = lhs.check_primary_key(referenced_tables)? {
if lhs_table_index == table_index {
return Ok((true, *rhs));
ast::Expr::Binary(mut lhs, operator, mut rhs) => {
if lhs.is_primary_key_of(table) {
match operator {
ast::Operator::Equals => {
return Ok(Either::Right(Search::PrimaryKeyEq { cmp_expr: *rhs }));
}
ast::Operator::Greater
| ast::Operator::GreaterEquals
| ast::Operator::Less
| ast::Operator::LessEquals => {
return Ok(Either::Right(Search::PrimaryKeySearch {
cmp_op: operator,
cmp_expr: *rhs,
}));
}
_ => {}
}
}
if let Some(rhs_table_index) = rhs.check_primary_key(referenced_tables)? {
if rhs_table_index == table_index {
return Ok((true, *lhs));
if rhs.is_primary_key_of(table) {
match operator {
ast::Operator::Equals => {
return Ok(Either::Right(Search::PrimaryKeyEq { cmp_expr: *lhs }));
}
ast::Operator::Greater
| ast::Operator::GreaterEquals
| ast::Operator::Less
| ast::Operator::LessEquals => {
return Ok(Either::Right(Search::PrimaryKeySearch {
cmp_op: operator,
cmp_expr: *lhs,
}));
}
_ => {}
}
}
Ok((false, ast::Expr::Binary(lhs, ast::Operator::Equals, rhs)))
if let Some(index_index) = lhs.check_index_scan(table, available_indexes)? {
match operator {
ast::Operator::Equals
| ast::Operator::Greater
| ast::Operator::GreaterEquals
| ast::Operator::Less
| ast::Operator::LessEquals => {
return Ok(Either::Right(Search::IndexSearch {
index: available_indexes[index_index].clone(),
cmp_op: operator,
cmp_expr: *rhs,
}));
}
_ => {}
}
}
if let Some(index_index) = rhs.check_index_scan(table, available_indexes)? {
match operator {
ast::Operator::Equals
| ast::Operator::Greater
| ast::Operator::GreaterEquals
| ast::Operator::Less
| ast::Operator::LessEquals => {
return Ok(Either::Right(Search::IndexSearch {
index: available_indexes[index_index].clone(),
cmp_op: operator,
cmp_expr: *lhs,
}));
}
_ => {}
}
}
Ok(Either::Left(ast::Expr::Binary(lhs, operator, rhs)))
}
_ => Ok((false, expr)),
_ => Ok(Either::Left(expr)),
}
}

126
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 {
@@ -58,19 +64,6 @@ pub enum Operator {
source: Box<Operator>,
predicates: Vec<ast::Expr>,
},
// SeekRowid operator
// This operator is used to retrieve a single row from a table by its rowid.
// rowid_predicate is an expression that produces the comparison value for the rowid.
// e.g. rowid = 5, or rowid = other_table.foo
// predicates is an optional list of additional predicates to evaluate.
SeekRowid {
id: usize,
table: Rc<BTreeTable>,
table_identifier: String,
rowid_predicate: ast::Expr,
predicates: Option<Vec<ast::Expr>>,
step: usize,
},
// Limit operator
// This operator is used to limit the number of rows returned by the source operator.
Limit {
@@ -123,12 +116,43 @@ pub enum Operator {
predicates: Option<Vec<ast::Expr>>,
step: usize,
},
// Search operator
// This operator is used to search for a row in a table using an index
// (i.e. a primary key or a secondary index)
Search {
id: usize,
table: Rc<BTreeTable>,
table_identifier: String,
search: Search,
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.
Nothing,
}
/// An enum that represents a search operation that can be used to search for a row in a table using an index
/// (i.e. a primary key or a secondary index)
#[derive(Clone, Debug)]
pub enum Search {
/// A primary key equality search. This is a special case of the primary key search
/// that uses the SeekRowid bytecode instruction.
PrimaryKeyEq { cmp_expr: ast::Expr },
/// A primary key search. Uses bytecode instructions like SeekGT, SeekGE etc.
PrimaryKeySearch {
cmp_op: ast::Operator,
cmp_expr: ast::Expr,
},
/// A secondary index search. Uses bytecode instructions like SeekGE, SeekGT etc.
IndexSearch {
index: Rc<Index>,
cmp_op: ast::Operator,
cmp_expr: ast::Expr,
},
}
#[derive(Clone, Debug)]
pub enum ProjectionColumn {
Column(ast::Expr),
@@ -161,7 +185,6 @@ impl Operator {
..
} => aggregates.len() + group_by.as_ref().map_or(0, |g| g.len()),
Operator::Filter { source, .. } => source.column_count(referenced_tables),
Operator::SeekRowid { table, .. } => table.columns.len(),
Operator::Limit { source, .. } => source.column_count(referenced_tables),
Operator::Join { left, right, .. } => {
left.column_count(referenced_tables) + right.column_count(referenced_tables)
@@ -172,6 +195,7 @@ impl Operator {
.map(|e| e.column_count(referenced_tables))
.sum(),
Operator::Scan { table, .. } => table.columns.len(),
Operator::Search { table, .. } => table.columns.len(),
Operator::Nothing => 0,
}
}
@@ -203,9 +227,6 @@ impl Operator {
names
}
Operator::Filter { source, .. } => source.column_names(),
Operator::SeekRowid { table, .. } => {
table.columns.iter().map(|c| c.name.clone()).collect()
}
Operator::Limit { source, .. } => source.column_names(),
Operator::Join { left, right, .. } => {
let mut names = left.column_names();
@@ -226,6 +247,9 @@ impl Operator {
})
.collect(),
Operator::Scan { table, .. } => table.columns.iter().map(|c| c.name.clone()).collect(),
Operator::Search { table, .. } => {
table.columns.iter().map(|c| c.name.clone()).collect()
}
Operator::Nothing => vec![],
}
}
@@ -234,12 +258,12 @@ impl Operator {
match self {
Operator::Aggregate { id, .. } => *id,
Operator::Filter { id, .. } => *id,
Operator::SeekRowid { id, .. } => *id,
Operator::Limit { id, .. } => *id,
Operator::Join { id, .. } => *id,
Operator::Order { id, .. } => *id,
Operator::Projection { id, .. } => *id,
Operator::Scan { id, .. } => *id,
Operator::Search { id, .. } => *id,
Operator::Nothing => unreachable!(),
}
}
@@ -322,33 +346,6 @@ impl Display for Operator {
writeln!(f, "{}FILTER {}", indent, predicates_string)?;
fmt_operator(source, f, level + 1, true)
}
Operator::SeekRowid {
table,
rowid_predicate,
predicates,
..
} => {
match predicates {
Some(ps) => {
let predicates_string = ps
.iter()
.map(|p| p.to_string())
.collect::<Vec<String>>()
.join(" AND ");
writeln!(
f,
"{}SEEK {}.rowid ON rowid={} FILTER {}",
indent, &table.name, rowid_predicate, predicates_string
)?;
}
None => writeln!(
f,
"{}SEEK {}.rowid ON rowid={}",
indent, &table.name, rowid_predicate
)?,
}
Ok(())
}
Operator::Limit { source, limit, .. } => {
writeln!(f, "{}TAKE {}", indent, limit)?;
fmt_operator(source, f, level + 1, true)
@@ -429,6 +426,29 @@ impl Display for Operator {
}?;
Ok(())
}
Operator::Search {
table_identifier,
search,
..
} => {
match search {
Search::PrimaryKeyEq { .. } | Search::PrimaryKeySearch { .. } => {
writeln!(
f,
"{}SEARCH {} USING INTEGER PRIMARY KEY (rowid=?)",
indent, table_identifier
)?;
}
Search::IndexSearch { index, .. } => {
writeln!(
f,
"{}SEARCH {} USING INDEX {}",
indent, table_identifier, index.name
)?;
}
}
Ok(())
}
Operator::Nothing => Ok(()),
}
}
@@ -462,13 +482,6 @@ pub fn get_table_ref_bitmask_for_operator<'a>(
table_refs_mask |= get_table_ref_bitmask_for_ast_expr(tables, predicate)?;
}
}
Operator::SeekRowid { table, .. } => {
table_refs_mask |= 1
<< tables
.iter()
.position(|(t, _)| Rc::ptr_eq(t, table))
.unwrap();
}
Operator::Limit { source, .. } => {
table_refs_mask |= get_table_ref_bitmask_for_operator(tables, source)?;
}
@@ -489,6 +502,13 @@ pub fn get_table_ref_bitmask_for_operator<'a>(
.position(|(t, _)| Rc::ptr_eq(t, table))
.unwrap();
}
Operator::Search { table, .. } => {
table_refs_mask |= 1
<< tables
.iter()
.position(|(t, _)| Rc::ptr_eq(t, table))
.unwrap();
}
Operator::Nothing => {}
}
Ok(table_refs_mask)

7
core/translate/planner.rs
View File

@@ -281,6 +281,13 @@ pub fn prepare_select_plan<'a>(schema: &Schema, select: ast::Select) -> Result<P
Ok(Plan {
root_operator: operator,
referenced_tables,
available_indexes: schema
.indexes
.clone()
.into_iter()
.map(|(_, v)| v)
.flatten()
.collect(),
})
}
_ => todo!(),

15
core/types.rs
View File

@@ -409,13 +409,26 @@ pub enum CursorResult<T> {
IO,
}
#[derive(Clone, PartialEq, Debug)]
pub enum SeekOp {
EQ,
GE,
GT,
}
#[derive(Clone, PartialEq, Debug)]
pub enum SeekKey<'a> {
TableRowId(u64),
IndexKey(&'a OwnedRecord),
}
pub trait Cursor {
fn is_empty(&self) -> bool;
fn rewind(&mut self) -> Result<CursorResult<()>>;
fn next(&mut self) -> Result<CursorResult<()>>;
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(&mut self, key: SeekKey, op: SeekOp) -> 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);
}

70
core/vdbe/explain.rs
View File

@@ -521,6 +521,76 @@ 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 {
is_index,
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 {
is_index,
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,

297
core/vdbe/mod.rs
View File

@@ -30,7 +30,9 @@ use crate::pseudo::PseudoCursor;
use crate::schema::Table;
use crate::storage::sqlite3_ondisk::DatabaseHeader;
use crate::storage::{btree::BTreeCursor, pager::Pager};
use crate::types::{AggContext, Cursor, CursorResult, OwnedRecord, OwnedValue, Record};
use crate::types::{
AggContext, Cursor, CursorResult, OwnedRecord, OwnedValue, Record, SeekKey, SeekOp,
};
use crate::{Result, DATABASE_VERSION};
use datetime::{exec_date, exec_time, exec_unixepoch};
@@ -298,6 +300,53 @@ 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,
},
// If cursor_id refers to an SQL table (B-Tree that uses integer keys), use the value in start_reg as the key.
// If cursor_id refers to an SQL index, then start_reg is the first in an array of num_regs registers that are used as an unpacked index key.
// 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 {
is_index: bool,
cursor_id: CursorID,
start_reg: usize,
num_regs: usize,
target_pc: BranchOffset,
},
// If cursor_id refers to an SQL table (B-Tree that uses integer keys), use the value in start_reg as the key.
// If cursor_id refers to an SQL index, then start_reg is the first in an array of num_regs registers that are used as an unpacked index key.
// 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 {
is_index: bool,
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 +493,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 +508,7 @@ impl ProgramState {
cursors,
registers,
last_compare: None,
deferred_seek: None,
ended_coroutine: false,
regex_cache: RegexCache::new(),
}
@@ -958,6 +1009,19 @@ 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();
match table_cursor.seek(SeekKey::TableRowId(rowid.unwrap()), SeekOp::EQ)? {
CursorResult::Ok(_) => {}
CursorResult::IO => {
state.deferred_seek = Some((index_cursor_id, table_cursor_id));
return Ok(StepResult::IO);
}
}
}
let cursor = cursors.get_mut(cursor_id).unwrap();
if let Some(ref record) = *cursor.record()? {
let null_flag = cursor.get_null_flag();
@@ -1085,6 +1149,19 @@ 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();
match table_cursor.seek(SeekKey::TableRowId(rowid.unwrap()), SeekOp::EQ)? {
CursorResult::Ok(_) => {}
CursorResult::IO => {
state.deferred_seek = Some((index_cursor_id, table_cursor_id));
return Ok(StepResult::IO);
}
}
}
let cursor = cursors.get_mut(cursor_id).unwrap();
if let Some(ref rowid) = cursor.rowid()? {
state.registers[*dest] = OwnedValue::Integer(*rowid as i64);
@@ -1101,13 +1178,17 @@ impl Program {
let cursor = cursors.get_mut(cursor_id).unwrap();
let rowid = match &state.registers[*src_reg] {
OwnedValue::Integer(rowid) => *rowid as u64,
_ => {
OwnedValue::Null => {
state.pc = *target_pc;
continue;
}
other => {
return Err(LimboError::InternalError(
"SeekRowid: the value in the register is not an integer".into(),
format!("SeekRowid: the value in the register is not an integer or NULL: {}", other)
));
}
};
match cursor.seek_rowid(rowid)? {
match cursor.seek(SeekKey::TableRowId(rowid), SeekOp::EQ)? {
CursorResult::Ok(found) => {
if !found {
state.pc = *target_pc;
@@ -1121,6 +1202,180 @@ 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,
is_index,
} => {
if *is_index {
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(SeekKey::IndexKey(&record_from_regs), SeekOp::GE)? {
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);
}
}
} else {
let cursor = cursors.get_mut(cursor_id).unwrap();
let rowid = match &state.registers[*start_reg] {
OwnedValue::Null => {
// All integer values are greater than null so we just rewind the cursor
match cursor.rewind()? {
CursorResult::Ok(()) => {}
CursorResult::IO => {
// If there is I/O, the instruction is restarted.
return Ok(StepResult::IO);
}
}
state.pc += 1;
continue;
}
OwnedValue::Integer(rowid) => *rowid as u64,
_ => {
return Err(LimboError::InternalError(
"SeekGE: the value in the register is not an integer".into(),
));
}
};
match cursor.seek(SeekKey::TableRowId(rowid), SeekOp::GE)? {
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,
is_index,
} => {
if *is_index {
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(SeekKey::IndexKey(&record_from_regs), SeekOp::GT)? {
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);
}
}
} else {
let cursor = cursors.get_mut(cursor_id).unwrap();
let rowid = match &state.registers[*start_reg] {
OwnedValue::Null => {
// All integer values are greater than null so we just rewind the cursor
match cursor.rewind()? {
CursorResult::Ok(()) => {}
CursorResult::IO => {
// If there is I/O, the instruction is restarted.
return Ok(StepResult::IO);
}
}
state.pc += 1;
continue;
}
OwnedValue::Integer(rowid) => *rowid as u64,
_ => {
return Err(LimboError::InternalError(
"SeekGT: the value in the register is not an integer".into(),
));
}
};
match cursor.seek(SeekKey::TableRowId(rowid), SeekOp::GT)? {
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] {
@@ -1805,7 +2060,7 @@ fn get_new_rowid<R: Rng>(cursor: &mut Box<dyn Cursor>, mut rng: R) -> Result<Cur
let max_attempts = 100;
for count in 0..max_attempts {
rowid = distribution.sample(&mut rng).try_into().unwrap();
match cursor.seek_rowid(rowid)? {
match cursor.seek(SeekKey::TableRowId(rowid), SeekOp::EQ)? {
CursorResult::Ok(false) => break, // Found a non-existing rowid
CursorResult::Ok(true) => {
if count == max_attempts - 1 {
@@ -1869,7 +2124,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 {
@@ -2380,6 +2638,8 @@ fn execute_sqlite_version(version_integer: i64) -> String {
#[cfg(test)]
mod tests {
use crate::types::{SeekKey, SeekOp};
use super::{
exec_abs, exec_char, exec_hex, exec_if, exec_instr, exec_length, exec_like, exec_lower,
exec_ltrim, exec_max, exec_min, exec_nullif, exec_quote, exec_random, exec_randomblob,
@@ -2394,6 +2654,7 @@ mod tests {
mock! {
Cursor {
fn seek_to_last(&mut self) -> Result<CursorResult<()>>;
fn seek<'a>(&mut self, key: SeekKey<'a>, op: SeekOp) -> Result<CursorResult<bool>>;
fn rowid(&self) -> Result<Option<u64>>;
fn seek_rowid(&mut self, rowid: u64) -> Result<CursorResult<bool>>;
}
@@ -2408,8 +2669,8 @@ mod tests {
self.rowid()
}
fn seek_rowid(&mut self, rowid: u64) -> Result<CursorResult<bool>> {
self.seek_rowid(rowid)
fn seek(&mut self, key: SeekKey<'_>, op: SeekOp) -> Result<CursorResult<bool>> {
self.seek(key, op)
}
fn rewind(&mut self) -> Result<CursorResult<()>> {
@@ -2484,10 +2745,10 @@ mod tests {
.return_once(|| Ok(CursorResult::Ok(())));
mock.expect_rowid()
.return_once(|| Ok(Some(std::i64::MAX as u64)));
mock.expect_seek_rowid()
.with(predicate::always())
.returning(|rowid| {
if rowid == 50 {
mock.expect_seek()
.with(predicate::always(), predicate::always())
.returning(|rowid, _| {
if rowid == SeekKey::TableRowId(50) {
Ok(CursorResult::Ok(false))
} else {
Ok(CursorResult::Ok(true))
@@ -2505,9 +2766,9 @@ mod tests {
.return_once(|| Ok(CursorResult::Ok(())));
mock.expect_rowid()
.return_once(|| Ok(Some(std::i64::MAX as u64)));
mock.expect_seek_rowid()
.with(predicate::always())
.return_once(|_| Ok(CursorResult::IO));
mock.expect_seek()
.with(predicate::always(), predicate::always())
.return_once(|_, _| Ok(CursorResult::IO));
let result = get_new_rowid(&mut (Box::new(mock) as Box<dyn Cursor>), thread_rng());
assert!(matches!(result, Ok(CursorResult::IO)));
@@ -2518,9 +2779,9 @@ mod tests {
.return_once(|| Ok(CursorResult::Ok(())));
mock.expect_rowid()
.return_once(|| Ok(Some(std::i64::MAX as u64)));
mock.expect_seek_rowid()
.with(predicate::always())
.returning(|_| Ok(CursorResult::Ok(true)));
mock.expect_seek()
.with(predicate::always(), predicate::always())
.returning(|_, _| Ok(CursorResult::Ok(true)));
// Mock the random number generation
let result = get_new_rowid(&mut (Box::new(mock) as Box<dyn Cursor>), StepRng::new(1, 1));

4
core/vdbe/sorter.rs
View File

@@ -1,5 +1,5 @@
use crate::{
types::{Cursor, CursorResult, OwnedRecord, OwnedValue},
types::{Cursor, CursorResult, OwnedRecord, OwnedValue, SeekKey, SeekOp},
Result,
};
use std::cell::{Ref, RefCell};
@@ -49,7 +49,7 @@ impl Cursor for Sorter {
todo!();
}
fn seek_rowid(&mut self, _: u64) -> Result<CursorResult<bool>> {
fn seek(&mut self, _: SeekKey<'_>, _: SeekOp) -> Result<CursorResult<bool>> {
unimplemented!();
}

4
testing/agg-functions.test
View File

@@ -28,8 +28,8 @@ do_execsql_test select-total-text {
} {0.0}
do_execsql_test select-limit {
SELECT id FROM users LIMIT 1;
} {1}
SELECT typeof(id) FROM users LIMIT 1;
} {integer}
do_execsql_test select-count {
SELECT count(id) FROM users;

10
testing/cmdlineshell.test
View File

@@ -5,7 +5,7 @@ source $testdir/tester.tcl
do_execsql_test schema {
.schema
} {{CREATE TABLE users (
} {"CREATE TABLE users (
id INTEGER PRIMARY KEY,
first_name TEXT,
last_name TEXT,
@@ -21,11 +21,12 @@ CREATE TABLE products (
id INTEGER PRIMARY KEY,
name TEXT,
price REAL
);}}
);
CREATE INDEX age_idx on users (age);"}
do_execsql_test schema-1 {
.schema users
} {{CREATE TABLE users (
} {"CREATE TABLE users (
id INTEGER PRIMARY KEY,
first_name TEXT,
last_name TEXT,
@@ -36,7 +37,8 @@ do_execsql_test schema-1 {
state TEXT,
zipcode TEXT,
age INTEGER
);}}
);
CREATE INDEX age_idx on users (age);"}
do_execsql_test schema-2 {
.schema products

6
testing/join.test
View File

@@ -200,3 +200,9 @@ Jamie||Edward}
do_execsql_test left-join-constant-condition-true-inner-join-constant-condition-false {
select u.first_name, p.name, u2.first_name from users u left join products as p on 1 join users u2 on 0 limit 5;
} {}
do_execsql_test join-utilizing-both-seekrowid-and-secondary-index {
select u.first_name, p.name from users u join products p on u.id = p.id and u.age > 70;
} {Matthew|boots
Nicholas|shorts
Jamie|hat}

4
testing/select.test
View File

@@ -44,8 +44,8 @@ do_execsql_test table-star {
} {1|hat|79.0|hat}
do_execsql_test table-star-2 {
select p.*, u.age from users u join products p limit 1;
} {1|hat|79.0|94}
select p.*, u.first_name from users u join products p on u.id = p.id limit 1;
} {1|hat|79.0|Jamie}
do_execsql_test seekrowid {
select * from users u where u.id = 5;

BIN
testing/testing.db
View File

Binary file not shown.

58
testing/where.test
View File

@@ -92,24 +92,25 @@ do_execsql_test where-clause-no-table-constant-condition-false-7 {
select 1 where 'hamburger';
} {}
# this test functions as an assertion that the index on users.age is being used, since the results are ordered by age without an order by.
do_execsql_test select-where-and {
select first_name, age from users where first_name = 'Jamie' and age > 80
} {Jamie|94
} {Jamie|87
Jamie|88
Jamie|88
Jamie|99
Jamie|92
Jamie|87
Jamie|88
Jamie|94
Jamie|99
}
do_execsql_test select-where-or {
select first_name, age from users where first_name = 'Jamie' and age > 80
} {Jamie|94
} {Jamie|87
Jamie|88
Jamie|88
Jamie|99
Jamie|92
Jamie|87
Jamie|88
Jamie|94
Jamie|99
}
do_execsql_test select-where-and-or {
@@ -267,3 +268,44 @@ do_execsql_test where-complex-parentheses {
select id, name from products where ((id = 5 and name = 'sweatshirt') or (id = 1 and name = 'hat')) and (name = 'sweatshirt' or name = 'hat') ORDER BY id;
} {1|hat
5|sweatshirt}
# regression test for primary key index behavior
do_execsql_test where-id-index-seek-regression-test {
select id from users where id > 9995;
} {9996
9997
9998
9999
10000}
do_execsql_test where-id-index-seek-regression-test-2 {
select count(1) from users where id > 0;
} {10000}
# regression test for secondary index (users.age) behavior
do_execsql_test where-age-index-seek-regression-test {
select age from users where age >= 100 limit 20;
} {100
100
100
100
100
100
100
100
100
100
100
100
100
100
100
100
100
100
100
100}
do_execsql_test where-age-index-seek-regression-test-2 {
select count(1) from users where age > 0;
} {10000}