Skip to content

Commit 95d381e

Browse files
PThorpe92PThorpe92
committed
WIP - Update support
1 parent 38d2afc commit 95d381e

File tree

5 files changed

+429
-40
lines changed

5 files changed

+429
-40
lines changed

‎core/storage/btree.rs

+158-35
Original file line numberDiff line numberDiff line change
@@ -820,43 +820,64 @@ impl BTreeCursor {
820820
(self.find_cell(page, int_key), page.page_type())
821821
};
822822

823-
// TODO: if overwrite drop cell
824-
825-
// insert cell
826-
let mut cell_payload: Vec<u8> = Vec::new();
827-
fill_cell_payload(
828-
page_type,
829-
Some(int_key),
830-
&mut cell_payload,
831-
record,
832-
self.usable_space() as u16,
833-
self.pager.clone(),
834-
);
835-
836-
// insert
837-
let overflow = {
838-
let contents = page.get().contents.as_mut().unwrap();
839-
debug!(
840-
"insert_into_page(overflow, cell_count={})",
841-
contents.cell_count()
842-
);
843-
844-
insert_into_cell(
845-
contents,
846-
cell_payload.as_slice(),
823+
// check if we have an exact match on the cell at cell_idx:
824+
// (`update` support/insert record to same existing rowid)
825+
if cell_idx < page.get_contents().cell_count() {
826+
// parse the existing cell to get the stored key
827+
if let BTreeCell::TableLeafCell(tbl_leaf) = page.get_contents().cell_get(
847828
cell_idx,
848-
self.usable_space() as u16,
849-
)?;
850-
contents.overflow_cells.len()
851-
};
852-
let write_info = self
853-
.state
854-
.mut_write_info()
855-
.expect("can't count while inserting");
856-
if overflow > 0 {
857-
write_info.state = WriteState::BalanceStart;
829+
self.pager.clone(),
830+
payload_overflow_threshold_max(page_type, self.usable_space() as u16),
831+
payload_overflow_threshold_min(page_type, self.usable_space() as u16),
832+
self.usable_space(),
833+
)? {
834+
if tbl_leaf._rowid == int_key {
835+
tracing::debug!("insert_into_page: found exact match with cell_idx={cell_idx}, overwriting");
836+
// we are "overwriting in place" the old cell, not removing the pointers
837+
self.overwrite_cell(page.clone(), cell_idx, record)?;
838+
self.state
839+
.mut_write_info()
840+
.expect("expected write info")
841+
.state = WriteState::Finish;
842+
}
843+
}
858844
} else {
859-
write_info.state = WriteState::Finish;
845+
// insert cell
846+
let mut cell_payload: Vec<u8> = Vec::new();
847+
fill_cell_payload(
848+
page_type,
849+
Some(int_key),
850+
&mut cell_payload,
851+
record,
852+
self.usable_space() as u16,
853+
self.pager.clone(),
854+
);
855+
856+
// insert
857+
let overflow = {
858+
let contents = page.get().contents.as_mut().unwrap();
859+
debug!(
860+
"insert_into_page(overflow, cell_count={})",
861+
contents.cell_count()
862+
);
863+
864+
insert_into_cell(
865+
contents,
866+
cell_payload.as_slice(),
867+
cell_idx,
868+
self.usable_space() as u16,
869+
)?;
870+
contents.overflow_cells.len()
871+
};
872+
let write_info = self
873+
.state
874+
.mut_write_info()
875+
.expect("can't count while inserting");
876+
if overflow > 0 {
877+
write_info.state = WriteState::BalanceStart;
878+
} else {
879+
write_info.state = WriteState::Finish;
880+
}
860881
}
861882
}
862883
WriteState::BalanceStart
@@ -2200,6 +2221,108 @@ impl BTreeCursor {
22002221
pub fn table_id(&self) -> usize {
22012222
self.root_page
22022223
}
2224+
2225+
pub fn overwrite_cell(
2226+
&mut self,
2227+
page_ref: PageRef,
2228+
cell_idx: usize,
2229+
record: &Record,
2230+
) -> Result<CursorResult<()>> {
2231+
// build the new payload
2232+
let page_type = page_ref.get().contents.as_ref().unwrap().page_type();
2233+
let mut new_payload = Vec::new();
2234+
fill_cell_payload(
2235+
page_type,
2236+
self.rowid.get(),
2237+
&mut new_payload,
2238+
record,
2239+
self.usable_space() as u16,
2240+
self.pager.clone(),
2241+
);
2242+
2243+
// figure out old cell offset & size
2244+
let (old_offset, old_local_size) = {
2245+
let page = page_ref.get().contents.as_ref().unwrap();
2246+
page.cell_get_raw_region(
2247+
cell_idx,
2248+
payload_overflow_threshold_max(page_type, self.usable_space() as u16),
2249+
payload_overflow_threshold_min(page_type, self.usable_space() as u16),
2250+
self.usable_space(),
2251+
)
2252+
};
2253+
2254+
// if it all fits in local space and old_local_size is enough, do an in-place overwrite
2255+
if new_payload.len() <= old_local_size {
2256+
self.overwrite_content(
2257+
page_ref.clone(),
2258+
old_offset,
2259+
&new_payload,
2260+
0,
2261+
new_payload.len(),
2262+
)?;
2263+
// if there's leftover local space (old_local_size > new_payload.len()), zero it or free it
2264+
let remaining = old_local_size - new_payload.len();
2265+
if remaining > 0 {
2266+
let buf = page_ref.get().contents.as_mut().unwrap().as_ptr();
2267+
for i in 0..remaining {
2268+
buf[old_offset + new_payload.len() + i] = 0;
2269+
}
2270+
}
2271+
Ok(CursorResult::Ok(()))
2272+
} else {
2273+
// doesn't fit, drop it and insert a new one
2274+
drop_cell(
2275+
page_ref.get_contents(),
2276+
cell_idx,
2277+
self.usable_space() as u16,
2278+
)?;
2279+
insert_into_cell(
2280+
page_ref.get_contents(),
2281+
&new_payload,
2282+
cell_idx,
2283+
self.usable_space() as u16,
2284+
)?;
2285+
Ok(CursorResult::Ok(()))
2286+
}
2287+
}
2288+
2289+
pub fn overwrite_content(
2290+
&mut self,
2291+
page_ref: PageRef,
2292+
dest_offset: usize,
2293+
new_payload: &[u8],
2294+
src_offset: usize,
2295+
amount: usize,
2296+
) -> Result<CursorResult<()>> {
2297+
return_if_locked!(page_ref);
2298+
page_ref.set_dirty();
2299+
self.pager.add_dirty(page_ref.get().id);
2300+
let buf = page_ref.get().contents.as_mut().unwrap().as_ptr();
2301+
2302+
// if new_payload doesn't have enough data, we fill with zeros
2303+
let n_data = new_payload.len().saturating_sub(src_offset);
2304+
if n_data == 0 {
2305+
// everything is zeros
2306+
for i in 0..amount {
2307+
if buf[dest_offset + i] != 0 {
2308+
buf[dest_offset + i] = 0;
2309+
}
2310+
}
2311+
} else {
2312+
let copy_len = n_data.min(amount);
2313+
// copy the overlapping portion
2314+
buf[dest_offset..dest_offset + copy_len]
2315+
.copy_from_slice(&new_payload[src_offset..src_offset + copy_len]);
2316+
2317+
// if copy_len < amount => fill remainder with 0
2318+
if copy_len < amount {
2319+
for i in copy_len..amount {
2320+
buf[dest_offset + i] = 0;
2321+
}
2322+
}
2323+
}
2324+
Ok(CursorResult::Ok(()))
2325+
}
22032326
}
22042327

