mirrors/limbo
1
0
Fork 0
mirror of https://github.com/tursodatabase/limbo.git synced 2025-08-03 17:48:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Merge 'Support backwards index scan and seeks + utilize indexes in removing ORDER BY' from Jussi Saurio

Browse source 
## Main stuff
- Support iterating an index backwards
- Support scanning an index (instead of seeking with a condition)
- Support backwards index seeks
- Support backwards rowid seeks
- Fix existing backwards iteration logic for table btrees
- Remove ORDER BY entirely if any index satisfies the ordering
- Add fuzz tests for rowid seeks, 1 and 2 column index seeks
## Bytecode examples (note the lack of order by sorting):
one column index order by, forwards:
```sql
limbo> explain select first_name from users order by age;
addr  opcode             p1    p2    p3    p4             p5  comment
----  -----------------  ----  ----  ----  -------------  --  -------
0     Init               0     13    0                    0   Start at 13
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     RewindAsync        1     0     0                    0
6     RewindAwait        1     12    0                    0   Rewind table age_idx
7       DeferredSeek     1     0     0                    0
8       Column           0     1     1                    0   r[1]=users.first_name
9       ResultRow        1     1     0                    0   output=r[1]
10    NextAsync          1     0     0                    0
11    NextAwait          1     7     0                    0
12    Halt               0     0     0                    0
13    Transaction        0     0     0                    0   write=false
14    Goto               0     1     0                    0
```
one column index order by, backwards:
```sql
limbo> explain select first_name from users order by age desc;
addr  opcode             p1    p2    p3    p4             p5  comment
----  -----------------  ----  ----  ----  -------------  --  -------
0     Init               0     13    0                    0   Start at 13
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     LastAsync          1     0     0                    0
6     LastAwait          1     0     0                    0
7       DeferredSeek     1     0     0                    0
8       Column           0     1     1                    0   r[1]=users.first_name
9       ResultRow        1     1     0                    0   output=r[1]
10    PrevAsync          1     0     0                    0
11    PrevAwait          1     0     0                    0
12    Halt               0     0     0                    0
13    Transaction        0     0     0                    0   write=false
14    Goto               0     1     0                    0
```
rowid seek, backwards:
```sql
limbo> explain select * from users where id < 100 order by id desc;
addr  opcode             p1    p2    p3    p4             p5  comment
----  -----------------  ----  ----  ----  -------------  --  -------
0     Init               0     19    0                    0   Start at 19
1     OpenReadAsync      0     2     0                    0   table=users, root=2
2     OpenReadAwait      0     0     0                    0
3     Integer            100   11    0                    0   r[11]=100
4     SeekLT             0     18    11                   0
5       RowId            0     1     0                    0   r[1]=users.rowid
6       Column           0     1     2                    0   r[2]=users.first_name
7       Column           0     2     3                    0   r[3]=users.last_name
8       Column           0     3     4                    0   r[4]=users.email
9       Column           0     4     5                    0   r[5]=users.phone_number
10      Column           0     5     6                    0   r[6]=users.address
11      Column           0     6     7                    0   r[7]=users.city
12      Column           0     7     8                    0   r[8]=users.state
13      Column           0     8     9                    0   r[9]=users.zipcode
14      Column           0     9     10                   0   r[10]=users.age
15      ResultRow        1     10    0                    0   output=r[1..10]
16    PrevAsync          0     0     0                    0
17    PrevAwait          0     0     0                    0
18    Halt               0     0     0                    0
19    Transaction        0     0     0                    0   write=false
20    Goto               0     1     0                    0
```
two column order by, setup:
```sql
cargo run dualpk.db

Limbo v0.0.18-pre.3
Enter ".help" for usage hints.
limbo> .schema
CREATE TABLE a(b,c,d,e, primary key (d,c));
```
two column order by, forwards:
```sql
limbo> explain select * from a order by d,c;
addr  opcode             p1    p2    p3    p4             p5  comment
----  -----------------  ----  ----  ----  -------------  --  -------
0     Init               0     16    0                    0   Start at 16
1     OpenReadAsync      0     2     0                    0   table=a, root=2
2     OpenReadAwait      0     0     0                    0
3     OpenReadAsync      1     3     0                    0   table=sqlite_autoindex_a_1, root=3
4     OpenReadAwait      0     0     0                    0
5     RewindAsync        1     0     0                    0
6     RewindAwait        1     15    0                    0   Rewind table sqlite_autoindex_a_1
7       DeferredSeek     1     0     0                    0
8       Column           0     0     1                    0   r[1]=a.b
9       Column           0     1     2                    0   r[2]=a.c
10      Column           0     2     3                    0   r[3]=a.d
11      Column           0     3     4                    0   r[4]=a.e
12      ResultRow        1     4     0                    0   output=r[1..4]
13    NextAsync          1     0     0                    0
14    NextAwait          1     7     0                    0
15    Halt               0     0     0                    0
16    Transaction        0     0     0                    0   write=false
17    Goto               0     1     0                    0
```
two column order by, forwards with index seek:
```sql
limbo> explain select * from a where d > 100 order by d,c;
addr  opcode             p1    p2    p3    p4             p5  comment
----  -----------------  ----  ----  ----  -------------  --  -------
0     Init               0     16    0                    0   Start at 16
1     OpenReadAsync      0     2     0                    0   table=a, root=2
2     OpenReadAwait      0     0     0                    0
3     OpenReadAsync      1     3     0                    0   table=sqlite_autoindex_a_1, root=3
4     OpenReadAwait      0     0     0                    0
5     Integer            100   5     0                    0   r[5]=100
6     SeekGT             1     15    5                    0
7       DeferredSeek     1     0     0                    0
8       Column           0     0     1                    0   r[1]=a.b
9       Column           0     1     2                    0   r[2]=a.c
10      Column           0     2     3                    0   r[3]=a.d
11      Column           0     3     4                    0   r[4]=a.e
12      ResultRow        1     4     0                    0   output=r[1..4]
13    NextAsync          1     0     0                    0
14    NextAwait          1     7     0                    0
15    Halt               0     0     0                    0
16    Transaction        0     0     0                    0   write=false
17    Goto               0     1     0                    0
```
two column order by, forwards with index scan and termination condition:
```sql
limbo> explain select * from a where d < 100 order by d,c;
addr  opcode             p1    p2    p3    p4             p5  comment
----  -----------------  ----  ----  ----  -------------  --  -------
0     Init               0     18    0                    0   Start at 18
1     OpenReadAsync      0     2     0                    0   table=a, root=2
2     OpenReadAwait      0     0     0                    0
3     OpenReadAsync      1     3     0                    0   table=sqlite_autoindex_a_1, root=3
4     OpenReadAwait      0     0     0                    0
5     Null               0     5     0                    0   r[5]=NULL
6     SeekGT             1     17    5                    0
7       Integer          100   6     0                    0   r[6]=100
8       IdxGE            1     17    6                    0
9       DeferredSeek     1     0     0                    0
10      Column           0     0     1                    0   r[1]=a.b
11      Column           0     1     2                    0   r[2]=a.c
12      Column           0     2     3                    0   r[3]=a.d
13      Column           0     3     4                    0   r[4]=a.e
14      ResultRow        1     4     0                    0   output=r[1..4]
15    NextAsync          1     0     0                    0
16    NextAwait          1     7     0                    0
17    Halt               0     0     0                    0
18    Transaction        0     0     0                    0   write=false
19    Goto               0     1     0                    0
```
two column order by, backwards:
```sql
limbo> explain select * from a order by d desc,c desc;
addr  opcode             p1    p2    p3    p4             p5  comment
----  -----------------  ----  ----  ----  -------------  --  -------
0     Init               0     16    0                    0   Start at 16
1     OpenReadAsync      0     2     0                    0   table=a, root=2
2     OpenReadAwait      0     0     0                    0
3     OpenReadAsync      1     3     0                    0   table=sqlite_autoindex_a_1, root=3
4     OpenReadAwait      0     0     0                    0
5     LastAsync          1     0     0                    0
6     LastAwait          1     0     0                    0
7       DeferredSeek     1     0     0                    0
8       Column           0     0     1                    0   r[1]=a.b
9       Column           0     1     2                    0   r[2]=a.c
10      Column           0     2     3                    0   r[3]=a.d
11      Column           0     3     4                    0   r[4]=a.e
12      ResultRow        1     4     0                    0   output=r[1..4]
13    PrevAsync          1     0     0                    0
14    PrevAwait          1     0     0                    0
15    Halt               0     0     0                    0
16    Transaction        0     0     0                    0   write=false
17    Goto               0     1     0                    0
```
two column order by, backwards with index seek:
```sql
limbo> explain select * from a where d < 100 order by d desc,c desc;
addr  opcode             p1    p2    p3    p4             p5  comment
----  -----------------  ----  ----  ----  -------------  --  -------
0     Init               0     16    0                    0   Start at 16
1     OpenReadAsync      0     2     0                    0   table=a, root=2
2     OpenReadAwait      0     0     0                    0
3     OpenReadAsync      1     3     0                    0   table=sqlite_autoindex_a_1, root=3
4     OpenReadAwait      0     0     0                    0
5     Integer            100   5     0                    0   r[5]=100
6     SeekLT             1     15    5                    0
7       DeferredSeek     1     0     0                    0
8       Column           0     0     1                    0   r[1]=a.b
9       Column           0     1     2                    0   r[2]=a.c
10      Column           0     2     3                    0   r[3]=a.d
11      Column           0     3     4                    0   r[4]=a.e
12      ResultRow        1     4     0                    0   output=r[1..4]
13    PrevAsync          1     0     0                    0
14    PrevAwait          1     0     0                    0
15    Halt               0     0     0                    0
16    Transaction        0     0     0                    0   write=false
17    Goto               0     1     0                    0
```
two column order by, backwards with index scan and termination
condition:
```sql
limbo> explain select * from a where d > 100 order by d desc,c desc;
addr  opcode             p1    p2    p3    p4             p5  comment
----  -----------------  ----  ----  ----  -------------  --  -------
0     Init               0     18    0                    0   Start at 18
1     OpenReadAsync      0     2     0                    0   table=a, root=2
2     OpenReadAwait      0     0     0                    0
3     OpenReadAsync      1     3     0                    0   table=sqlite_autoindex_a_1, root=3
4     OpenReadAwait      0     0     0                    0
5     LastAsync          1     0     0                    0
6     LastAwait          1     0     0                    0
7       Integer          100   6     0                    0   r[6]=100
8       IdxLE            1     17    6                    0
9       DeferredSeek     1     0     0                    0
10      Column           0     0     1                    0   r[1]=a.b
11      Column           0     1     2                    0   r[2]=a.c
12      Column           0     2     3                    0   r[3]=a.d
13      Column           0     3     4                    0   r[4]=a.e
14      ResultRow        1     4     0                    0   output=r[1..4]
15    PrevAsync          1     0     0                    0
16    PrevAwait          1     0     0                    0
17    Halt               0     0     0                    0
18    Transaction        0     0     0                    0   write=false
19    Goto               0     1     0                    0
```

Reviewed-by: Preston Thorpe (@PThorpe92)

Closes #1209
This commit is contained in:
Jussi Saurio 2025-04-09 12:03:14 +03:00
commit aa6e2d853a
15 changed files with 1367 additions and 403 deletions
Download patch file Download diff file Expand all files Collapse all files

530
core/storage/btree.rs
View file

