aljaz/turso
1
0
Fork 0
mirror of https://github.com/aljazceru/turso.git synced 2025-12-25 20:14:21 +01:00
Code Issues Packages Projects Releases Wiki Activity

Merge 'Support offset clause in Update queries ' from Preston Thorpe

Browse Source 
Closes #1245

Reviewed-by: Jussi Saurio <jussi.saurio@gmail.com>

Closes #1252
This commit is contained in:
Pekka Enberg
2025-04-06 19:37:18 +03:00
parent 2d3fd01f91 62e0a137dd
commit 400ebe6b60
8 changed files with 111 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
core/translate/delete.rs
View File

@@ -68,7 +68,7 @@ pub fn prepare_delete_plan(
)?;
// Parse the LIMIT/OFFSET clause
let (resolved_limit, resolved_offset) = limit.map_or(Ok((None, None)), |l| parse_limit(*l))?;
let (resolved_limit, resolved_offset) = limit.map_or(Ok((None, None)), |l| parse_limit(&l))?;
let plan = DeletePlan {
table_references,
@@ -86,7 +86,5 @@ pub fn prepare_delete_plan(
fn estimate_num_instructions(plan: &DeletePlan) -> usize {
let base = 20;
let num_instructions = base + plan.table_references.len() * 10;
num_instructions
base + plan.table_references.len() * 10
}

88
core/translate/emitter.rs
View File

@@ -442,25 +442,11 @@ fn emit_program_for_update(
// Exit on LIMIT 0
if let Some(0) = plan.limit {
epilogue(program, init_label, start_offset, TransactionMode::Read)?;
epilogue(program, init_label, start_offset, TransactionMode::None)?;
program.result_columns = plan.returning.unwrap_or_default();
program.table_references = plan.table_references;
return Ok(());
}
let after_main_loop_label = program.allocate_label();
t_ctx.label_main_loop_end = Some(after_main_loop_label);
if plan.contains_constant_false_condition {
program.emit_insn(Insn::Goto {
target_pc: after_main_loop_label,
});
}
let skip_label = program.allocate_label();
init_loop(
program,
&mut t_ctx,
&plan.table_references,
OperationMode::UPDATE,
)?;
if t_ctx.reg_limit.is_none() && plan.limit.is_some() {
let reg = program.alloc_register();
t_ctx.reg_limit = Some(reg);
@@ -469,7 +455,36 @@ fn emit_program_for_update(
dest: reg,
});
program.mark_last_insn_constant();
if t_ctx.reg_offset.is_none() && plan.offset.is_some_and(|n| n.ne(&0)) {
let reg = program.alloc_register();
t_ctx.reg_offset = Some(reg);
program.emit_insn(Insn::Integer {
value: plan.offset.unwrap() as i64,
dest: reg,
});
program.mark_last_insn_constant();
let combined_reg = program.alloc_register();
t_ctx.reg_limit_offset_sum = Some(combined_reg);
program.emit_insn(Insn::OffsetLimit {
limit_reg: t_ctx.reg_limit.unwrap(),
offset_reg: reg,
combined_reg,
});
}
}
let after_main_loop_label = program.allocate_label();
t_ctx.label_main_loop_end = Some(after_main_loop_label);
if plan.contains_constant_false_condition {
program.emit_insn(Insn::Goto {
target_pc: after_main_loop_label,
});
}
init_loop(
program,
&mut t_ctx,
&plan.table_references,
OperationMode::UPDATE,
)?;
open_loop(
program,
&mut t_ctx,
@@ -477,7 +492,6 @@ fn emit_program_for_update(
&plan.where_clause,
)?;
emit_update_insns(&plan, &t_ctx, program)?;
program.resolve_label(skip_label, program.offset());
close_loop(program, &mut t_ctx, &plan.table_references)?;
program.resolve_label(after_main_loop_label, program.offset());
@@ -495,6 +509,7 @@ fn emit_update_insns(
program: &mut ProgramBuilder,
) -> crate::Result<()> {
let table_ref = &plan.table_references.first().unwrap();
let loop_labels = t_ctx.labels_main_loop.first().unwrap();
let (cursor_id, index) = match &table_ref.op {
Operation::Scan { .. } => (program.resolve_cursor_id(&table_ref.identifier), None),
Operation::Search(search) => match search {
@@ -508,24 +523,6 @@ fn emit_update_insns(
},
_ => return Ok(()),
};
for cond in plan.where_clause.iter().filter(|c| c.is_constant()) {
let jump_target = program.allocate_label();
let meta = ConditionMetadata {
jump_if_condition_is_true: false,
jump_target_when_true: jump_target,
jump_target_when_false: t_ctx.label_main_loop_end.unwrap(),
};
translate_condition_expr(
program,
&plan.table_references,
&cond.expr,
meta,
&t_ctx.resolver,
)?;
program.resolve_label(jump_target, program.offset());
}
let first_col_reg = program.alloc_registers(table_ref.table.columns().len());
let rowid_reg = program.alloc_register();
program.emit_insn(Insn::RowId {
cursor_id,
@@ -537,6 +534,29 @@ fn emit_update_insns(
target_pc: t_ctx.label_main_loop_end.unwrap(),
});
if let Some(offset) = t_ctx.reg_offset {
program.emit_insn(Insn::IfPos {
reg: offset,
target_pc: loop_labels.next,
decrement_by: 1,
});
}
for cond in plan.where_clause.iter().filter(|c| c.is_constant()) {
let meta = ConditionMetadata {
jump_if_condition_is_true: false,
jump_target_when_true: BranchOffset::Placeholder,
jump_target_when_false: loop_labels.next,
};
translate_condition_expr(
program,
&plan.table_references,
&cond.expr,
meta,
&t_ctx.resolver,
)?;
}
let first_col_reg = program.alloc_registers(table_ref.table.columns().len());
// we scan a column at a time, loading either the column's values, or the new value
// from the Set expression, into registers so we can emit a MakeRecord and update the row.
for idx in 0..table_ref.columns().len() {

6
core/translate/main_loop.rs
View File

@@ -37,11 +37,11 @@ pub struct LeftJoinMetadata {
#[derive(Debug, Clone, Copy)]
pub struct LoopLabels {
/// jump to the start of the loop body
loop_start: BranchOffset,
pub loop_start: BranchOffset,
/// jump to the NextAsync instruction (or equivalent)
next: BranchOffset,
pub next: BranchOffset,
/// jump to the end of the loop, exiting it
loop_end: BranchOffset,
pub loop_end: BranchOffset,
}
impl LoopLabels {

2
core/translate/plan.rs
View File

@@ -176,8 +176,8 @@ pub struct UpdatePlan {
pub set_clauses: Vec<(usize, ast::Expr)>,
pub where_clause: Vec<WhereTerm>,
pub order_by: Option<Vec<(ast::Expr, Direction)>>,
// TODO: support OFFSET
pub limit: Option<isize>,
pub offset: Option<isize>,
// TODO: optional RETURNING clause
pub returning: Option<Vec<ResultSetColumn>>,
// whether the WHERE clause is always false

23
core/translate/planner.rs
View File

@@ -850,30 +850,33 @@ fn parse_join<'a>(
Ok(())
}
pub fn parse_limit(limit: Limit) -> Result<(Option<isize>, Option<isize>)> {
let offset_val = match limit.offset {
pub fn parse_limit(limit: &Limit) -> Result<(Option<isize>, Option<isize>)> {
let offset_val = match &limit.offset {
Some(offset_expr) => match offset_expr {
Expr::Literal(ast::Literal::Numeric(n)) => n.parse().ok(),
// If OFFSET is negative, the result is as if OFFSET is zero
Expr::Unary(UnaryOperator::Negative, expr) => match *expr {
Expr::Literal(ast::Literal::Numeric(n)) => n.parse::<isize>().ok().map(|num| -num),
_ => crate::bail_parse_error!("Invalid OFFSET clause"),
},
Expr::Unary(UnaryOperator::Negative, expr) => {
if let Expr::Literal(ast::Literal::Numeric(ref n)) = &**expr {
n.parse::<isize>().ok().map(|num| -num)
} else {
crate::bail_parse_error!("Invalid OFFSET clause");
}
}
_ => crate::bail_parse_error!("Invalid OFFSET clause"),
},
None => Some(0),
};
if let Expr::Literal(ast::Literal::Numeric(n)) = limit.expr {
if let Expr::Literal(ast::Literal::Numeric(n)) = &limit.expr {
Ok((n.parse().ok(), offset_val))
} else if let Expr::Unary(UnaryOperator::Negative, expr) = limit.expr {
if let Expr::Literal(ast::Literal::Numeric(n)) = *expr {
} else if let Expr::Unary(UnaryOperator::Negative, expr) = &limit.expr {
if let Expr::Literal(ast::Literal::Numeric(n)) = &**expr {
let limit_val = n.parse::<isize>().ok().map(|num| -num);
Ok((limit_val, offset_val))
} else {
crate::bail_parse_error!("Invalid LIMIT clause");
}
} else if let Expr::Id(id) = limit.expr {
} else if let Expr::Id(id) = &limit.expr {
if id.0.eq_ignore_ascii_case("true") {
Ok((Some(1), offset_val))
} else if id.0.eq_ignore_ascii_case("false") {

2
core/translate/select.rs
View File

@@ -370,7 +370,7 @@ pub fn prepare_select_plan<'a>(
// Parse the LIMIT/OFFSET clause
(plan.limit, plan.offset) =
select.limit.map_or(Ok((None, None)), |l| parse_limit(*l))?;
select.limit.map_or(Ok((None, None)), |l| parse_limit(&l))?;
// Return the unoptimized query plan
Ok(Plan::Select(plan))

15
core/translate/update.rs
View File

@@ -151,17 +151,21 @@ pub fn prepare_update_plan(schema: &Schema, body: &mut Update) -> crate::Result<
})
.collect()
});
// Parse the WHERE clause
parse_where(
body.where_clause.as_ref().map(|w| *w.clone()),
&table_references,
Some(&result_columns),
&mut where_clause,
)?;
let limit = if let Some(Ok((limit, _))) = body.limit.as_ref().map(|l| parse_limit(*l.clone())) {
limit
} else {
None
};
// Parse the LIMIT/OFFSET clause
let (limit, offset) = body
.limit
.as_ref()
.map(|l| parse_limit(l))
.unwrap_or(Ok((None, None)))?;
Ok(Plan::Update(UpdatePlan {
table_references,
set_clauses,
@@ -169,6 +173,7 @@ pub fn prepare_update_plan(schema: &Schema, body: &mut Update) -> crate::Result<
returning: Some(result_columns),
order_by,
limit,
offset,
contains_constant_false_condition: false,
}))
}

27
testing/cli_tests/cli_test_cases.py
View File

@@ -264,6 +264,32 @@ def test_update_with_limit():
limbo.quit()
def test_update_with_limit_and_offset():
limbo = TestLimboShell(
"CREATE TABLE t (a,b,c); insert into t values (1,2,3), (4,5,6), (7,8,9), (1,2,3),(4,5,6), (7,8,9);"
)
limbo.run_test("update-limit-offset", "UPDATE t SET a = 10 LIMIT 1 OFFSET 3;", "")
limbo.run_test(
"update-limit-offset-result", "SELECT COUNT(*) from t WHERE a = 10;", "1"
)
limbo.run_test("update-limit-result", "SELECT a from t LIMIT 4;", "1\n4\n7\n10")
limbo.run_test(
"update-limit-offset-zero", "UPDATE t SET a = 100 LIMIT 0 OFFSET 0;", ""
)
limbo.run_test(
"update-limit-zero-result", "SELECT COUNT(*) from t WHERE a = 100;", "0"
)
limbo.run_test("update-limit-all", "UPDATE t SET a = 100 LIMIT -1 OFFSET 1;", "")
limbo.run_test("update-limit-result", "SELECT COUNT(*) from t WHERE a = 100;", "5")
limbo.run_test(
"udpate-limit-where", "UPDATE t SET a = 333 WHERE b = 5 LIMIT 1 OFFSET 2;", ""
)
limbo.run_test(
"update-limit-where-result", "SELECT COUNT(*) from t WHERE a = 333;", "0"
)
limbo.quit()
def test_insert_default_values():
limbo = TestLimboShell(
"CREATE TABLE t (a integer default(42),b integer default (43),c integer default(44));"
@@ -292,4 +318,5 @@ if __name__ == "__main__":
test_import_csv_skip()
test_table_patterns()
test_update_with_limit()
test_update_with_limit_and_offset()
print("All tests have passed")