diff --git a/core/translate/compound_select.rs b/core/translate/compound_select.rs index d17a28b79..973c5f27f 100644 --- a/core/translate/compound_select.rs +++ b/core/translate/compound_select.rs @@ -22,6 +22,7 @@ pub fn emit_program_for_compound_select( left: _left, right_most, limit, + offset, .. } = &plan else { @@ -39,8 +40,8 @@ pub fn emit_program_for_compound_select( } } - // Each subselect shares the same limit_ctx, because the LIMIT applies to the entire compound select, - // not just a single subselect. + // Each subselect shares the same limit_ctx and offset, because the LIMIT, OFFSET applies to + // the entire compound select, not just a single subselect. let limit_ctx = limit.map(|limit| { let reg = program.alloc_register(); program.emit_insn(Insn::Integer { @@ -49,6 +50,22 @@ pub fn emit_program_for_compound_select( }); LimitCtx::new_shared(reg) }); + let offset_reg = offset.map(|offset| { + let reg = program.alloc_register(); + program.emit_insn(Insn::Integer { + value: offset as i64, + dest: reg, + }); + + let combined_reg = program.alloc_register(); + program.emit_insn(Insn::OffsetLimit { + offset_reg: reg, + combined_reg, + limit_reg: limit_ctx.unwrap().reg_limit, + }); + + reg + }); // When a compound SELECT is part of a query that yields results to a coroutine (e.g. within an INSERT clause), // we must allocate registers for the result columns to be yielded. Each subselect will then yield to @@ -67,6 +84,7 @@ pub fn emit_program_for_compound_select( schema, syms, limit_ctx, + offset_reg, yield_reg, reg_result_cols_start, )?; @@ -80,12 +98,14 @@ pub fn emit_program_for_compound_select( // Emits bytecode for a compound SELECT statement. This function processes the rightmost part of // the compound SELECT and handles the left parts recursively based on the compound operator type. +#[allow(clippy::too_many_arguments)] fn emit_compound_select( program: &mut ProgramBuilder, plan: Plan, schema: &Schema, syms: &SymbolTable, limit_ctx: Option, + offset_reg: Option, yield_reg: Option, reg_result_cols_start: Option, ) -> crate::Result<()> { @@ -130,6 +150,7 @@ fn emit_compound_select( schema, syms, limit_ctx, + offset_reg, yield_reg, reg_result_cols_start, )?; @@ -144,6 +165,10 @@ fn emit_compound_select( right_most.limit = limit; right_most_ctx.limit_ctx = Some(limit_ctx); } + if offset_reg.is_some() { + right_most.offset = offset; + right_most_ctx.reg_offset = offset_reg; + } emit_query(program, &mut right_most, &mut right_most_ctx)?; program.preassign_label_to_next_insn(label_next_select); } @@ -176,6 +201,7 @@ fn emit_compound_select( schema, syms, None, + None, yield_reg, reg_result_cols_start, )?; @@ -193,6 +219,7 @@ fn emit_compound_select( dedupe_index.0, dedupe_index.1.as_ref(), limit_ctx, + offset_reg, yield_reg, ); } @@ -225,6 +252,7 @@ fn emit_compound_select( schema, syms, None, + None, yield_reg, reg_result_cols_start, )?; @@ -244,6 +272,7 @@ fn emit_compound_select( right_cursor_id, target_cursor_id, limit_ctx, + offset_reg, yield_reg, ); } @@ -276,6 +305,7 @@ fn emit_compound_select( schema, syms, None, + None, yield_reg, reg_result_cols_start, )?; @@ -287,7 +317,7 @@ fn emit_compound_select( emit_query(program, &mut right_most, &mut right_most_ctx)?; if new_index { read_deduplicated_union_or_except_rows( - program, cursor_id, &index, limit_ctx, yield_reg, + program, cursor_id, &index, limit_ctx, offset_reg, yield_reg, ); } } @@ -297,6 +327,10 @@ fn emit_compound_select( right_most_ctx.limit_ctx = Some(limit_ctx); right_most.limit = limit; } + if offset_reg.is_some() { + right_most.offset = offset; + right_most_ctx.reg_offset = offset_reg; + } emit_query(program, &mut right_most, &mut right_most_ctx)?; } } @@ -351,6 +385,7 @@ fn read_deduplicated_union_or_except_rows( dedupe_cursor_id: usize, dedupe_index: &Index, limit_ctx: Option, + offset_reg: Option, yield_reg: Option, ) { let label_close = program.allocate_label(); @@ -362,6 +397,13 @@ fn read_deduplicated_union_or_except_rows( pc_if_empty: label_dedupe_next, }); program.preassign_label_to_next_insn(label_dedupe_loop_start); + if let Some(reg) = offset_reg { + program.emit_insn(Insn::IfPos { + reg, + target_pc: label_dedupe_next, + decrement_by: 1, + }); + } for col_idx in 0..dedupe_index.columns.len() { let start_reg = if let Some(yield_reg) = yield_reg { // Need to reuse the yield_reg for the column being emitted @@ -406,6 +448,7 @@ fn read_deduplicated_union_or_except_rows( } // Emits the bytecode for Reading rows from the intersection of two cursors. +#[allow(clippy::too_many_arguments)] fn read_intersect_rows( program: &mut ProgramBuilder, left_cursor_id: usize, @@ -413,6 +456,7 @@ fn read_intersect_rows( right_cursor_id: usize, target_cursor: Option, limit_ctx: Option, + offset_reg: Option, yield_reg: Option, ) { let label_close = program.allocate_label(); @@ -435,6 +479,13 @@ fn read_intersect_rows( record_reg: row_content_reg, num_regs: 0, }); + if let Some(reg) = offset_reg { + program.emit_insn(Insn::IfPos { + reg, + target_pc: label_next, + decrement_by: 1, + }); + } let column_count = index.columns.len(); let cols_start_reg = if let Some(yield_reg) = yield_reg { yield_reg + 1 diff --git a/core/translate/emitter.rs b/core/translate/emitter.rs index 955427e4b..6a2d90661 100644 --- a/core/translate/emitter.rs +++ b/core/translate/emitter.rs @@ -266,7 +266,7 @@ pub fn emit_query<'a>( t_ctx: &mut TranslateCtx<'a>, ) -> Result { if !plan.values.is_empty() { - let reg_result_cols_start = emit_values(program, plan, &t_ctx.resolver, t_ctx.limit_ctx)?; + let reg_result_cols_start = emit_values(program, plan, t_ctx)?; return Ok(reg_result_cols_start); } diff --git a/core/translate/order_by.rs b/core/translate/order_by.rs index 3fb9ec21a..a2ab74331 100644 --- a/core/translate/order_by.rs +++ b/core/translate/order_by.rs @@ -113,7 +113,7 @@ pub fn emit_order_by( }); program.preassign_label_to_next_insn(sort_loop_start_label); - emit_offset(program, plan, sort_loop_next_label, t_ctx.reg_offset)?; + emit_offset(program, plan, sort_loop_next_label, t_ctx.reg_offset); program.emit_insn(Insn::SorterData { cursor_id: sort_cursor, diff --git a/core/translate/result_row.rs b/core/translate/result_row.rs index 4a7b78890..f2b722988 100644 --- a/core/translate/result_row.rs +++ b/core/translate/result_row.rs @@ -30,7 +30,7 @@ pub fn emit_select_result( limit_ctx: Option, ) -> Result<()> { if let (Some(jump_to), Some(_)) = (offset_jump_to, label_on_limit_reached) { - emit_offset(program, plan, jump_to, reg_offset)?; + emit_offset(program, plan, jump_to, reg_offset); } let start_reg = reg_result_cols_start; @@ -163,7 +163,7 @@ pub fn emit_offset( plan: &SelectPlan, jump_to: BranchOffset, reg_offset: Option, -) -> Result<()> { +) { match plan.offset { Some(offset) if offset > 0 => { program.add_comment(program.offset(), "OFFSET"); @@ -175,5 +175,4 @@ pub fn emit_offset( } _ => {} } - Ok(()) } diff --git a/core/translate/select.rs b/core/translate/select.rs index bec600890..0c4888d45 100644 --- a/core/translate/select.rs +++ b/core/translate/select.rs @@ -154,10 +154,6 @@ pub fn prepare_select_plan( } let (limit, offset) = select.limit.map_or(Ok((None, None)), |l| parse_limit(&l))?; - // FIXME: handle OFFSET for compound selects - if offset.is_some_and(|o| o > 0) { - crate::bail_parse_error!("OFFSET is not supported for compound SELECTs yet"); - } // FIXME: handle ORDER BY for compound selects if select.order_by.is_some() { crate::bail_parse_error!("ORDER BY is not supported for compound SELECTs yet"); diff --git a/core/translate/values.rs b/core/translate/values.rs index 73a33d5eb..8315290e6 100644 --- a/core/translate/values.rs +++ b/core/translate/values.rs @@ -1,6 +1,7 @@ -use crate::translate::emitter::{LimitCtx, Resolver}; +use crate::translate::emitter::{Resolver, TranslateCtx}; use crate::translate::expr::{translate_expr_no_constant_opt, NoConstantOptReason}; use crate::translate::plan::{QueryDestination, SelectPlan}; +use crate::translate::result_row::emit_offset; use crate::vdbe::builder::ProgramBuilder; use crate::vdbe::insn::{IdxInsertFlags, Insn}; use crate::vdbe::BranchOffset; @@ -9,22 +10,19 @@ use crate::Result; pub fn emit_values( program: &mut ProgramBuilder, plan: &SelectPlan, - resolver: &Resolver, - limit_ctx: Option, + t_ctx: &TranslateCtx, ) -> Result { if plan.values.len() == 1 { - let start_reg = emit_values_when_single_row(program, plan, resolver, limit_ctx)?; + let start_reg = emit_values_when_single_row(program, plan, t_ctx)?; return Ok(start_reg); } let reg_result_cols_start = match plan.query_destination { - QueryDestination::ResultRows => emit_toplevel_values(program, plan, resolver, limit_ctx)?, + QueryDestination::ResultRows => emit_toplevel_values(program, plan, t_ctx)?, QueryDestination::CoroutineYield { yield_reg, .. } => { - emit_values_in_subquery(program, plan, resolver, yield_reg)? - } - QueryDestination::EphemeralIndex { .. } => { - emit_toplevel_values(program, plan, resolver, limit_ctx)? + emit_values_in_subquery(program, plan, &t_ctx.resolver, yield_reg)? } + QueryDestination::EphemeralIndex { .. } => emit_toplevel_values(program, plan, t_ctx)?, QueryDestination::EphemeralTable { .. } => unreachable!(), }; Ok(reg_result_cols_start) @@ -33,9 +31,10 @@ pub fn emit_values( fn emit_values_when_single_row( program: &mut ProgramBuilder, plan: &SelectPlan, - resolver: &Resolver, - limit_ctx: Option, + t_ctx: &TranslateCtx, ) -> Result { + let end_label = program.allocate_label(); + emit_offset(program, plan, end_label, t_ctx.reg_offset); let first_row = &plan.values[0]; let row_len = first_row.len(); let start_reg = program.alloc_registers(row_len); @@ -45,12 +44,11 @@ fn emit_values_when_single_row( None, v, start_reg + i, - resolver, + &t_ctx.resolver, NoConstantOptReason::RegisterReuse, )?; } - let end_label = program.allocate_label(); - emit_values_to_destination(program, plan, start_reg, row_len, limit_ctx, end_label); + emit_values_to_destination(program, plan, t_ctx, start_reg, row_len, end_label); program.preassign_label_to_next_insn(end_label); Ok(start_reg) } @@ -58,8 +56,7 @@ fn emit_values_when_single_row( fn emit_toplevel_values( program: &mut ProgramBuilder, plan: &SelectPlan, - resolver: &Resolver, - limit_ctx: Option, + t_ctx: &TranslateCtx, ) -> Result { let yield_reg = program.alloc_register(); let definition_label = program.allocate_label(); @@ -71,7 +68,7 @@ fn emit_toplevel_values( }); program.preassign_label_to_next_insn(start_offset_label); - let start_reg = emit_values_in_subquery(program, plan, resolver, yield_reg)?; + let start_reg = emit_values_in_subquery(program, plan, &t_ctx.resolver, yield_reg)?; program.emit_insn(Insn::EndCoroutine { yield_reg }); program.preassign_label_to_next_insn(definition_label); @@ -82,12 +79,15 @@ fn emit_toplevel_values( start_offset: start_offset_label, }); let end_label = program.allocate_label(); - let goto_label = program.allocate_label(); - program.preassign_label_to_next_insn(goto_label); + let yield_label = program.allocate_label(); + program.preassign_label_to_next_insn(yield_label); program.emit_insn(Insn::Yield { yield_reg, end_offset: end_label, }); + + let goto_label = program.allocate_label(); + emit_offset(program, plan, goto_label, t_ctx.reg_offset); let row_len = plan.values[0].len(); let copy_start_reg = program.alloc_registers(row_len); for i in 0..row_len { @@ -98,10 +98,11 @@ fn emit_toplevel_values( }); } - emit_values_to_destination(program, plan, copy_start_reg, row_len, limit_ctx, end_label); + emit_values_to_destination(program, plan, t_ctx, copy_start_reg, row_len, end_label); + program.preassign_label_to_next_insn(goto_label); program.emit_insn(Insn::Goto { - target_pc: goto_label, + target_pc: yield_label, }); program.preassign_label_to_next_insn(end_label); @@ -139,9 +140,9 @@ fn emit_values_in_subquery( fn emit_values_to_destination( program: &mut ProgramBuilder, plan: &SelectPlan, + t_ctx: &TranslateCtx, start_reg: usize, row_len: usize, - limit_ctx: Option, end_label: BranchOffset, ) { match &plan.query_destination { @@ -150,7 +151,7 @@ fn emit_values_to_destination( start_reg, count: row_len, }); - if let Some(limit_ctx) = limit_ctx { + if let Some(limit_ctx) = t_ctx.limit_ctx { program.emit_insn(Insn::DecrJumpZero { reg: limit_ctx.reg_limit, target_pc: end_label, diff --git a/core/vdbe/explain.rs b/core/vdbe/explain.rs index 05d47bd98..76a969698 100644 --- a/core/vdbe/explain.rs +++ b/core/vdbe/explain.rs @@ -207,7 +207,7 @@ pub fn insn_to_str( "IfPos", *reg as i32, target_pc.as_debug_int(), - 0, + *decrement_by as i32, Value::build_text(""), 0, format!( diff --git a/testing/select.test b/testing/select.test index e9ecf7f73..0a0acbab3 100755 --- a/testing/select.test +++ b/testing/select.test @@ -384,6 +384,24 @@ if {[info exists ::env(SQLITE_EXEC)] && ($::env(SQLITE_EXEC) eq "scripts/limbo-s x|x y|y} + do_execsql_test_on_specific_db {:memory:} select-union-all-with-offset { + CREATE TABLE t (x TEXT, y TEXT); + CREATE TABLE u (x TEXT, y TEXT); + INSERT INTO t VALUES('x','x'),('y','y'); + INSERT INTO u VALUES('x','x'),('y','y'),('z', 'z'); + + select * from t UNION ALL select * from u limit 1 offset 1; + } {y|y} + + do_execsql_test_on_specific_db {:memory:} select-union-with-offset { + CREATE TABLE t (x TEXT, y TEXT); + CREATE TABLE u (x TEXT, y TEXT); + INSERT INTO t VALUES('x','x'),('y','y'); + INSERT INTO u VALUES('x','x'),('y','y'),('z', 'z'); + + select * from t UNION select * from u limit 1 offset 1; + } {y|y} + do_execsql_test_on_specific_db {:memory:} select-intersect-1 { CREATE TABLE t (x TEXT, y TEXT); CREATE TABLE u (x TEXT, y TEXT); @@ -461,6 +479,16 @@ if {[info exists ::env(SQLITE_EXEC)] && ($::env(SQLITE_EXEC) eq "scripts/limbo-s } {x|x y|y} + do_execsql_test_on_specific_db {:memory:} select-intersect-with-offset { + CREATE TABLE t (x TEXT, y TEXT); + CREATE TABLE u (x TEXT, y TEXT); + INSERT INTO t VALUES('x','x'),('y','y'), ('z','z'); + INSERT INTO u VALUES('x','x'),('y','y'), ('z','z'); + + select * from t INTERSECT select * from u limit 2 offset 1; + } {y|y + z|z} + do_execsql_test_on_specific_db {:memory:} select-intersect-union-with-limit { CREATE TABLE t (x TEXT, y TEXT); CREATE TABLE u (x TEXT, y TEXT); @@ -616,7 +644,7 @@ if {[info exists ::env(SQLITE_EXEC)] && ($::env(SQLITE_EXEC) eq "scripts/limbo-s x y} - do_execsql_test_on_specific_db {:memory:} select-values-union-all-limit-2 { + do_execsql_test_on_specific_db {:memory:} select-values-union-all-limit-1 { CREATE TABLE t (x TEXT); INSERT INTO t VALUES('x'), ('y'), ('z'); @@ -624,6 +652,24 @@ if {[info exists ::env(SQLITE_EXEC)] && ($::env(SQLITE_EXEC) eq "scripts/limbo-s } {a b x} + + do_execsql_test_on_specific_db {:memory:} select-values-union-all-offset { + CREATE TABLE t (x TEXT); + INSERT INTO t VALUES('x'), ('y'), ('z'); + + values('a'), ('b') UNION ALL select * from t limit 3 offset 1; + } {b + x + y} + + do_execsql_test_on_specific_db {:memory:} select-values-union-all-offset-1 { + CREATE TABLE t (x TEXT); + INSERT INTO t VALUES('i'), ('j'), ('x'), ('y'), ('z'); + + values('a') UNION ALL select * from t limit 3 offset 1; + } {i + j + x} } do_execsql_test_on_specific_db {:memory:} select-no-match-in-leaf-page { diff --git a/tests/integration/fuzz/mod.rs b/tests/integration/fuzz/mod.rs index 5b669f455..5dd504aff 100644 --- a/tests/integration/fuzz/mod.rs +++ b/tests/integration/fuzz/mod.rs @@ -607,7 +607,13 @@ mod tests { // if the right most SELECT is a VALUES clause, no limit is not allowed if rng.random_bool(0.8) && !has_right_most_values { let limit_val = rng.random_range(0..=MAX_LIMIT_VALUE); // LIMIT 0 is valid - query = format!("{query} LIMIT {limit_val}"); + + if rng.random_bool(0.8) { + query = format!("{query} LIMIT {limit_val}"); + } else { + let offset_val = rng.random_range(0..=MAX_LIMIT_VALUE); + query = format!("{query} LIMIT {limit_val} OFFSET {offset_val}"); + } } log::debug!(