@ -4,6 +4,7 @@ use crate::storage::pager::Pager;
use crate::storage::sqlite3_ondisk::{
read_u32, read_varint, BTreeCell, PageContent, PageType, TableInteriorCell, TableLeafCell,
};
use crate::translate::plan::IterationDirection;
use crate::MvCursor;
use crate::types::{
@ -312,6 +313,17 @@ enum OverflowState {
Done,
}
/// Iteration state of the cursor. Can only be set once.
/// Once a SeekGT or SeekGE is performed, the cursor must iterate forwards and calling prev() is an error.
/// Similarly, once a SeekLT or SeekLE is performed, the cursor must iterate backwards and calling next() is an error.
/// When a SeekEQ or SeekRowid is performed, the cursor is NOT allowed to iterate further.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum IterationState {
Unset,
Iterating(IterationDirection),
IterationNotAllowed,
}
pub struct BTreeCursor {
/// The multi-version cursor that is used to read and write to the database file.
mv_cursor: Option<Rc<RefCell<MvCursor>>>,
@ -337,6 +349,8 @@ pub struct BTreeCursor {
/// Reusable immutable record, used to allow better allocation strategy.
reusable_immutable_record: RefCell<Option<ImmutableRecord>>,
empty_record: Cell<bool>,
pub iteration_state: IterationState,
}
/// Stack of pages representing the tree traversal order.
@ -385,6 +399,7 @@ impl BTreeCursor {
},
reusable_immutable_record: RefCell::new(None),
empty_record: Cell::new(true),
iteration_state: IterationState::Unset,
}
}
@ -404,7 +419,10 @@ impl BTreeCursor {
/// Move the cursor to the previous record and return it.
/// Used in backwards iteration.
fn get_prev_record(&mut self) -> Result<CursorResult<Option<u64>>> {
fn get_prev_record(
&mut self,
predicate: Option<(SeekKey<'_>, SeekOp)>,
) -> Result<CursorResult<Option<u64>>> {
loop {
let page = self.stack.top();
let cell_idx = self.stack.current_cell_index();
@ -413,11 +431,11 @@ impl BTreeCursor {
// todo: find a better way to flag moved to end or begin of page
if self.stack.current_cell_index_less_than_min() {
loop {
if self.stack.current_cell_index() > 0 {
self.stack.retreat();
if self.stack.current_cell_index() >= 0 {
break;
}
if self.stack.has_parent() {
self.going_upwards = true;
self.stack.pop();
} else {
// moved to begin of btree
@ -429,11 +447,6 @@ impl BTreeCursor {
}
let cell_idx = cell_idx as usize;
tracing::trace!(
"get_prev_record current id={} cell={}",
page.get().id,
cell_idx
);
return_if_locked!(page);
if !page.is_loaded() {
self.pager.load_page(page.clone())?;
@ -442,13 +455,24 @@ impl BTreeCursor {
let contents = page.get().contents.as_ref().unwrap();
let cell_count = contents.cell_count();
// If we are at the end of the page and we haven't just come back from the right child,
// we now need to move to the rightmost child.
if cell_idx as i32 == i32::MAX && !self.going_upwards {
let rightmost_pointer = contents.rightmost_pointer();
if let Some(rightmost_pointer) = rightmost_pointer {
self.stack
.push_backwards(self.pager.read_page(rightmost_pointer as usize)?);
continue;
}
}
let cell_idx = if cell_idx >= cell_count {
self.stack.set_cell_index(cell_count as i32 - 1);
cell_count - 1
} else {
cell_idx
};
let cell = contents.cell_get(
cell_idx,
payload_overflow_threshold_max(contents.page_type(), self.usable_space() as u16),
@ -462,9 +486,7 @@ impl BTreeCursor {
_rowid,
}) => {
let mem_page = self.pager.read_page(_left_child_page as usize)?;
self.stack.push(mem_page);
// use cell_index = i32::MAX to tell next loop to go to the end of the current page
self.stack.set_cell_index(i32::MAX);
self.stack.push_backwards(mem_page);
continue;
}
BTreeCell::TableLeafCell(TableLeafCell {
@ -484,8 +506,135 @@ impl BTreeCursor {
self.stack.retreat();
return Ok(CursorResult::Ok(Some(_rowid)));
}
BTreeCell::IndexInteriorCell(_) => todo!(),
BTreeCell::IndexLeafCell(_) => todo!(),
BTreeCell::IndexInteriorCell(IndexInteriorCell {
payload,
left_child_page,
first_overflow_page,
payload_size,
}) => {
if !self.going_upwards {
// In backwards iteration, if we haven't just moved to this interior node from the
// right child, but instead are about to move to the left child, we need to retreat
// so that we don't come back to this node again.
// For example:
// this parent: key 666
// left child has: key 663, key 664, key 665
// we need to move to the previous parent (with e.g. key 662) when iterating backwards.
self.stack.retreat();
let mem_page = self.pager.read_page(left_child_page as usize)?;
self.stack.push(mem_page);
// use cell_index = i32::MAX to tell next loop to go to the end of the current page
self.stack.set_cell_index(i32::MAX);
continue;
}
if let Some(next_page) = first_overflow_page {
return_if_io!(self.process_overflow_read(payload, next_page, payload_size))
} else {
crate::storage::sqlite3_ondisk::read_record(
payload,
self.get_immutable_record_or_create().as_mut().unwrap(),
)?
};
// Going upwards = we just moved to an interior cell from the right child.
// On the first pass we must take the record from the interior cell (since unlike table btrees, index interior cells have payloads)
// We then mark going_upwards=false so that we go back down the tree on the next invocation.
self.going_upwards = false;
if predicate.is_none() {
let rowid = match self.get_immutable_record().as_ref().unwrap().last_value()
{
Some(RefValue::Integer(rowid)) => *rowid as u64,
_ => unreachable!("index cells should have an integer rowid"),
};
return Ok(CursorResult::Ok(Some(rowid)));
}
let (key, op) = predicate.as_ref().unwrap();
let SeekKey::IndexKey(index_key) = key else {
unreachable!("index seek key should be a record");
};
let order = {
let record = self.get_immutable_record();
let record = record.as_ref().unwrap();
let record_values = record.get_values();
let record_slice_same_num_cols =
&record_values[..index_key.get_values().len()];
let order =
compare_immutable(record_slice_same_num_cols, index_key.get_values());
order
};
let found = match op {
SeekOp::EQ => order.is_eq(),
SeekOp::LE => order.is_le(),
SeekOp::LT => order.is_lt(),
_ => unreachable!("Seek GT/GE should not happen in get_prev_record() because we are iterating backwards"),
};
if found {
let rowid = match self.get_immutable_record().as_ref().unwrap().last_value()
{
Some(RefValue::Integer(rowid)) => *rowid as u64,
_ => unreachable!("index cells should have an integer rowid"),
};
return Ok(CursorResult::Ok(Some(rowid)));
} else {
continue;
}
}
BTreeCell::IndexLeafCell(IndexLeafCell {
payload,
first_overflow_page,
payload_size,
}) => {
if let Some(next_page) = first_overflow_page {
return_if_io!(self.process_overflow_read(payload, next_page, payload_size))
} else {
crate::storage::sqlite3_ondisk::read_record(
payload,
self.get_immutable_record_or_create().as_mut().unwrap(),
)?
};
self.stack.retreat();
if predicate.is_none() {
let rowid = match self.get_immutable_record().as_ref().unwrap().last_value()
{
Some(RefValue::Integer(rowid)) => *rowid as u64,
_ => unreachable!("index cells should have an integer rowid"),
};
return Ok(CursorResult::Ok(Some(rowid)));
}
let (key, op) = predicate.as_ref().unwrap();
let SeekKey::IndexKey(index_key) = key else {
unreachable!("index seek key should be a record");
};
let order = {
let record = self.get_immutable_record();
let record = record.as_ref().unwrap();
let record_values = record.get_values();
let record_slice_same_num_cols =
&record_values[..index_key.get_values().len()];
let order =
compare_immutable(record_slice_same_num_cols, index_key.get_values());
order
};
let found = match op {
SeekOp::EQ => order.is_eq(),
SeekOp::LE => order.is_le(),
SeekOp::LT => order.is_lt(),
_ => unreachable!("Seek GT/GE should not happen in get_prev_record() because we are iterating backwards"),
};
if found {
let rowid = match self.get_immutable_record().as_ref().unwrap().last_value()
{
Some(RefValue::Integer(rowid)) => *rowid as u64,
_ => unreachable!("index cells should have an integer rowid"),
};
return Ok(CursorResult::Ok(Some(rowid)));
} else {
continue;
}
}
}
}
}
@ -720,6 +869,7 @@ impl BTreeCursor {
SeekOp::GT => order.is_gt(),
SeekOp::GE => order.is_ge(),
SeekOp::EQ => order.is_eq(),
_ => unreachable!("Seek LE/LT should not happen in get_next_record() because we are iterating forwards"),
};
if found {
let rowid = match self.get_immutable_record().as_ref().unwrap().last_value()
@ -771,6 +921,7 @@ impl BTreeCursor {
SeekOp::GT => order.is_lt(),
SeekOp::GE => order.is_le(),
SeekOp::EQ => order.is_le(),
_ => todo!("not implemented: {:?}", op),
};
if found {
let rowid = match self.get_immutable_record().as_ref().unwrap().last_value()
@ -792,6 +943,35 @@ impl BTreeCursor {
/// or e.g. find the first record greater than the seek key in a range query (e.g. SELECT * FROM table WHERE col > 10).
/// We don't include the rowid in the comparison and that's why the last value from the record is not included.
fn do_seek(&mut self, key: SeekKey<'_>, op: SeekOp) -> Result<CursorResult<Option<u64>>> {
assert!(
self.iteration_state != IterationState::Unset,
"iteration state must have been set before do_seek() is called"
);
let valid_op = match (self.iteration_state, op) {
(IterationState::Iterating(IterationDirection::Forwards), SeekOp::GE | SeekOp::GT) => {
true
}
(IterationState::Iterating(IterationDirection::Backwards), SeekOp::LE | SeekOp::LT) => {
true
}
(IterationState::IterationNotAllowed, SeekOp::EQ) => true,
_ => false,
};
assert!(
valid_op,
"invalid seek op for iteration state: {:?} {:?}",
self.iteration_state, op
);
let cell_iter_dir = match self.iteration_state {
IterationState::Iterating(IterationDirection::Forwards)
| IterationState::IterationNotAllowed => IterationDirection::Forwards,
IterationState::Iterating(IterationDirection::Backwards) => {
IterationDirection::Backwards
}
IterationState::Unset => {
unreachable!("iteration state must have been set before do_seek() is called");
}
};
return_if_io!(self.move_to(key.clone(), op.clone()));
{
@ -800,9 +980,27 @@ impl BTreeCursor {
let contents = page.get().contents.as_ref().unwrap();
for cell_idx in 0..contents.cell_count() {
let cell_count = contents.cell_count();
let mut cell_idx: isize = if cell_iter_dir == IterationDirection::Forwards {
0
} else {
cell_count as isize - 1
};
let end = if cell_iter_dir == IterationDirection::Forwards {
cell_count as isize - 1
} else {
0
};
self.stack.set_cell_index(cell_idx as i32);
while cell_count > 0
&& (if cell_iter_dir == IterationDirection::Forwards {
cell_idx <= end
} else {
cell_idx >= end
})
{
let cell = contents.cell_get(
cell_idx,
cell_idx as usize,
payload_overflow_threshold_max(
contents.page_type(),
self.usable_space() as u16,
@ -827,6 +1025,8 @@ impl BTreeCursor {
SeekOp::GT => *cell_rowid > rowid_key,
SeekOp::GE => *cell_rowid >= rowid_key,
SeekOp::EQ => *cell_rowid == rowid_key,
SeekOp::LE => *cell_rowid <= rowid_key,
SeekOp::LT => *cell_rowid < rowid_key,
};
if found {
if let Some(next_page) = first_overflow_page {
@ -841,10 +1041,10 @@ impl BTreeCursor {
self.get_immutable_record_or_create().as_mut().unwrap(),
)?
};
self.stack.advance();
self.stack.next_cell_in_direction(cell_iter_dir);
return Ok(CursorResult::Ok(Some(*cell_rowid)));
} else {
self.stack.advance();
self.stack.next_cell_in_direction(cell_iter_dir);
}
}
BTreeCell::IndexLeafCell(IndexLeafCell {
@ -869,14 +1069,17 @@ impl BTreeCursor {
};
let record = self.get_immutable_record();
let record = record.as_ref().unwrap();
let without_rowid = &record.get_values().as_slice()[..record.len() - 1];
let order = without_rowid.cmp(index_key.get_values());
let record_slice_equal_number_of_cols =
&record.get_values().as_slice()[..index_key.get_values().len()];
let order = record_slice_equal_number_of_cols.cmp(index_key.get_values());
let found = match op {
SeekOp::GT => order.is_gt(),
SeekOp::GE => order.is_ge(),
SeekOp::EQ => order.is_eq(),
SeekOp::LE => order.is_le(),
SeekOp::LT => order.is_lt(),
};
self.stack.advance();
self.stack.next_cell_in_direction(cell_iter_dir);
if found {
let rowid = match record.last_value() {
Some(RefValue::Integer(rowid)) => *rowid as u64,
@ -889,6 +1092,11 @@ impl BTreeCursor {
unreachable!("unexpected cell type: {:?}", cell_type);
}
}
if cell_iter_dir == IterationDirection::Forwards {
cell_idx += 1;
} else {
cell_idx -= 1;
}
}
}
@ -909,7 +1117,20 @@ impl BTreeCursor {
// 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)));
match self.iteration_state {
IterationState::Iterating(IterationDirection::Forwards) => {
return self.get_next_record(Some((key, op)));
}
IterationState::Iterating(IterationDirection::Backwards) => {
return self.get_prev_record(Some((key, op)));
}
IterationState::Unset => {
unreachable!("iteration state must not be unset");
}
IterationState::IterationNotAllowed => {
unreachable!("iteration state must not be IterationNotAllowed");
}
}
}
Ok(CursorResult::Ok(None))
@ -983,6 +1204,13 @@ impl BTreeCursor {
// 6. If we find the cell, we return the record. Otherwise, we return an empty result.
self.move_to_root();
let iter_dir = match self.iteration_state {
IterationState::Iterating(IterationDirection::Backwards) => {
IterationDirection::Backwards
}
_ => IterationDirection::Forwards,
};
loop {
let page = self.stack.top();
return_if_locked!(page);
@ -994,7 +1222,7 @@ impl BTreeCursor {
let mut found_cell = false;
for cell_idx in 0..contents.cell_count() {
match &contents.cell_get(
let cell = contents.cell_get(
cell_idx,
payload_overflow_threshold_max(
contents.page_type(),
@ -1005,25 +1233,78 @@ impl BTreeCursor {
self.usable_space() as u16,
),
self.usable_space(),
)? {
)?;
match &cell {
BTreeCell::TableInteriorCell(TableInteriorCell {
_left_child_page,
_rowid,
_rowid: cell_rowid,
}) => {
let SeekKey::TableRowId(rowid_key) = key else {
unreachable!("table seek key should be a rowid");
};
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,
// in sqlite btrees left child pages have <= keys.
// table btrees can have a duplicate rowid in the interior cell, so for example if we are looking for rowid=10,
// and we find an interior cell with rowid=10, we need to move to the left page since (due to the <= rule of sqlite btrees)
// the left page may have a rowid=10.
// Logic table for determining if target leaf page is in left subtree
//
// Forwards iteration (looking for first match in tree):
// OP | Current Cell vs Seek Key | Action? | Explanation
// GT | > | go left | First > key is in left subtree
// GT | = or < | go right | First > key is in right subtree
// GE | > or = | go left | First >= key is in left subtree
// GE | < | go right | First >= key is in right subtree
//
// Backwards iteration (looking for last match in tree):
// OP | Current Cell vs Seek Key | Action? | Explanation
// LE | > or = | go left | Last <= key is in left subtree
// LE | < | go right | Last <= key is in right subtree
// LT | > or = | go left | Last < key is in left subtree
// LT | < | go right?| Last < key is in right subtree, except if cell rowid is exactly 1 less
//
// No iteration (point query):
// EQ | > or = | go left | Last = key is in left subtree
// EQ | < | go right | Last = key is in right subtree
let target_leaf_page_is_in_left_subtree = match (self.iteration_state, cmp)
{
(
IterationState::Iterating(IterationDirection::Forwards),
SeekOp::GT,
) => *cell_rowid > rowid_key,
(
IterationState::Iterating(IterationDirection::Forwards),
SeekOp::GE,
) => *cell_rowid >= rowid_key,
(
IterationState::Iterating(IterationDirection::Backwards),
SeekOp::LE,
) => *cell_rowid >= rowid_key,
(
IterationState::Iterating(IterationDirection::Backwards),
SeekOp::LT,
) => *cell_rowid >= rowid_key || *cell_rowid == rowid_key - 1,
(_any, SeekOp::EQ) => *cell_rowid >= rowid_key,
_ => unreachable!(
"invalid combination of seek op and iteration state: {:?} {:?}",
cmp, self.iteration_state
),
};
self.stack.advance();
if target_leaf_page_is_in_left_subtree {
// If we found our target rowid in the left subtree,
// we need to move the parent cell pointer forwards or backwards depending on the iteration direction.
// For example: since the internal node contains the max rowid of the left subtree, we need to move the
// parent pointer backwards in backwards iteration so that we don't come back to the parent again.
// E.g.
// this parent: rowid 666
// left child has: 664,665,666
// we need to move to the previous parent (with e.g. rowid 663) when iterating backwards.
self.stack.next_cell_in_direction(iter_dir);
let mem_page = self.pager.read_page(*_left_child_page as usize)?;
self.stack.push(mem_page);
found_cell = true;
break;
} else {
self.stack.advance();
}
}
BTreeCell::TableLeafCell(TableLeafCell {
@ -1057,17 +1338,84 @@ impl BTreeCursor {
self.get_immutable_record_or_create().as_mut().unwrap(),
)?
};
let order = compare_immutable(
let record = self.get_immutable_record();
let record = record.as_ref().unwrap();
let record_slice_equal_number_of_cols =
&record.get_values().as_slice()[..index_key.get_values().len()];
let interior_cell_vs_index_key = compare_immutable(
record_slice_equal_number_of_cols,
index_key.get_values(),
self.get_immutable_record().as_ref().unwrap().get_values(),
);
let target_leaf_page_is_in_the_left_subtree = match cmp {
SeekOp::GT => order.is_lt(),
SeekOp::GE => order.is_le(),
SeekOp::EQ => order.is_le(),
// in sqlite btrees left child pages have <= keys.
// in general, in forwards iteration we want to find the first key that matches the seek condition.
// in backwards iteration we want to find the last key that matches the seek condition.
//
// Logic table for determining if target leaf page is in left subtree.
// For index b-trees this is a bit more complicated since the interior cells contain payloads (the key is the payload).
// and for non-unique indexes there might be several cells with the same key.
//
// Forwards iteration (looking for first match in tree):
// OP | Current Cell vs Seek Key | Action? | Explanation
// GT | > | go left | First > key could be exactly this one, or in left subtree
// GT | = or < | go right | First > key must be in right subtree
// GE | > | go left | First >= key could be exactly this one, or in left subtree
// GE | = | go left | First >= key could be exactly this one, or in left subtree
// GE | < | go right | First >= key must be in right subtree
//
// Backwards iteration (looking for last match in tree):
// OP | Current Cell vs Seek Key | Action? | Explanation
// LE | > | go left | Last <= key must be in left subtree
// LE | = | go right | Last <= key is either this one, or somewhere to the right of this one. So we need to go right to make sure
// LE | < | go right | Last <= key must be in right subtree
// LT | > | go left | Last < key must be in left subtree
// LT | = | go left | Last < key must be in left subtree since we want strictly less than
// LT | < | go right | Last < key could be exactly this one, or in right subtree
//
// No iteration (point query):
// EQ | > | go left | First = key must be in left subtree
// EQ | = | go left | First = key could be exactly this one, or in left subtree
// EQ | < | go right | First = key must be in right subtree
assert!(
self.iteration_state != IterationState::Unset,
"iteration state must have been set before move_to() is called"
);
let target_leaf_page_is_in_left_subtree = match (cmp, self.iteration_state)
{
(
SeekOp::GT,
IterationState::Iterating(IterationDirection::Forwards),
) => interior_cell_vs_index_key.is_gt(),
(
SeekOp::GE,
IterationState::Iterating(IterationDirection::Forwards),
) => interior_cell_vs_index_key.is_ge(),
(SeekOp::EQ, IterationState::IterationNotAllowed) => {
interior_cell_vs_index_key.is_ge()
}
(
SeekOp::LE,
IterationState::Iterating(IterationDirection::Backwards),
) => interior_cell_vs_index_key.is_gt(),
(
SeekOp::LT,
IterationState::Iterating(IterationDirection::Backwards),
) => interior_cell_vs_index_key.is_ge(),
_ => unreachable!(
"invalid combination of seek op and iteration state: {:?} {:?}",
cmp, self.iteration_state
),
};
if target_leaf_page_is_in_the_left_subtree {
// we don't advance in case of index tree internal nodes because we will visit this node going up
if target_leaf_page_is_in_left_subtree {
// we don't advance in case of forward iteration and index tree internal nodes because we will visit this node going up.
// in backwards iteration, we must retreat because otherwise we would unnecessarily visit this node again.
// Example:
// this parent: key 666, and we found the target key in the left child.
// left child has: key 663, key 664, key 665
// we need to move to the previous parent (with e.g. key 662) when iterating backwards so that we don't end up back here again.
if iter_dir == IterationDirection::Backwards {
self.stack.retreat();
}
let mem_page = self.pager.read_page(*left_child_page as usize)?;
self.stack.push(mem_page);
found_cell = true;
@ -2607,6 +2955,14 @@ impl BTreeCursor {
}
pub fn rewind(&mut self) -> Result<CursorResult<()>> {
assert!(
matches!(
self.iteration_state,
IterationState::Unset | IterationState::Iterating(IterationDirection::Forwards)
),
"iteration state must be unset or Iterating(Forwards) when rewind() is called"
);
self.iteration_state = IterationState::Iterating(IterationDirection::Forwards);
if self.mv_cursor.is_some() {
let rowid = return_if_io!(self.get_next_record(None));
self.rowid.replace(rowid);
@ -2622,6 +2978,14 @@ impl BTreeCursor {
}
pub fn last(&mut self) -> Result<CursorResult<()>> {
assert!(
matches!(
self.iteration_state,
IterationState::Unset | IterationState::Iterating(IterationDirection::Backwards)
),
"iteration state must be unset or Iterating(Backwards) when last() is called"
);
self.iteration_state = IterationState::Iterating(IterationDirection::Backwards);
assert!(self.mv_cursor.is_none());
match self.move_to_rightmost()? {
CursorResult::Ok(_) => self.prev(),
@ -2630,6 +2994,14 @@ impl BTreeCursor {
}
pub fn next(&mut self) -> Result<CursorResult<()>> {
assert!(
matches!(
self.iteration_state,
IterationState::Iterating(IterationDirection::Forwards)
),
"iteration state must be Iterating(Forwards) when next() is called, but it was {:?}",
self.iteration_state
);
let rowid = return_if_io!(self.get_next_record(None));
self.rowid.replace(rowid);
self.empty_record.replace(rowid.is_none());
@ -2637,8 +3009,15 @@ impl BTreeCursor {
}
pub fn prev(&mut self) -> Result<CursorResult<()>> {
assert!(
matches!(
self.iteration_state,
IterationState::Iterating(IterationDirection::Backwards)
),
"iteration state must be Iterating(Backwards) when prev() is called"
);
assert!(self.mv_cursor.is_none());
match self.get_prev_record()? {
match self.get_prev_record(None)? {
CursorResult::Ok(rowid) => {
self.rowid.replace(rowid);
self.empty_record.replace(rowid.is_none());
@ -2663,6 +3042,38 @@ impl BTreeCursor {
pub fn seek(&mut self, key: SeekKey<'_>, op: SeekOp) -> Result<CursorResult<bool>> {
assert!(self.mv_cursor.is_none());
match op {
SeekOp::GE | SeekOp::GT => {
if self.iteration_state == IterationState::Unset {
self.iteration_state = IterationState::Iterating(IterationDirection::Forwards);
} else {
assert!(matches!(
self.iteration_state,
IterationState::Iterating(IterationDirection::Forwards)
));
}
}
SeekOp::LE | SeekOp::LT => {
if self.iteration_state == IterationState::Unset {
self.iteration_state = IterationState::Iterating(IterationDirection::Backwards);
} else {
assert!(matches!(
self.iteration_state,
IterationState::Iterating(IterationDirection::Backwards)
));
}
}
SeekOp::EQ => {
if self.iteration_state == IterationState::Unset {
self.iteration_state = IterationState::IterationNotAllowed;
} else {
assert!(matches!(
self.iteration_state,
IterationState::IterationNotAllowed
));
}
}
};
let rowid = return_if_io!(self.do_seek(key, op));
self.rowid.replace(rowid);
self.empty_record.replace(rowid.is_none());
@ -3025,7 +3436,7 @@ impl BTreeCursor {
/// Search for a key in an Index Btree. Looking up indexes that need to be unique, we cannot compare the rowid
pub fn key_exists_in_index(&mut self, key: &ImmutableRecord) -> Result<CursorResult<bool>> {
return_if_io!(self.do_seek(SeekKey::IndexKey(key), SeekOp::GE));
return_if_io!(self.seek(SeekKey::IndexKey(key), SeekOp::GE));
let record_opt = self.record();
match record_opt.as_ref() {
@ -3056,7 +3467,7 @@ impl BTreeCursor {
OwnedValue::Integer(i) => i,
_ => unreachable!("btree tables are indexed by integers!"),
};
return_if_io!(self.move_to(SeekKey::TableRowId(*int_key as u64), SeekOp::EQ));
let _ = return_if_io!(self.move_to(SeekKey::TableRowId(*int_key as u64), SeekOp::EQ));
let page = self.stack.top();
// TODO(pere): request load
return_if_locked!(page);
@ -3485,7 +3896,7 @@ impl PageStack {
}
/// Push a new page onto the stack.
/// This effectively means traversing to a child page.
fn push(&self, page: PageRef) {
fn _push(&self, page: PageRef, starting_cell_idx: i32) {
tracing::trace!(
"pagestack::push(current={}, new_page_id={})",
self.current_page.get(),
@ -3498,7 +3909,15 @@ impl PageStack {
"corrupted database, stack is bigger than expected"
);
self.stack.borrow_mut()[current as usize] = Some(page);
self.cell_indices.borrow_mut()[current as usize] = 0;
self.cell_indices.borrow_mut()[current as usize] = starting_cell_idx;
}
fn push(&self, page: PageRef) {
self._push(page, 0);
}
fn push_backwards(&self, page: PageRef) {
self._push(page, i32::MAX);
}
/// Pop a page off the stack.
@ -3558,6 +3977,18 @@ impl PageStack {
self.cell_indices.borrow_mut()[current] -= 1;
}
/// Move the cursor to the next cell in the current page according to the iteration direction.
fn next_cell_in_direction(&self, iteration_direction: IterationDirection) {
match iteration_direction {
IterationDirection::Forwards => {
self.advance();
}
IterationDirection::Backwards => {
self.retreat();
}
}
}
fn set_cell_index(&self, idx: i32) {
let current = self.current();
self.cell_indices.borrow_mut()[current] = idx
@ -4824,7 +5255,7 @@ mod tests {
run_until_done(
|| {
let key = SeekKey::TableRowId(key as u64);
cursor.move_to(key, SeekOp::EQ)
cursor.seek(key, SeekOp::EQ)
},
pager.deref(),
)
@ -4841,6 +5272,8 @@ mod tests {
// FIXME: add sorted vector instead, should be okay for small amounts of keys for now :P, too lazy to fix right now
keys.sort();
cursor.move_to_root();
// hack to allow bypassing our internal invariant of not allowing cursor iteration after SeekOp::EQ
cursor.iteration_state = IterationState::Iterating(IterationDirection::Forwards);
let mut valid = true;
for key in keys.iter() {
tracing::trace!("seeking key: {}", key);
@ -4852,6 +5285,7 @@ mod tests {
break;
}
}
cursor.iteration_state = IterationState::Unset;
// let's validate btree too so that we undertsand where the btree failed
if matches!(validate_btree(pager.clone(), root_page), (_, false)) || !valid {
let btree_after = format_btree(pager.clone(), root_page, 0);
@ -4869,6 +5303,8 @@ mod tests {
}
keys.sort();
cursor.move_to_root();
// hack to allow bypassing our internal invariant of not allowing cursor iteration after SeekOp::EQ
cursor.iteration_state = IterationState::Iterating(IterationDirection::Forwards);
for key in keys.iter() {
tracing::trace!("seeking key: {}", key);
run_until_done(|| cursor.next(), pager.deref()).unwrap();
@ -5740,7 +6176,7 @@ mod tests {
run_until_done(
|| {
let key = SeekKey::TableRowId(i as u64);
cursor.move_to(key, SeekOp::EQ)
cursor.seek(key, SeekOp::EQ)
},
pager.deref(),
)
@ -5820,7 +6256,7 @@ mod tests {
run_until_done(
|| {
let key = SeekKey::TableRowId(i as u64);
cursor.move_to(key, SeekOp::EQ)
cursor.seek(key, SeekOp::EQ)
},
pager.deref(),
)
@ -5902,7 +6338,7 @@ mod tests {
run_until_done(
|| {
let key = SeekKey::TableRowId(i as u64);
cursor.move_to(key, SeekOp::EQ)
cursor.seek(key, SeekOp::EQ)
},
pager.deref(),
)

7
core/translate/delete.rs
View file

@ -7,7 +7,7 @@ use crate::vdbe::builder::{ProgramBuilder, ProgramBuilderOpts, QueryMode};
use crate::{schema::Schema, Result, SymbolTable};
use limbo_sqlite3_parser::ast::{Expr, Limit, QualifiedName};
use super::plan::TableReference;
use super::plan::{IterationDirection, TableReference};
pub fn translate_delete(
query_mode: QueryMode,
@ -53,7 +53,10 @@ pub fn prepare_delete_plan(
let table_references = vec![TableReference {
table,
identifier: name,
op: Operation::Scan { iter_dir: None },
op: Operation::Scan {
iter_dir: IterationDirection::Forwards,
index: None,
},
join_info: None,
}];

521
core/translate/main_loop.rs
View file

@ -79,7 +79,7 @@ pub fn init_loop(
}
}
match &table.op {
Operation::Scan { .. } => {
Operation::Scan { index, .. } => {
let cursor_id = program.alloc_cursor_id(
Some(table.identifier.clone()),
match &table.table {
@ -90,6 +90,9 @@ pub fn init_loop(
other => panic!("Invalid table reference type in Scan: {:?}", other),
},
);
let index_cursor_id = index.as_ref().map(|i| {
program.alloc_cursor_id(Some(i.name.clone()), CursorType::BTreeIndex(i.clone()))
});
match (mode, &table.table) {
(OperationMode::SELECT, Table::BTree(btree)) => {
let root_page = btree.root_page;
@ -98,6 +101,13 @@ pub fn init_loop(
root_page,
});
program.emit_insn(Insn::OpenReadAwait {});
if let Some(index_cursor_id) = index_cursor_id {
program.emit_insn(Insn::OpenReadAsync {
cursor_id: index_cursor_id,
root_page: index.as_ref().unwrap().root_page,
});
program.emit_insn(Insn::OpenReadAwait {});
}
}
(OperationMode::DELETE, Table::BTree(btree)) => {
let root_page = btree.root_page;
@ -113,6 +123,12 @@ pub fn init_loop(
cursor_id,
root_page: root_page.into(),
});
if let Some(index_cursor_id) = index_cursor_id {
program.emit_insn(Insn::OpenWriteAsync {
cursor_id: index_cursor_id,
root_page: index.as_ref().unwrap().root_page.into(),
});
}
program.emit_insn(Insn::OpenWriteAwait {});
}
(OperationMode::SELECT, Table::Virtual(_)) => {
@ -282,36 +298,35 @@ pub fn open_loop(
program.resolve_label(jump_target_when_true, program.offset());
}
}
Operation::Scan { iter_dir } => {
Operation::Scan { iter_dir, index } => {
let cursor_id = program.resolve_cursor_id(&table.identifier);
let index_cursor_id = index.as_ref().map(|i| program.resolve_cursor_id(&i.name));
let iteration_cursor_id = index_cursor_id.unwrap_or(cursor_id);
if !matches!(&table.table, Table::Virtual(_)) {
if iter_dir
.as_ref()
.is_some_and(|dir| *dir == IterationDirection::Backwards)
{
program.emit_insn(Insn::LastAsync { cursor_id });
if *iter_dir == IterationDirection::Backwards {
program.emit_insn(Insn::LastAsync {
cursor_id: iteration_cursor_id,
});
} else {
program.emit_insn(Insn::RewindAsync { cursor_id });
program.emit_insn(Insn::RewindAsync {
cursor_id: iteration_cursor_id,
});
}
}
match &table.table {
Table::BTree(_) => program.emit_insn(
if iter_dir
.as_ref()
.is_some_and(|dir| *dir == IterationDirection::Backwards)
{
Table::BTree(_) => {
program.emit_insn(if *iter_dir == IterationDirection::Backwards {
Insn::LastAwait {
cursor_id,
cursor_id: iteration_cursor_id,
pc_if_empty: loop_end,
}
} else {
Insn::RewindAwait {
cursor_id,
cursor_id: iteration_cursor_id,
pc_if_empty: loop_end,
}
},
),
})
}
Table::Virtual(ref table) => {
let start_reg = program
.alloc_registers(table.args.as_ref().map(|a| a.len()).unwrap_or(0));
@ -337,6 +352,13 @@ pub fn open_loop(
}
program.resolve_label(loop_start, program.offset());
if let Some(index_cursor_id) = index_cursor_id {
program.emit_insn(Insn::DeferredSeek {
index_cursor_id,
table_cursor_id: cursor_id,
});
}
for cond in predicates
.iter()
.filter(|cond| cond.should_eval_at_loop(table_index))
@ -361,139 +383,6 @@ pub fn open_loop(
let table_cursor_id = program.resolve_cursor_id(&table.identifier);
// 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))
} else {
None
};
let cmp_reg = program.alloc_register();
let (cmp_expr, cmp_op) = match search {
Search::IndexSearch {
cmp_expr, cmp_op, ..
} => (cmp_expr, cmp_op),
Search::RowidSearch { cmp_expr, cmp_op } => (cmp_expr, cmp_op),
Search::RowidEq { .. } => unreachable!(),
};
// TODO this only handles ascending indexes
match cmp_op {
ast::Operator::Equals
| ast::Operator::Greater
| ast::Operator::GreaterEquals => {
translate_expr(
program,
Some(tables),
&cmp_expr.expr,
cmp_reg,
&t_ctx.resolver,
)?;
}
ast::Operator::Less | ast::Operator::LessEquals => {
program.emit_insn(Insn::Null {
dest: cmp_reg,
dest_end: None,
});
}
_ => unreachable!(),
}
// If we try to seek to a key that is not present in the table/index, we exit the loop entirely.
program.emit_insn(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: loop_end,
},
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: loop_end,
},
_ => unreachable!(),
});
if *cmp_op == ast::Operator::Less || *cmp_op == ast::Operator::LessEquals {
translate_expr(
program,
Some(tables),
&cmp_expr.expr,
cmp_reg,
&t_ctx.resolver,
)?;
}
program.resolve_label(loop_start, program.offset());
// TODO: We are currently only handling ascending indexes.
// For conditions like index_key > 10, we have already sought 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 sought 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 sought 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.
match cmp_op {
ast::Operator::Equals | ast::Operator::LessEquals => {
if let Some(index_cursor_id) = index_cursor_id {
program.emit_insn(Insn::IdxGT {
cursor_id: index_cursor_id,
start_reg: cmp_reg,
num_regs: 1,
target_pc: loop_end,
});
} else {
let rowid_reg = program.alloc_register();
program.emit_insn(Insn::RowId {
cursor_id: table_cursor_id,
dest: rowid_reg,
});
program.emit_insn(Insn::Gt {
lhs: rowid_reg,
rhs: cmp_reg,
target_pc: loop_end,
flags: CmpInsFlags::default(),
});
}
}
ast::Operator::Less => {
if let Some(index_cursor_id) = index_cursor_id {
program.emit_insn(Insn::IdxGE {
cursor_id: index_cursor_id,
start_reg: cmp_reg,
num_regs: 1,
target_pc: loop_end,
});
} else {
let rowid_reg = program.alloc_register();
program.emit_insn(Insn::RowId {
cursor_id: table_cursor_id,
dest: rowid_reg,
});
program.emit_insn(Insn::Ge {
lhs: rowid_reg,
rhs: cmp_reg,
target_pc: loop_end,
flags: CmpInsFlags::default(),
});
}
}
_ => {}
}
if let Some(index_cursor_id) = index_cursor_id {
program.emit_insn(Insn::DeferredSeek {
index_cursor_id,
table_cursor_id,
});
}
}
if let Search::RowidEq { cmp_expr } = search {
let src_reg = program.alloc_register();
translate_expr(
@ -508,7 +397,280 @@ pub fn open_loop(
src_reg,
target_pc: next,
});
} else {
// Otherwise, it's an index/rowid scan, i.e. first a seek is performed and then a scan until the comparison expression is not satisfied anymore.
let index_cursor_id = if let Search::IndexSearch { index, .. } = search {
Some(program.resolve_cursor_id(&index.name))
} else {
None
};
let (cmp_expr, cmp_op, iter_dir) = match search {
Search::IndexSearch {
cmp_expr,
cmp_op,
iter_dir,
..
} => (cmp_expr, cmp_op, iter_dir),
Search::RowidSearch {
cmp_expr,
cmp_op,
iter_dir,
} => (cmp_expr, cmp_op, iter_dir),
Search::RowidEq { .. } => unreachable!(),
};
// There are a few steps in an index seek:
// 1. Emit the comparison expression for the rowid/index seek. For example, if we a clause 'WHERE index_key >= 10', we emit the comparison expression 10 into cmp_reg.
//
// 2. Emit the seek instruction. SeekGE and SeekGT are used in forwards iteration, SeekLT and SeekLE are used in backwards iteration.
// All of the examples below assume an ascending index, because we do not support descending indexes yet.
// If we are scanning the ascending index:
// - Forwards, and have a GT/GE/EQ comparison, the comparison expression from step 1 is used as the value to seek to, because that is the lowest possible value that satisfies the clause.
// - Forwards, and have a LT/LE comparison, NULL is used as the comparison expression because we actually want to start scanning from the beginning of the index.
// - Backwards, and have a GT/GE comparison, no Seek instruction is emitted and we emit LastAsync instead, because we want to start scanning from the end of the index.
// - Backwards, and have a LT/LE/EQ comparison, we emit a Seek instruction with the comparison expression from step 1 as the value to seek to, since that is the highest possible
// value that satisfies the clause.
let seek_cmp_reg = program.alloc_register();
let mut comparison_expr_translated = false;
match (cmp_op, iter_dir) {
// Forwards, GT/GE/EQ -> use the comparison expression (i.e. seek to the first key where the cmp expr is satisfied, and then scan forwards)
(
ast::Operator::Equals
| ast::Operator::Greater
| ast::Operator::GreaterEquals,
IterationDirection::Forwards,
) => {
translate_expr(
program,
Some(tables),
&cmp_expr.expr,
seek_cmp_reg,
&t_ctx.resolver,
)?;
comparison_expr_translated = true;
match cmp_op {
ast::Operator::Equals | ast::Operator::GreaterEquals => {
program.emit_insn(Insn::SeekGE {
is_index: index_cursor_id.is_some(),
cursor_id: index_cursor_id.unwrap_or(table_cursor_id),
start_reg: seek_cmp_reg,
num_regs: 1,
target_pc: loop_end,
});
}
ast::Operator::Greater => {
program.emit_insn(Insn::SeekGT {
is_index: index_cursor_id.is_some(),
cursor_id: index_cursor_id.unwrap_or(table_cursor_id),
start_reg: seek_cmp_reg,
num_regs: 1,
target_pc: loop_end,
});
}
_ => unreachable!(),
}
}
// Forwards, LT/LE -> use NULL (i.e. start from the beginning of the index)
(
ast::Operator::Less | ast::Operator::LessEquals,
IterationDirection::Forwards,
) => {
program.emit_insn(Insn::Null {
dest: seek_cmp_reg,
dest_end: None,
});
program.emit_insn(Insn::SeekGT {
is_index: index_cursor_id.is_some(),
cursor_id: index_cursor_id.unwrap_or(table_cursor_id),
start_reg: seek_cmp_reg,
num_regs: 1,
target_pc: loop_end,
});
}
// Backwards, GT/GE -> no seek, emit LastAsync (i.e. start from the end of the index)
(
ast::Operator::Greater | ast::Operator::GreaterEquals,
IterationDirection::Backwards,
) => {
program.emit_insn(Insn::LastAsync {
cursor_id: index_cursor_id.unwrap_or(table_cursor_id),
});
program.emit_insn(Insn::LastAwait {
cursor_id: index_cursor_id.unwrap_or(table_cursor_id),
pc_if_empty: loop_end,
});
}
// Backwards, LT/LE/EQ -> use the comparison expression (i.e. seek from the end of the index until the cmp expr is satisfied, and then scan backwards)
(
ast::Operator::Less | ast::Operator::LessEquals | ast::Operator::Equals,
IterationDirection::Backwards,
) => {
translate_expr(
program,
Some(tables),
&cmp_expr.expr,
seek_cmp_reg,
&t_ctx.resolver,
)?;
comparison_expr_translated = true;
match cmp_op {
ast::Operator::Less => {
program.emit_insn(Insn::SeekLT {
is_index: index_cursor_id.is_some(),
cursor_id: index_cursor_id.unwrap_or(table_cursor_id),
start_reg: seek_cmp_reg,
num_regs: 1,
target_pc: loop_end,
});
}
ast::Operator::LessEquals | ast::Operator::Equals => {
program.emit_insn(Insn::SeekLE {
is_index: index_cursor_id.is_some(),
cursor_id: index_cursor_id.unwrap_or(table_cursor_id),
start_reg: seek_cmp_reg,
num_regs: 1,
target_pc: loop_end,
});
}
_ => unreachable!(),
}
}
_ => unreachable!(),
};
program.resolve_label(loop_start, program.offset());
let scan_terminating_cmp_reg = if comparison_expr_translated {
seek_cmp_reg
} else {
let reg = program.alloc_register();
translate_expr(
program,
Some(tables),
&cmp_expr.expr,
reg,
&t_ctx.resolver,
)?;
reg
};
// 3. Emit a scan-terminating comparison instruction (IdxGT, IdxGE, IdxLT, IdxLE if index; GT, GE, LT, LE if btree rowid scan).
// Here the comparison expression from step 1 is compared to the current index key and the loop is exited if the comparison is true.
// The comparison operator used in the Idx__ instruction is the inverse of the WHERE clause comparison operator.
// For example, if we are scanning forwards and have a clause 'WHERE index_key < 10', we emit IdxGE(10) since >=10 is the first key where our condition is not satisfied anymore.
match (cmp_op, iter_dir) {
// Forwards, <= -> terminate if >
(
ast::Operator::Equals | ast::Operator::LessEquals,
IterationDirection::Forwards,
) => {
if let Some(index_cursor_id) = index_cursor_id {
program.emit_insn(Insn::IdxGT {
cursor_id: index_cursor_id,
start_reg: scan_terminating_cmp_reg,
num_regs: 1,
target_pc: loop_end,
});
} else {
let rowid_reg = program.alloc_register();
program.emit_insn(Insn::RowId {
cursor_id: table_cursor_id,
dest: rowid_reg,
});
program.emit_insn(Insn::Gt {
lhs: rowid_reg,
rhs: scan_terminating_cmp_reg,
target_pc: loop_end,
flags: CmpInsFlags::default(),
});
}
}
// Forwards, < -> terminate if >=
(ast::Operator::Less, IterationDirection::Forwards) => {
if let Some(index_cursor_id) = index_cursor_id {
program.emit_insn(Insn::IdxGE {
cursor_id: index_cursor_id,
start_reg: scan_terminating_cmp_reg,
num_regs: 1,
target_pc: loop_end,
});
} else {
let rowid_reg = program.alloc_register();
program.emit_insn(Insn::RowId {
cursor_id: table_cursor_id,
dest: rowid_reg,
});
program.emit_insn(Insn::Ge {
lhs: rowid_reg,
rhs: scan_terminating_cmp_reg,
target_pc: loop_end,
flags: CmpInsFlags::default(),
});
}
}
// Backwards, >= -> terminate if <
(
ast::Operator::Equals | ast::Operator::GreaterEquals,
IterationDirection::Backwards,
) => {
if let Some(index_cursor_id) = index_cursor_id {
program.emit_insn(Insn::IdxLT {
cursor_id: index_cursor_id,
start_reg: scan_terminating_cmp_reg,
num_regs: 1,
target_pc: loop_end,
});
} else {
let rowid_reg = program.alloc_register();
program.emit_insn(Insn::RowId {
cursor_id: table_cursor_id,
dest: rowid_reg,
});
program.emit_insn(Insn::Lt {
lhs: rowid_reg,
rhs: scan_terminating_cmp_reg,
target_pc: loop_end,
flags: CmpInsFlags::default(),
});
}
}
// Backwards, > -> terminate if <=
(ast::Operator::Greater, IterationDirection::Backwards) => {
if let Some(index_cursor_id) = index_cursor_id {
program.emit_insn(Insn::IdxLE {
cursor_id: index_cursor_id,
start_reg: scan_terminating_cmp_reg,
num_regs: 1,
target_pc: loop_end,
});
} else {
let rowid_reg = program.alloc_register();
program.emit_insn(Insn::RowId {
cursor_id: table_cursor_id,
dest: rowid_reg,
});
program.emit_insn(Insn::Le {
lhs: rowid_reg,
rhs: scan_terminating_cmp_reg,
target_pc: loop_end,
flags: CmpInsFlags::default(),
});
}
}
// Forwards, > and >= -> we already did a seek to the first key where the cmp expr is satisfied, so we dont have a terminating condition
// Backwards, < and <= -> we already did a seek to the last key where the cmp expr is satisfied, so we dont have a terminating condition
_ => {}
}
if let Some(index_cursor_id) = index_cursor_id {
// Don't do a btree table seek until it's actually necessary to read from the table.
program.emit_insn(Insn::DeferredSeek {
index_cursor_id,
table_cursor_id,
});
}
}
for cond in predicates
.iter()
.filter(|cond| cond.should_eval_at_loop(table_index))
@ -813,30 +975,33 @@ pub fn close_loop(
target_pc: loop_labels.loop_start,
});
}
Operation::Scan { iter_dir, .. } => {
Operation::Scan {
index, iter_dir, ..
} => {
program.resolve_label(loop_labels.next, program.offset());
let cursor_id = program.resolve_cursor_id(&table.identifier);
let index_cursor_id = index.as_ref().map(|i| program.resolve_cursor_id(&i.name));
let iteration_cursor_id = index_cursor_id.unwrap_or(cursor_id);
match &table.table {
Table::BTree(_) => {
if iter_dir
.as_ref()
.is_some_and(|dir| *dir == IterationDirection::Backwards)
{
program.emit_insn(Insn::PrevAsync { cursor_id });
if *iter_dir == IterationDirection::Backwards {
program.emit_insn(Insn::PrevAsync {
cursor_id: iteration_cursor_id,
});
} else {
program.emit_insn(Insn::NextAsync { cursor_id });
program.emit_insn(Insn::NextAsync {
cursor_id: iteration_cursor_id,
});
}
if iter_dir
.as_ref()
.is_some_and(|dir| *dir == IterationDirection::Backwards)
{
if *iter_dir == IterationDirection::Backwards {
program.emit_insn(Insn::PrevAwait {
cursor_id,
cursor_id: iteration_cursor_id,
pc_if_next: loop_labels.loop_start,
});
} else {
program.emit_insn(Insn::NextAwait {
cursor_id,
cursor_id: iteration_cursor_id,
pc_if_next: loop_labels.loop_start,
});
}
@ -854,17 +1019,29 @@ pub fn close_loop(
program.resolve_label(loop_labels.next, program.offset());
// Rowid equality point lookups are handled with a SeekRowid instruction which does not loop, so there is no need to emit a NextAsync instruction.
if !matches!(search, Search::RowidEq { .. }) {
let cursor_id = match search {
Search::IndexSearch { index, .. } => program.resolve_cursor_id(&index.name),
Search::RowidSearch { .. } => program.resolve_cursor_id(&table.identifier),
let (cursor_id, iter_dir) = match search {
Search::IndexSearch {
index, iter_dir, ..
} => (program.resolve_cursor_id(&index.name), *iter_dir),
Search::RowidSearch { iter_dir, .. } => {
(program.resolve_cursor_id(&table.identifier), *iter_dir)
}
Search::RowidEq { .. } => unreachable!(),
};
program.emit_insn(Insn::NextAsync { cursor_id });
program.emit_insn(Insn::NextAwait {
cursor_id,
pc_if_next: loop_labels.loop_start,
});
if iter_dir == IterationDirection::Backwards {
program.emit_insn(Insn::PrevAsync { cursor_id });
program.emit_insn(Insn::PrevAwait {
cursor_id,
pc_if_next: loop_labels.loop_start,
});
} else {
program.emit_insn(Insn::NextAsync { cursor_id });
program.emit_insn(Insn::NextAwait {
cursor_id,
pc_if_next: loop_labels.loop_start,
});
}
}
}
}

191
core/translate/optimizer.rs
View file

@ -1,6 +1,6 @@
use std::{collections::HashMap, sync::Arc};
use limbo_sqlite3_parser::ast;
use limbo_sqlite3_parser::ast::{self, Expr, SortOrder};
use crate::{
schema::{Index, Schema},
@ -9,8 +9,8 @@ use crate::{
};
use super::plan::{
DeletePlan, Direction, IterationDirection, Operation, Plan, Search, SelectPlan, TableReference,
UpdatePlan, WhereTerm,
DeletePlan, Direction, GroupBy, IterationDirection, Operation, Plan, Search, SelectPlan,
TableReference, UpdatePlan, WhereTerm,
};
pub fn optimize_plan(plan: &mut Plan, schema: &Schema) -> Result<()> {
@ -40,10 +40,10 @@ fn optimize_select_plan(plan: &mut SelectPlan, schema: &Schema) -> Result<()> {
&mut plan.table_references,
&schema.indexes,
&mut plan.where_clause,
&mut plan.order_by,
&plan.group_by,
)?;
eliminate_unnecessary_orderby(plan, schema)?;
eliminate_orderby_like_groupby(plan)?;
Ok(())
@ -62,6 +62,8 @@ fn optimize_delete_plan(plan: &mut DeletePlan, schema: &Schema) -> Result<()> {
&mut plan.table_references,
&schema.indexes,
&mut plan.where_clause,
&mut plan.order_by,
&None,
)?;
Ok(())
@ -79,6 +81,8 @@ fn optimize_update_plan(plan: &mut UpdatePlan, schema: &Schema) -> Result<()> {
&mut plan.table_references,
&schema.indexes,
&mut plan.where_clause,
&mut plan.order_by,
&None,
)?;
Ok(())
}
@ -93,33 +97,6 @@ fn optimize_subqueries(plan: &mut SelectPlan, schema: &Schema) -> Result<()> {
Ok(())
}
fn query_is_already_ordered_by(
table_references: &[TableReference],
key: &mut ast::Expr,
available_indexes: &HashMap<String, Vec<Arc<Index>>>,
) -> Result<bool> {
let first_table = table_references.first();
if first_table.is_none() {
return Ok(false);
}
let table_reference = first_table.unwrap();
match &table_reference.op {
Operation::Scan { .. } => Ok(key.is_rowid_alias_of(0)),
Operation::Search(search) => match search {
Search::RowidEq { .. } => Ok(key.is_rowid_alias_of(0)),
Search::RowidSearch { .. } => Ok(key.is_rowid_alias_of(0)),
Search::IndexSearch { index, .. } => {
let index_rc = key.check_index_scan(0, table_reference, available_indexes)?;
let index_is_the_same = index_rc
.map(|irc| Arc::ptr_eq(index, &irc))
.unwrap_or(false);
Ok(index_is_the_same)
}
},
_ => Ok(false),
}
}
fn eliminate_orderby_like_groupby(plan: &mut SelectPlan) -> Result<()> {
if plan.order_by.is_none() | plan.group_by.is_none() {
return Ok(());
@ -185,36 +162,117 @@ fn eliminate_orderby_like_groupby(plan: &mut SelectPlan) -> Result<()> {
Ok(())
}
fn eliminate_unnecessary_orderby(plan: &mut SelectPlan, schema: &Schema) -> Result<()> {
if plan.order_by.is_none() {
fn eliminate_unnecessary_orderby(
table_references: &mut [TableReference],
available_indexes: &HashMap<String, Vec<Arc<Index>>>,
order_by: &mut Option<Vec<(ast::Expr, Direction)>>,
group_by: &Option<GroupBy>,
) -> Result<()> {
let Some(order) = order_by else {
return Ok(());
}
if plan.table_references.is_empty() {
};
let Some(first_table_reference) = table_references.first_mut() else {
return Ok(());
}
};
let Some(btree_table) = first_table_reference.btree() else {
return Ok(());
};
// If GROUP BY clause is present, we can't rely on already ordered columns because GROUP BY reorders the data
// This early return prevents the elimination of ORDER BY when GROUP BY exists, as sorting must be applied after grouping
// And if ORDER BY clause duplicates GROUP BY we handle it later in fn eliminate_orderby_like_groupby
if plan.group_by.is_some() {
if group_by.is_some() {
return Ok(());
}
let Operation::Scan {
index, iter_dir, ..
} = &mut first_table_reference.op
else {
return Ok(());
};
assert!(
index.is_none(),
"Nothing shouldve transformed the scan to use an index yet"
);
// Special case: if ordering by just the rowid, we can remove the ORDER BY clause
if order.len() == 1 && order[0].0.is_rowid_alias_of(0) {
*iter_dir = match order[0].1 {
Direction::Ascending => IterationDirection::Forwards,
Direction::Descending => IterationDirection::Backwards,
};
*order_by = None;
return Ok(());
}
let o = plan.order_by.as_mut().unwrap();
// Find the best matching index for the ORDER BY columns
let table_name = &btree_table.name;
let mut best_index = (None, 0);
if o.len() != 1 {
// TODO: handle multiple order by keys
return Ok(());
for (_, indexes) in available_indexes.iter() {
for index_candidate in indexes.iter().filter(|i| &i.table_name == table_name) {
let matching_columns = index_candidate.columns.iter().enumerate().take_while(|(i, c)| {
if let Some((Expr::Column { table, column, .. }, _)) = order.get(*i) {
let col_idx_in_table = btree_table
.columns
.iter()
.position(|tc| tc.name.as_ref() == Some(&c.name));
matches!(col_idx_in_table, Some(col_idx) if *table == 0 && *column == col_idx)
} else {
false
}
}).count();
if matching_columns > best_index.1 {
best_index = (Some(index_candidate), matching_columns);
}
}
}
let (key, direction) = o.first_mut().unwrap();
let Some(matching_index) = best_index.0 else {
return Ok(());
};
let match_count = best_index.1;
let already_ordered =
query_is_already_ordered_by(&plan.table_references, key, &schema.indexes)?;
// If we found a matching index, use it for scanning
*index = Some(matching_index.clone());
// If the order by direction matches the index direction, we can iterate the index in forwards order.
// If they don't, we must iterate the index in backwards order.
let index_direction = &matching_index.columns.first().as_ref().unwrap().order;
*iter_dir = match (index_direction, order[0].1) {
(SortOrder::Asc, Direction::Ascending) | (SortOrder::Desc, Direction::Descending) => {
IterationDirection::Forwards
}
(SortOrder::Asc, Direction::Descending) | (SortOrder::Desc, Direction::Ascending) => {
IterationDirection::Backwards
}
};
if already_ordered {
push_scan_direction(&mut plan.table_references[0], direction);
plan.order_by = None;
// If the index covers all ORDER BY columns, and one of the following applies:
// - the ORDER BY directions exactly match the index orderings,
// - the ORDER by directions are the exact opposite of the index orderings,
// we can remove the ORDER BY clause.
if match_count == order.len() {
let full_match = {
let mut all_match_forward = true;
let mut all_match_reverse = true;
for (i, (_, direction)) in order.iter().enumerate() {
match (&matching_index.columns[i].order, direction) {
(SortOrder::Asc, Direction::Ascending)
| (SortOrder::Desc, Direction::Descending) => {
all_match_reverse = false;
}
(SortOrder::Asc, Direction::Descending)
| (SortOrder::Desc, Direction::Ascending) => {
all_match_forward = false;
}
}
}
all_match_forward || all_match_reverse
};
if full_match {
*order_by = None;
}
}
Ok(())
@ -222,24 +280,25 @@ fn eliminate_unnecessary_orderby(plan: &mut SelectPlan, schema: &Schema) -> Resu
/**
* Use indexes where possible.
* Right now we make decisions about using indexes ONLY based on condition expressions, not e.g. ORDER BY or others.
* This is just because we are WIP.
*
* When this function is called, condition expressions from both the actual WHERE clause and the JOIN clauses are in the where_clause vector.
* If we find a condition that can be used to index scan, we pop it off from the where_clause vector and put it into a Search operation.
* We put it there simply because it makes it a bit easier to track during translation.
*
* In this function we also try to eliminate ORDER BY clauses if there is an index that satisfies the ORDER BY clause.
*/
fn use_indexes(
table_references: &mut [TableReference],
available_indexes: &HashMap<String, Vec<Arc<Index>>>,
where_clause: &mut Vec<WhereTerm>,
order_by: &mut Option<Vec<(ast::Expr, Direction)>>,
group_by: &Option<GroupBy>,
) -> Result<()> {
if where_clause.is_empty() {
return Ok(());
}
// Try to use indexes for eliminating ORDER BY clauses
eliminate_unnecessary_orderby(table_references, available_indexes, order_by, group_by)?;
// Try to use indexes for WHERE conditions
'outer: for (table_index, table_reference) in table_references.iter_mut().enumerate() {
if let Operation::Scan { .. } = &mut table_reference.op {
if let Operation::Scan { iter_dir, .. } = &table_reference.op {
let mut i = 0;
while i < where_clause.len() {
let cond = where_clause.get_mut(i).unwrap();
@ -248,6 +307,7 @@ fn use_indexes(
table_index,
table_reference,
available_indexes,
*iter_dir,
)? {
where_clause.remove(i);
table_reference.op = Operation::Search(index_search);
@ -296,20 +356,6 @@ fn eliminate_constant_conditions(
Ok(ConstantConditionEliminationResult::Continue)
}
fn push_scan_direction(table: &mut TableReference, direction: &Direction) {
if let Operation::Scan {
ref mut iter_dir, ..
} = table.op
{
if iter_dir.is_none() {
match direction {
Direction::Ascending => *iter_dir = Some(IterationDirection::Forwards),
Direction::Descending => *iter_dir = Some(IterationDirection::Backwards),
}
}
}
}
fn rewrite_exprs_select(plan: &mut SelectPlan) -> Result<()> {
for rc in plan.result_columns.iter_mut() {
rewrite_expr(&mut rc.expr)?;
@ -611,6 +657,7 @@ pub fn try_extract_index_search_expression(
table_index: usize,
table_reference: &TableReference,
available_indexes: &HashMap<String, Vec<Arc<Index>>>,
iter_dir: IterationDirection,
) -> Result<Option<Search>> {
if !cond.should_eval_at_loop(table_index) {
return Ok(None);
@ -641,6 +688,7 @@ pub fn try_extract_index_search_expression(
from_outer_join: cond.from_outer_join,
eval_at: cond.eval_at,
},
iter_dir,
}));
}
_ => {}
@ -671,6 +719,7 @@ pub fn try_extract_index_search_expression(
from_outer_join: cond.from_outer_join,
eval_at: cond.eval_at,
},
iter_dir,
}));
}
_ => {}
@ -695,6 +744,7 @@ pub fn try_extract_index_search_expression(
from_outer_join: cond.from_outer_join,
eval_at: cond.eval_at,
},
iter_dir,
}));
}
_ => {}
@ -719,6 +769,7 @@ pub fn try_extract_index_search_expression(
from_outer_join: cond.from_outer_join,
eval_at: cond.eval_at,
},
iter_dir,
}));
}
_ => {}

11
core/translate/plan.rs
View file

@ -259,12 +259,11 @@ pub struct TableReference {
pub enum Operation {
// Scan operation
// This operation is used to scan a table.
// The iter_dir are uset to indicate the direction of the iterator.
// The use of Option for iter_dir is aimed at implementing a conservative optimization strategy: it only pushes
// iter_dir down to Scan when iter_dir is None, to prevent potential result set errors caused by multiple
// assignments. for more detailed discussions, please refer to https://github.com/tursodatabase/limbo/pull/376
// The iter_dir is used to indicate the direction of the iterator.
Scan {
iter_dir: Option<IterationDirection>,
iter_dir: IterationDirection,
/// The index that we are using to scan the table, if any.
index: Option<Arc<Index>>,
},
// Search operation
// This operation is used to search for a row in a table using an index
@ -337,12 +336,14 @@ pub enum Search {
RowidSearch {
cmp_op: ast::Operator,
cmp_expr: WhereTerm,
iter_dir: IterationDirection,
},
/// A secondary index search. Uses bytecode instructions like SeekGE, SeekGT etc.
IndexSearch {
index: Arc<Index>,
cmp_op: ast::Operator,
cmp_expr: WhereTerm,
iter_dir: IterationDirection,
},
}

14
core/translate/planner.rs
View file

@ -1,7 +1,7 @@
use super::{
plan::{
Aggregate, EvalAt, JoinInfo, Operation, Plan, ResultSetColumn, SelectPlan, SelectQueryType,
TableReference, WhereTerm,
Aggregate, EvalAt, IterationDirection, JoinInfo, Operation, Plan, ResultSetColumn,
SelectPlan, SelectQueryType, TableReference, WhereTerm,
},
select::prepare_select_plan,
SymbolTable,
@ -320,7 +320,10 @@ fn parse_from_clause_table<'a>(
));
};
scope.tables.push(TableReference {
op: Operation::Scan { iter_dir: None },
op: Operation::Scan {
iter_dir: IterationDirection::Forwards,
index: None,
},
table: tbl_ref,
identifier: alias.unwrap_or(normalized_qualified_name),
join_info: None,
@ -399,7 +402,10 @@ fn parse_from_clause_table<'a>(
.unwrap_or(normalized_name.to_string());
scope.tables.push(TableReference {
op: Operation::Scan { iter_dir: None },
op: Operation::Scan {
iter_dir: IterationDirection::Forwards,
index: None,
},
join_info: None,
table: Table::Virtual(vtab),
identifier: alias,

24
core/translate/update.rs
View file

@ -13,7 +13,8 @@ use super::optimizer::optimize_plan;
use super::plan::{
Direction, IterationDirection, Plan, ResultSetColumn, TableReference, UpdatePlan,
};
use super::planner::{bind_column_references, parse_limit, parse_where};
use super::planner::bind_column_references;
use super::planner::{parse_limit, parse_where};
/*
* Update is simple. By default we scan the table, and for each row, we check the WHERE
@ -72,18 +73,25 @@ pub fn prepare_update_plan(schema: &Schema, body: &mut Update) -> crate::Result<
let Some(btree_table) = table.btree() else {
bail_parse_error!("Error: {} is not a btree table", table_name);
};
let iter_dir: Option<IterationDirection> = body.order_by.as_ref().and_then(|order_by| {
order_by.first().and_then(|ob| {
ob.order.map(|o| match o {
SortOrder::Asc => IterationDirection::Forwards,
SortOrder::Desc => IterationDirection::Backwards,
let iter_dir = body
.order_by
.as_ref()
.and_then(|order_by| {
order_by.first().and_then(|ob| {
ob.order.map(|o| match o {
SortOrder::Asc => IterationDirection::Forwards,
SortOrder::Desc => IterationDirection::Backwards,
})
})
})
});
.unwrap_or(IterationDirection::Forwards);
let table_references = vec![TableReference {
table: Table::BTree(btree_table.clone()),
identifier: table_name.0.clone(),
op: Operation::Scan { iter_dir },
op: Operation::Scan {
iter_dir,
index: None,
},
join_info: None,
}];
let set_clauses = body

4
core/types.rs
View file

@ -1203,11 +1203,13 @@ pub enum CursorResult<T> {
IO,
}
#[derive(Clone, PartialEq, Debug)]
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum SeekOp {
EQ,
GE,
GT,
LE,
LT,
}
#[derive(Clone, PartialEq, Debug)]

6
core/vdbe/builder.rs
View file

@ -413,6 +413,12 @@ impl ProgramBuilder {
Insn::SeekGT { target_pc, .. } => {
resolve(target_pc, "SeekGT");
}
Insn::SeekLE { target_pc, .. } => {
resolve(target_pc, "SeekLE");
}
Insn::SeekLT { target_pc, .. } => {
resolve(target_pc, "SeekLT");
}
Insn::IdxGE { target_pc, .. } => {
resolve(target_pc, "IdxGE");
}

113
core/vdbe/execute.rs
View file

@ -1892,97 +1892,69 @@ pub fn op_deferred_seek(
Ok(InsnFunctionStepResult::Step)
}
pub fn op_seek_ge(
pub fn op_seek(
program: &Program,
state: &mut ProgramState,
insn: &Insn,
pager: &Rc<Pager>,
mv_store: Option<&Rc<MvStore>>,
) -> Result<InsnFunctionStepResult> {
let Insn::SeekGE {
let (Insn::SeekGE {
cursor_id,
start_reg,
num_regs,
target_pc,
is_index,
} = insn
else {
unreachable!("unexpected Insn {:?}", insn)
};
assert!(target_pc.is_offset());
if *is_index {
let found = {
let mut cursor = state.get_cursor(*cursor_id);
let cursor = cursor.as_btree_mut();
let record_from_regs = make_record(&state.registers, start_reg, num_regs);
let found =
return_if_io!(cursor.seek(SeekKey::IndexKey(&record_from_regs), SeekOp::GE));
found
};
if !found {
state.pc = target_pc.to_offset_int();
} else {
state.pc += 1;
}
} else {
let pc = {
let mut cursor = state.get_cursor(*cursor_id);
let cursor = cursor.as_btree_mut();
let rowid = match state.registers[*start_reg].get_owned_value() {
OwnedValue::Null => {
// All integer values are greater than null so we just rewind the cursor
return_if_io!(cursor.rewind());
None
}
OwnedValue::Integer(rowid) => Some(*rowid as u64),
_ => {
return Err(LimboError::InternalError(
"SeekGE: the value in the register is not an integer".into(),
));
}
};
match rowid {
Some(rowid) => {
let found = return_if_io!(cursor.seek(SeekKey::TableRowId(rowid), SeekOp::GE));
if !found {
target_pc.to_offset_int()
} else {
state.pc + 1
}
}
None => state.pc + 1,
}
};
state.pc = pc;
}
Ok(InsnFunctionStepResult::Step)
}
pub fn op_seek_gt(
program: &Program,
state: &mut ProgramState,
insn: &Insn,
pager: &Rc<Pager>,
mv_store: Option<&Rc<MvStore>>,
) -> Result<InsnFunctionStepResult> {
let Insn::SeekGT {
| Insn::SeekGT {
cursor_id,
start_reg,
num_regs,
target_pc,
is_index,
} = insn
}
| Insn::SeekLE {
cursor_id,
start_reg,
num_regs,
target_pc,
is_index,
}
| Insn::SeekLT {
cursor_id,
start_reg,
num_regs,
target_pc,
is_index,
}) = insn
else {
unreachable!("unexpected Insn {:?}", insn)
};
assert!(target_pc.is_offset());
assert!(
target_pc.is_offset(),
"target_pc should be an offset, is: {:?}",
target_pc
);
let op = match insn {
Insn::SeekGE { .. } => SeekOp::GE,
Insn::SeekGT { .. } => SeekOp::GT,
Insn::SeekLE { .. } => SeekOp::LE,
Insn::SeekLT { .. } => SeekOp::LT,
_ => unreachable!("unexpected Insn {:?}", insn),
};
let op_name = match op {
SeekOp::GE => "SeekGE",
SeekOp::GT => "SeekGT",
SeekOp::LE => "SeekLE",
SeekOp::LT => "SeekLT",
_ => unreachable!("unexpected SeekOp {:?}", op),
};
if *is_index {
let found = {
let mut cursor = state.get_cursor(*cursor_id);
let cursor = cursor.as_btree_mut();
let record_from_regs = make_record(&state.registers, start_reg, num_regs);
let found =
return_if_io!(cursor.seek(SeekKey::IndexKey(&record_from_regs), SeekOp::GT));
let found = return_if_io!(cursor.seek(SeekKey::IndexKey(&record_from_regs), op));
found
};
if !found {
@ -2002,14 +1974,15 @@ pub fn op_seek_gt(
}
OwnedValue::Integer(rowid) => Some(*rowid as u64),
_ => {
return Err(LimboError::InternalError(
"SeekGT: the value in the register is not an integer".into(),
));
return Err(LimboError::InternalError(format!(
"{}: the value in the register is not an integer",
op_name
)));
}
};
let found = match rowid {
Some(rowid) => {
let found = return_if_io!(cursor.seek(SeekKey::TableRowId(rowid), SeekOp::GT));
let found = return_if_io!(cursor.seek(SeekKey::TableRowId(rowid), op));
if !found {
target_pc.to_offset_int()
} else {

50
core/vdbe/explain.rs
View file

@ -736,23 +736,35 @@ pub fn insn_to_str(
start_reg,
num_regs: _,
target_pc,
} => (
"SeekGT",
*cursor_id as i32,
target_pc.to_debug_int(),
*start_reg as i32,
OwnedValue::build_text(""),
0,
"".to_string(),
),
Insn::SeekGE {
}
| Insn::SeekGE {
is_index: _,
cursor_id,
start_reg,
num_regs: _,
target_pc,
}
| Insn::SeekLE {
is_index: _,
cursor_id,
start_reg,
num_regs: _,
target_pc,
}
| Insn::SeekLT {
is_index: _,
cursor_id,
start_reg,
num_regs: _,
target_pc,
} => (
"SeekGE",
match insn {
Insn::SeekGT { .. } => "SeekGT",
Insn::SeekGE { .. } => "SeekGE",
Insn::SeekLE { .. } => "SeekLE",
Insn::SeekLT { .. } => "SeekLT",
_ => unreachable!(),
},
*cursor_id as i32,
target_pc.to_debug_int(),
*start_reg as i32,
@ -1213,9 +1225,9 @@ pub fn insn_to_str(
0,
"".to_string(),
),
Insn::LastAsync { .. } => (
Insn::LastAsync { cursor_id } => (
"LastAsync",
0,
*cursor_id as i32,
0,
0,
OwnedValue::build_text(""),
@ -1240,27 +1252,27 @@ pub fn insn_to_str(
0,
where_clause.clone(),
),
Insn::LastAwait { .. } => (
Insn::LastAwait { cursor_id, .. } => (
"LastAwait",
0,
*cursor_id as i32,
0,
0,
OwnedValue::build_text(""),
0,
"".to_string(),
),
Insn::PrevAsync { .. } => (
Insn::PrevAsync { cursor_id } => (
"PrevAsync",
0,
*cursor_id as i32,
0,
0,
OwnedValue::build_text(""),
0,
"".to_string(),
),
Insn::PrevAwait { .. } => (
Insn::PrevAwait { cursor_id, .. } => (
"PrevAwait",
0,
*cursor_id as i32,
0,
0,
OwnedValue::build_text(""),

30
core/vdbe/insn.rs
View file

@ -501,6 +501,30 @@ pub enum Insn {
/// 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.
// 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 less than or equal to the given key. If not found, jump to the given PC. Otherwise, continue to the next instruction.
SeekLE {
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 less than the given key. If not found, jump to the given PC. Otherwise, continue to the next instruction.
SeekLT {
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,
@ -1306,8 +1330,10 @@ impl Insn {
Insn::SeekRowid { .. } => execute::op_seek_rowid,
Insn::DeferredSeek { .. } => execute::op_deferred_seek,
Insn::SeekGE { .. } => execute::op_seek_ge,
Insn::SeekGT { .. } => execute::op_seek_gt,
Insn::SeekGE { .. } => execute::op_seek,
Insn::SeekGT { .. } => execute::op_seek,
Insn::SeekLE { .. } => execute::op_seek,
Insn::SeekLT { .. } => execute::op_seek,
Insn::SeekEnd { .. } => execute::op_seek_end,
Insn::IdxGE { .. } => execute::op_idx_ge,
Insn::IdxGT { .. } => execute::op_idx_gt,

5
core/vdbe/mod.rs
View file

@ -561,7 +561,10 @@ fn get_indent_count(indent_count: usize, curr_insn: &Insn, prev_insn: Option<&In
| Insn::LastAwait { .. }
| Insn::SorterSort { .. }
| Insn::SeekGE { .. }
| Insn::SeekGT { .. } => indent_count + 1,
| Insn::SeekGT { .. }
| Insn::SeekLE { .. }
| Insn::SeekLT { .. } => indent_count + 1,
_ => indent_count,
}
} else {

59
testing/orderby.test
View file

@ -141,3 +141,62 @@ Collin|15}
do_execsql_test case-insensitive-alias {
select u.first_name as fF, count(1) > 0 as cC from users u where fF = 'Jamie' group by fF order by cC;
} {Jamie|1}
do_execsql_test age_idx_order_desc {
select first_name from users order by age desc limit 3;
} {Robert
Sydney
Matthew}
do_execsql_test rowid_or_integer_pk_desc {
select first_name from users order by id desc limit 3;
} {Nicole
Gina
Dorothy}
# These two following tests may seem dumb but they verify that index scanning by age_idx doesn't drop any rows due to BTree bugs
do_execsql_test orderby_asc_verify_rows {
select count(1) from (select * from users order by age desc)
} {10000}
do_execsql_test orderby_desc_verify_rows {
select count(1) from (select * from users order by age desc)
} {10000}
do_execsql_test orderby_desc_with_offset {
select first_name, age from users order by age desc limit 3 offset 666;
} {Francis|94
Matthew|94
Theresa|94}
do_execsql_test orderby_desc_with_filter {
select first_name, age from users where age <= 50 order by age desc limit 5;
} {Gerald|50
Nicole|50
Tammy|50
Marissa|50
Daniel|50}
do_execsql_test orderby_asc_with_filter_range {
select first_name, age from users where age <= 50 and age >= 49 order by age asc limit 5;
} {William|49
Jennifer|49
Robert|49
David|49
Stephanie|49}
do_execsql_test orderby_desc_with_filter_id_lt {
select id from users where id < 6666 order by id desc limit 5;
} {6665
6664
6663
6662
6661}
do_execsql_test orderby_desc_with_filter_id_le {
select id from users where id <= 6666 order by id desc limit 5;
} {6666
6665
6664
6663
6662}

205
tests/integration/fuzz/mod.rs
View file

@ -2,9 +2,9 @@ pub mod grammar_generator;
#[cfg(test)]
mod tests {
use std::rc::Rc;
use std::{collections::HashSet, rc::Rc};
use rand::SeedableRng;
use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha8Rng;
use rusqlite::params;
@ -107,6 +107,207 @@ mod tests {
}
}
#[test]
pub fn rowid_seek_fuzz() {
let db = TempDatabase::new_with_rusqlite("CREATE TABLE t(x INTEGER PRIMARY KEY)"); // INTEGER PRIMARY KEY is a rowid alias, so an index is not created
let sqlite_conn = rusqlite::Connection::open(db.path.clone()).unwrap();
let insert = format!(
"INSERT INTO t VALUES {}",
(1..10000)
.map(|x| format!("({})", x))
.collect::<Vec<_>>()
.join(", ")
);
sqlite_conn.execute(&insert, params![]).unwrap();
sqlite_conn.close().unwrap();
let sqlite_conn = rusqlite::Connection::open(db.path.clone()).unwrap();
let limbo_conn = db.connect_limbo();
const COMPARISONS: [&str; 4] = ["<", "<=", ">", ">="];
const ORDER_BY: [Option<&str>; 4] = [
None,
Some("ORDER BY x"),
Some("ORDER BY x DESC"),
Some("ORDER BY x ASC"),
];
for comp in COMPARISONS.iter() {
for order_by in ORDER_BY.iter() {
for max in 0..=10000 {
let query = format!(
"SELECT * FROM t WHERE x {} {} {} LIMIT 3",
comp,
max,
order_by.unwrap_or("")
);
log::trace!("query: {}", query);
let limbo = limbo_exec_rows(&db, &limbo_conn, &query);
let sqlite = sqlite_exec_rows(&sqlite_conn, &query);
assert_eq!(
limbo, sqlite,
"query: {}, limbo: {:?}, sqlite: {:?}",
query, limbo, sqlite
);
}
}
}
}
#[test]
pub fn index_scan_fuzz() {
let db = TempDatabase::new_with_rusqlite("CREATE TABLE t(x PRIMARY KEY)");
let sqlite_conn = rusqlite::Connection::open(db.path.clone()).unwrap();
let insert = format!(
"INSERT INTO t VALUES {}",
(0..10000)
.map(|x| format!("({})", x))
.collect::<Vec<_>>()
.join(", ")
);
sqlite_conn.execute(&insert, params![]).unwrap();
sqlite_conn.close().unwrap();
let sqlite_conn = rusqlite::Connection::open(db.path.clone()).unwrap();
let limbo_conn = db.connect_limbo();
const COMPARISONS: [&str; 5] = ["=", "<", "<=", ">", ">="];
const ORDER_BY: [Option<&str>; 4] = [
None,
Some("ORDER BY x"),
Some("ORDER BY x DESC"),
Some("ORDER BY x ASC"),
];
for comp in COMPARISONS.iter() {
for order_by in ORDER_BY.iter() {
for max in 0..=10000 {
let query = format!(
"SELECT * FROM t WHERE x {} {} {} LIMIT 3",
comp,
max,
order_by.unwrap_or(""),
);
let limbo = limbo_exec_rows(&db, &limbo_conn, &query);
let sqlite = sqlite_exec_rows(&sqlite_conn, &query);
assert_eq!(
limbo, sqlite,
"query: {}, limbo: {:?}, sqlite: {:?}",
query, limbo, sqlite
);
}
}
}
}
#[test]
pub fn index_scan_compound_key_fuzz() {
let (mut rng, seed) = if std::env::var("SEED").is_ok() {
let seed = std::env::var("SEED").unwrap().parse::<u64>().unwrap();
(ChaCha8Rng::seed_from_u64(seed), seed)
} else {
rng_from_time()
};
let db = TempDatabase::new_with_rusqlite("CREATE TABLE t(x, y, z, PRIMARY KEY (x, y))");
let sqlite_conn = rusqlite::Connection::open(db.path.clone()).unwrap();
let mut pk_tuples = HashSet::new();
while pk_tuples.len() < 100000 {
pk_tuples.insert((rng.random_range(0..3000), rng.random_range(0..3000)));
}
let mut tuples = Vec::new();
for pk_tuple in pk_tuples {
tuples.push(format!(
"({}, {}, {})",
pk_tuple.0,
pk_tuple.1,
rng.random_range(0..2000)
));
}
let insert = format!("INSERT INTO t VALUES {}", tuples.join(", "));
sqlite_conn.execute(&insert, params![]).unwrap();
sqlite_conn.close().unwrap();
let sqlite_conn = rusqlite::Connection::open(db.path.clone()).unwrap();
let limbo_conn = db.connect_limbo();
const COMPARISONS: [&str; 5] = ["=", "<", "<=", ">", ">="];
const ORDER_BY: [Option<&str>; 4] = [
None,
Some("ORDER BY x"),
Some("ORDER BY x DESC"),
Some("ORDER BY x ASC"),
];
let print_dump_on_fail = |insert: &str, seed: u64| {
let comment = format!("-- seed: {}; dump for manual debugging:", seed);
let pragma_journal_mode = "PRAGMA journal_mode = wal;";
let create_table = "CREATE TABLE t(x, y, z, PRIMARY KEY (x, y));";
let dump = format!(
"{}\n{}\n{}\n{}\n{}",
comment, pragma_journal_mode, create_table, comment, insert
);
println!("{}", dump);
};
for comp in COMPARISONS.iter() {
for order_by in ORDER_BY.iter() {
for max in 0..=3000 {
// see comment below about ordering and the '=' comparison operator; omitting LIMIT for that reason
// we mainly have LIMIT here for performance reasons but for = we want to get all the rows to ensure
// correctness in the = case
let limit = if *comp == "=" { "" } else { "LIMIT 5" };
let query = format!(
"SELECT * FROM t WHERE x {} {} {} {}",
comp,
max,
order_by.unwrap_or(""),
limit
);
log::trace!("query: {}", query);
let limbo = limbo_exec_rows(&db, &limbo_conn, &query);
let sqlite = sqlite_exec_rows(&sqlite_conn, &query);
let is_equal = limbo == sqlite;
if !is_equal {
// if the condition is = and the same rows are present but in different order, then we accept that
// e.g. sqlite doesn't bother iterating in reverse order if "WHERE X = 3 ORDER BY X DESC", but we currently do.
if *comp == "=" {
let limbo_row_count = limbo.len();
let sqlite_row_count = sqlite.len();
if limbo_row_count == sqlite_row_count {
for limbo_row in limbo.iter() {
if !sqlite.contains(limbo_row) {
// save insert to file and print the filename for debugging
let error_msg = format!("row not found in sqlite: query: {}, limbo: {:?}, sqlite: {:?}, seed: {}", query, limbo, sqlite, seed);
print_dump_on_fail(&insert, seed);
panic!("{}", error_msg);
}
}
for sqlite_row in sqlite.iter() {
if !limbo.contains(sqlite_row) {
let error_msg = format!("row not found in limbo: query: {}, limbo: {:?}, sqlite: {:?}, seed: {}", query, limbo, sqlite, seed);
print_dump_on_fail(&insert, seed);
panic!("{}", error_msg);
}
}
continue;
} else {
print_dump_on_fail(&insert, seed);
let error_msg = format!("row count mismatch (limbo: {}, sqlite: {}): query: {}, limbo: {:?}, sqlite: {:?}, seed: {}", limbo_row_count, sqlite_row_count, query, limbo, sqlite, seed);
panic!("{}", error_msg);
}
}
print_dump_on_fail(&insert, seed);
panic!(
"query: {}, limbo: {:?}, sqlite: {:?}, seed: {}",
query, limbo, sqlite, seed
);
}
}
}
}
}
#[test]
pub fn arithmetic_expression_fuzz() {
let _ = env_logger::try_init();