diff --git a/core/translate/compound_select.rs b/core/translate/compound_select.rs index d17a28b79..f8d01ada2 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,16 @@ fn read_deduplicated_union_or_except_rows( pc_if_empty: label_dedupe_next, }); program.preassign_label_to_next_insn(label_dedupe_loop_start); + match offset_reg { + Some(reg) if reg > 0 => { + 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 +451,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 +459,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 +482,16 @@ fn read_intersect_rows( record_reg: row_content_reg, num_regs: 0, }); + match offset_reg { + Some(reg) if reg > 0 => { + 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/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/testing/select.test b/testing/select.test index ec434b538..7c163c6b1 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);