diff --git a/core/translate/delete.rs b/core/translate/delete.rs index 1e0d64a98..9652048fe 100644 --- a/core/translate/delete.rs +++ b/core/translate/delete.rs @@ -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 } diff --git a/core/translate/emitter.rs b/core/translate/emitter.rs index e4d05bfaa..fac46fee6 100644 --- a/core/translate/emitter.rs +++ b/core/translate/emitter.rs @@ -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() { diff --git a/core/translate/main_loop.rs b/core/translate/main_loop.rs index e7235e8f0..fd1a8ac9c 100644 --- a/core/translate/main_loop.rs +++ b/core/translate/main_loop.rs @@ -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 { diff --git a/core/translate/plan.rs b/core/translate/plan.rs index 3d93548de..af14f0352 100644 --- a/core/translate/plan.rs +++ b/core/translate/plan.rs @@ -176,8 +176,8 @@ pub struct UpdatePlan { pub set_clauses: Vec<(usize, ast::Expr)>, pub where_clause: Vec, pub order_by: Option>, - // TODO: support OFFSET pub limit: Option, + pub offset: Option, // TODO: optional RETURNING clause pub returning: Option>, // whether the WHERE clause is always false diff --git a/core/translate/planner.rs b/core/translate/planner.rs index 953a15e59..1b78c954c 100644 --- a/core/translate/planner.rs +++ b/core/translate/planner.rs @@ -850,30 +850,33 @@ fn parse_join<'a>( Ok(()) } -pub fn parse_limit(limit: Limit) -> Result<(Option, Option)> { - let offset_val = match limit.offset { +pub fn parse_limit(limit: &Limit) -> Result<(Option, Option)> { + 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::().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::().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::().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") { diff --git a/core/translate/select.rs b/core/translate/select.rs index fe7656f16..bde61880f 100644 --- a/core/translate/select.rs +++ b/core/translate/select.rs @@ -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)) diff --git a/core/translate/update.rs b/core/translate/update.rs index a9388ee13..f282fff24 100644 --- a/core/translate/update.rs +++ b/core/translate/update.rs @@ -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, })) } diff --git a/testing/cli_tests/cli_test_cases.py b/testing/cli_tests/cli_test_cases.py index 2bce1f165..17035b0bf 100755 --- a/testing/cli_tests/cli_test_cases.py +++ b/testing/cli_tests/cli_test_cases.py @@ -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")