add collations to btree cursor

This commit is contained in:
pedrocarlo
2025-05-16 18:41:14 -03:00
parent 5bd47d7462
commit f28ce2b757
3 changed files with 84 additions and 25 deletions

View File

@@ -407,6 +407,10 @@ pub struct BTreeCursor {
context: Option<CursorContext>,
/// Store whether the Cursor is in a valid state. Meaning if it is pointing to a valid cell index or not
valid_state: CursorValidState,
/// Colations for Index Btree constraint checks
/// Contains the Collation Seq for the whole Table
/// This Vec should be empty for Table Btree
collations: Vec<CollationSeq>,
}
impl BTreeCursor {
@@ -414,6 +418,7 @@ impl BTreeCursor {
mv_cursor: Option<Rc<RefCell<MvCursor>>>,
pager: Rc<Pager>,
root_page: usize,
collations: Vec<CollationSeq>,
) -> Self {
Self {
mv_cursor,
@@ -435,17 +440,27 @@ impl BTreeCursor {
count: 0,
context: None,
valid_state: CursorValidState::Valid,
collations,
}
}
pub fn new_table(
mv_cursor: Option<Rc<RefCell<MvCursor>>>,
pager: Rc<Pager>,
root_page: usize,
) -> Self {
Self::new(mv_cursor, pager, root_page, Vec::new())
}
pub fn new_index(
mv_cursor: Option<Rc<RefCell<MvCursor>>>,
pager: Rc<Pager>,
root_page: usize,
index: &Index,
collations: Vec<CollationSeq>,
) -> Self {
let index_key_sort_order = IndexKeySortOrder::from_index(index);
let mut cursor = Self::new(mv_cursor, pager, root_page);
let mut cursor = Self::new(mv_cursor, pager, root_page, collations);
cursor.index_key_sort_order = index_key_sort_order;
cursor
}
@@ -5953,7 +5968,7 @@ mod tests {
}
fn validate_btree(pager: Rc<Pager>, page_idx: usize) -> (usize, bool) {
let cursor = BTreeCursor::new(None, pager.clone(), page_idx);
let cursor = BTreeCursor::new_table(None, pager.clone(), page_idx);
let page = pager.read_page(page_idx).unwrap();
let page = page.get();
let contents = page.contents.as_ref().unwrap();
@@ -6043,7 +6058,7 @@ mod tests {
}
fn format_btree(pager: Rc<Pager>, page_idx: usize, depth: usize) -> String {
let cursor = BTreeCursor::new(None, pager.clone(), page_idx);
let cursor = BTreeCursor::new_table(None, pager.clone(), page_idx);
let page = pager.read_page(page_idx).unwrap();
let page = page.get();
let contents = page.contents.as_ref().unwrap();
@@ -6173,7 +6188,7 @@ mod tests {
.as_slice(),
] {
let (pager, root_page) = empty_btree();
let mut cursor = BTreeCursor::new(None, pager.clone(), root_page);
let mut cursor = BTreeCursor::new_table(None, pager.clone(), root_page);
for (key, size) in sequence.iter() {
run_until_done(
|| {
@@ -6240,7 +6255,7 @@ mod tests {
tracing::info!("super seed: {}", seed);
for _ in 0..attempts {
let (pager, root_page) = empty_btree();
let mut cursor = BTreeCursor::new(None, pager.clone(), root_page);
let mut cursor = BTreeCursor::new_table(None, pager.clone(), root_page);
let mut keys = SortedVec::new();
tracing::info!("seed: {}", seed);
for insert_id in 0..inserts {
@@ -6347,7 +6362,7 @@ mod tests {
let (pager, _) = empty_btree();
let index_root_page = pager.btree_create(&CreateBTreeFlags::new_index());
let index_root_page = index_root_page as usize;
let mut cursor = BTreeCursor::new(None, pager.clone(), index_root_page);
let mut cursor = BTreeCursor::new_table(None, pager.clone(), index_root_page);
let mut keys = SortedVec::new();
tracing::info!("seed: {}", seed);
for _ in 0..inserts {
@@ -6582,7 +6597,7 @@ mod tests {
#[ignore]
pub fn test_clear_overflow_pages() -> Result<()> {
let (pager, db_header) = setup_test_env(5);
let mut cursor = BTreeCursor::new(None, pager.clone(), 1);
let mut cursor = BTreeCursor::new_table(None, pager.clone(), 1);
let max_local = payload_overflow_threshold_max(PageType::TableLeaf, 4096);
let usable_size = cursor.usable_space();
@@ -6681,7 +6696,7 @@ mod tests {
#[test]
pub fn test_clear_overflow_pages_no_overflow() -> Result<()> {
let (pager, db_header) = setup_test_env(5);
let mut cursor = BTreeCursor::new(None, pager.clone(), 1);
let mut cursor = BTreeCursor::new_table(None, pager.clone(), 1);
let small_payload = vec![b'A'; 10];
@@ -6725,7 +6740,7 @@ mod tests {
fn test_btree_destroy() -> Result<()> {
let initial_size = 3;
let (pager, db_header) = setup_test_env(initial_size);
let mut cursor = BTreeCursor::new(None, pager.clone(), 2);
let mut cursor = BTreeCursor::new_table(None, pager.clone(), 2);
assert_eq!(
db_header.lock().database_size,
initial_size,
@@ -7382,7 +7397,7 @@ mod tests {
let (pager, root_page) = empty_btree();
let mut keys = Vec::new();
for i in 0..10000 {
let mut cursor = BTreeCursor::new(None, pager.clone(), root_page);
let mut cursor = BTreeCursor::new_table(None, pager.clone(), root_page);
tracing::info!("INSERT INTO t VALUES ({});", i,);
let value = ImmutableRecord::from_registers(&[Register::Value(Value::Integer(i))]);
tracing::trace!("before insert {}", i);
@@ -7409,7 +7424,7 @@ mod tests {
format_btree(pager.clone(), root_page, 0)
);
for key in keys.iter() {
let mut cursor = BTreeCursor::new(None, pager.clone(), root_page);
let mut cursor = BTreeCursor::new_table(None, pager.clone(), root_page);
let key = Value::Integer(*key);
let exists = run_until_done(|| cursor.exists(&key), pager.deref()).unwrap();
assert!(exists, "key not found {}", key);
@@ -7458,7 +7473,7 @@ mod tests {
// Insert 10,000 records in to the BTree.
for i in 1..=10000 {
let mut cursor = BTreeCursor::new(None, pager.clone(), root_page);
let mut cursor = BTreeCursor::new_table(None, pager.clone(), root_page);
let value = ImmutableRecord::from_registers(&[Register::Value(Value::Text(
Text::new("hello world"),
))]);
@@ -7486,7 +7501,7 @@ mod tests {
// Delete records with 500 <= key <= 3500
for i in 500..=3500 {
let mut cursor = BTreeCursor::new(None, pager.clone(), root_page);
let mut cursor = BTreeCursor::new_table(None, pager.clone(), root_page);
let seek_key = SeekKey::TableRowId(i as u64);
let found = run_until_done(|| cursor.seek(seek_key.clone(), SeekOp::EQ), pager.deref())
@@ -7503,7 +7518,7 @@ mod tests {
continue;
}
let mut cursor = BTreeCursor::new(None, pager.clone(), root_page);
let mut cursor = BTreeCursor::new_table(None, pager.clone(), root_page);
let key = Value::Integer(i);
let exists = run_until_done(|| cursor.exists(&key), pager.deref()).unwrap();
assert!(exists, "Key {} should exist but doesn't", i);
@@ -7511,7 +7526,7 @@ mod tests {
// Verify the deleted records don't exist.
for i in 500..=3500 {
let mut cursor = BTreeCursor::new(None, pager.clone(), root_page);
let mut cursor = BTreeCursor::new_table(None, pager.clone(), root_page);
let key = Value::Integer(i);
let exists = run_until_done(|| cursor.exists(&key), pager.deref()).unwrap();
assert!(!exists, "Deleted key {} still exists", i);
@@ -7533,7 +7548,7 @@ mod tests {
let (pager, root_page) = empty_btree();
for i in 0..iterations {
let mut cursor = BTreeCursor::new(None, pager.clone(), root_page);
let mut cursor = BTreeCursor::new_table(None, pager.clone(), root_page);
tracing::info!("INSERT INTO t VALUES ({});", i,);
let value = ImmutableRecord::from_registers(&[Register::Value(Value::Text(Text {
value: huge_texts[i].as_bytes().to_vec(),
@@ -7562,7 +7577,7 @@ mod tests {
format_btree(pager.clone(), root_page, 0)
);
}
let mut cursor = BTreeCursor::new(None, pager.clone(), root_page);
let mut cursor = BTreeCursor::new_table(None, pager.clone(), root_page);
cursor.move_to_root();
for i in 0..iterations {
let rowid = run_until_done(|| cursor.get_next_record(None), pager.deref()).unwrap();

View File

@@ -927,15 +927,36 @@ pub fn op_open_read(
let mut cursors = state.cursors.borrow_mut();
match cursor_type {
CursorType::BTreeTable(_) => {
let cursor = BTreeCursor::new(mv_cursor, pager.clone(), *root_page);
let cursor = BTreeCursor::new_table(mv_cursor, pager.clone(), *root_page);
cursors
.get_mut(*cursor_id)
.unwrap()
.replace(Cursor::new_btree(cursor));
}
CursorType::BTreeIndex(index) => {
let cursor =
BTreeCursor::new_index(mv_cursor, pager.clone(), *root_page, index.as_ref());
let table = program.table_references.iter().find_map(|table_ref| {
table_ref.btree().and_then(|table| {
if table.name == index.table_name {
Some(table)
} else {
None
}
})
});
let collations = table.map_or(Vec::new(), |table| {
table
.columns
.iter()
.map(|column| column.collation.unwrap_or_default())
.collect::<Vec<_>>()
});
let cursor = BTreeCursor::new_index(
mv_cursor,
pager.clone(),
*root_page,
index.as_ref(),
collations,
);
cursors
.get_mut(*cursor_id)
.unwrap()
@@ -4220,14 +4241,35 @@ pub fn op_open_write(
None => None,
};
if let Some(index) = maybe_index {
let cursor =
BTreeCursor::new_index(mv_cursor, pager.clone(), root_page as usize, index.as_ref());
let table = program.table_references.iter().find_map(|table_ref| {
table_ref.btree().and_then(|table| {
if table.name == index.table_name {
Some(table)
} else {
None
}
})
});
let collations = table.map_or(Vec::new(), |table| {
table
.columns
.iter()
.map(|column| column.collation.unwrap_or_default())
.collect::<Vec<_>>()
});
let cursor = BTreeCursor::new_index(
mv_cursor,
pager.clone(),
root_page as usize,
index.as_ref(),
collations,
);
cursors
.get_mut(*cursor_id)
.unwrap()
.replace(Cursor::new_btree(cursor));
} else {
let cursor = BTreeCursor::new(mv_cursor, pager.clone(), root_page as usize);
let cursor = BTreeCursor::new_table(mv_cursor, pager.clone(), root_page as usize);
cursors
.get_mut(*cursor_id)
.unwrap()
@@ -4297,7 +4339,8 @@ pub fn op_destroy(
if *is_temp == 1 {
todo!("temp databases not implemented yet.");
}
let mut cursor = BTreeCursor::new(None, pager.clone(), *root);
// TODO not sure if should be BTreeCursor::new_table or BTreeCursor::new_index here or neither and just pass an emtpy vec
let mut cursor = BTreeCursor::new(None, pager.clone(), *root, Vec::new());
cursor.btree_destroy()?;
state.pc += 1;
Ok(InsnFunctionStepResult::Step)
@@ -4671,7 +4714,7 @@ pub fn op_open_ephemeral(
}
None => None,
};
let mut cursor = BTreeCursor::new(mv_cursor, pager, root_page as usize);
let mut cursor = BTreeCursor::new_table(mv_cursor, pager, root_page as usize);
cursor.rewind()?; // Will never return io
let mut cursors: std::cell::RefMut<'_, Vec<Option<Cursor>>> = state.cursors.borrow_mut();

View File

@@ -1220,6 +1220,7 @@ mod tests {
);
let query = format!("INSERT INTO t VALUES ({}, {}, {})", x, y, z);
log::info!("insert: {}", query);
dbg!(&query);
assert_eq!(
limbo_exec_rows(&db, &limbo_conn, &query),
sqlite_exec_rows(&sqlite_conn, &query),