From 2323763a4f3bb0dee414ed271e850fbe7071dc88 Mon Sep 17 00:00:00 2001 From: Krishna Vishal Date: Fri, 20 Jun 2025 21:04:11 +0530 Subject: [PATCH] Integrate incremental column parsing into `btree.rs` --- core/storage/btree.rs | 53 +++++++++++++++++++++++++++++++------------ 1 file changed, 38 insertions(+), 15 deletions(-) diff --git a/core/storage/btree.rs b/core/storage/btree.rs index a360f2f16..6d245290b 100644 --- a/core/storage/btree.rs +++ b/core/storage/btree.rs @@ -13,7 +13,7 @@ use crate::{ }, translate::{collate::CollationSeq, plan::IterationDirection}, turso_assert, - types::{IndexKeyInfo, IndexKeySortOrder, ParseRecordState}, + types::{IndexKeyInfo, IndexKeySortOrder, ParseRecordState, RecordCursor}, MvCursor, }; @@ -37,6 +37,7 @@ use std::{ cmp::{Ordering, Reverse}, collections::BinaryHeap, fmt::Debug, + ops::DerefMut, pin::Pin, rc::Rc, sync::Arc, @@ -305,7 +306,7 @@ impl BTreeKey<'_> { } /// Assert that the key is an index key and return it. - fn to_index_key_values(&self) -> &'_ Vec { + fn to_index_key_values(&self) -> Vec { match self { BTreeKey::TableRowId(_) => panic!("BTreeKey::to_index_key called on TableRowId"), BTreeKey::IndexKey(key) => key.get_values(), @@ -528,6 +529,7 @@ pub struct BTreeCursor { read_overflow_state: RefCell>, /// Contains the current cell_idx for `find_cell` find_cell_state: FindCellState, + pub record_cursor: RefCell, } impl BTreeCursor { @@ -561,6 +563,7 @@ impl BTreeCursor { read_overflow_state: RefCell::new(None), find_cell_state: FindCellState(None), parse_record_state: RefCell::new(ParseRecordState::Init), + record_cursor: RefCell::new(RecordCursor::new()), } } @@ -602,8 +605,15 @@ impl BTreeCursor { if !self.has_rowid() { return None; } - let rowid = match self.get_immutable_record().as_ref().unwrap().last_value() { - Some(RefValue::Integer(rowid)) => *rowid, + let mut record_cursor_ref = self.record_cursor.borrow_mut(); + let record_cursor = record_cursor_ref.deref_mut(); + let rowid = match self + .get_immutable_record() + .as_ref() + .unwrap() + .last_value(record_cursor) + { + Some(Ok(RefValue::Integer(rowid))) => rowid as i64, _ => unreachable!( "index where has_rowid() is true should have an integer rowid as the last value" ), @@ -1561,11 +1571,12 @@ impl BTreeCursor { let (target_leaf_page_is_in_left_subtree, is_eq) = { let record = self.get_immutable_record(); let record = record.as_ref().unwrap(); + let record_values = record.get_values(); let record_slice_equal_number_of_cols = - &record.get_values().as_slice()[..index_key.get_values().len()]; + &record_values[..index_key.get_values().len()]; let interior_cell_vs_index_key = compare_immutable( record_slice_equal_number_of_cols, - index_key.get_values(), + index_key.get_values().as_slice(), self.key_sort_order(), &self.collations, ); @@ -1970,11 +1981,12 @@ impl BTreeCursor { let record = self.get_immutable_record(); let record = record.as_ref().unwrap(); tracing::debug!(?record); - let record_slice_equal_number_of_cols = - &record.get_values().as_slice()[..key.get_values().len()]; + let record_values = record.get_values(); + let key_values = key.get_values(); + let record_slice_equal_number_of_cols = &record_values[..key_values.len()]; compare_immutable( record_slice_equal_number_of_cols, - key.get_values(), + key_values.as_slice(), self.key_sort_order(), &self.collations, ) @@ -2065,7 +2077,7 @@ impl BTreeCursor { let record = bkey .get_record() .expect("expected record present on insert"); - + let record_values = record.get_values(); if let CursorState::None = &self.state { self.state = CursorState::Write(WriteInfo::new()); } @@ -2133,11 +2145,11 @@ impl BTreeCursor { BTreeCell::IndexLeafCell(..) => { // Not necessary to read record again here, as find_cell already does that for us let cmp = compare_immutable( - record.get_values(), + record_values.as_slice(), self.get_immutable_record() .as_ref() .unwrap() - .get_values(), + .get_values().as_slice(), self.key_sort_order(), &self.collations, ); @@ -2165,7 +2177,8 @@ impl BTreeCursor { } } // insert cell - let mut cell_payload: Vec = Vec::with_capacity(record.len() + 4); + + let mut cell_payload: Vec = Vec::with_capacity(record_values.len() + 4); fill_cell_payload( page_type, bkey.maybe_rowid(), @@ -4816,7 +4829,8 @@ impl BTreeCursor { ) -> Result> { // build the new payload let page_type = page_ref.get().get().contents.as_ref().unwrap().page_type(); - let mut new_payload = Vec::with_capacity(record.len()); + let serial_types_len = self.record_cursor.borrow_mut().len(record); + let mut new_payload = Vec::with_capacity(serial_types_len); let rowid = return_if_io!(self.rowid()); fill_cell_payload( page_type, @@ -4871,7 +4885,7 @@ impl BTreeCursor { fn get_immutable_record_or_create(&self) -> std::cell::RefMut<'_, Option> { if self.reusable_immutable_record.borrow().is_none() { - let record = ImmutableRecord::new(4096, 10); + let record = ImmutableRecord::new(4096); self.reusable_immutable_record.replace(Some(record)); } self.reusable_immutable_record.borrow_mut() @@ -5020,6 +5034,15 @@ impl BTreeCursor { self.pager .do_allocate_page(page_type, offset, BtreePageAllocMode::Any) } + + pub fn get_column(&self, idx: usize) -> Result { + let record_ref = self.reusable_immutable_record.borrow(); + let record = record_ref + .as_ref() + .ok_or_else(|| LimboError::InternalError("No record available".into()))?; + + self.record_cursor.borrow_mut().get_value(record, idx) + } } #[derive(Debug, thiserror::Error)]