22052328
impl PageStack {

‎core/translate/mod.rs

+3-1
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ pub(crate) mod result_row;
2323
pub(crate) mod select;
2424
pub(crate) mod subquery;
2525
pub(crate) mod transaction;
26+
pub(crate) mod update;
2627

2728
use crate::fast_lock::SpinLock;
2829
use crate::schema::Schema;
@@ -41,6 +42,7 @@ use std::fmt::Display;
4142
use std::rc::{Rc, Weak};
4243
use std::sync::Arc;
4344
use transaction::{translate_tx_begin, translate_tx_commit};
45+
use update::translate_update;
4446

4547
/// Translate SQL statement into bytecode program.
4648
pub fn translate(
@@ -109,7 +111,7 @@ pub fn translate(
109111
ast::Stmt::Rollback { .. } => bail_parse_error!("ROLLBACK not supported yet"),
110112
ast::Stmt::Savepoint(_) => bail_parse_error!("SAVEPOINT not supported yet"),
111113
ast::Stmt::Select(select) => translate_select(query_mode, schema, *select, syms)?,
112-
ast::Stmt::Update { .. } => bail_parse_error!("UPDATE not supported yet"),
114+
ast::Stmt::Update(update) => translate_update(query_mode, schema, &update, syms)?,
113115
ast::Stmt::Vacuum(_, _) => bail_parse_error!("VACUUM not supported yet"),
114116
ast::Stmt::Insert(insert) => {
115117
let Insert {

0 commit comments

Comments
 (0)