support the OFFSET clause for Compound select

This commit is contained in:
meteorgan
2025-07-29 00:41:12 +08:00
parent bac3add778
commit a0f5554b08
3 changed files with 88 additions and 7 deletions

View File

@@ -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<LimitCtx>,
offset_reg: Option<usize>,
yield_reg: Option<usize>,
reg_result_cols_start: Option<usize>,
) -> 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<LimitCtx>,
offset_reg: Option<usize>,
yield_reg: Option<usize>,
) {
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<usize>,
limit_ctx: Option<LimitCtx>,
offset_reg: Option<usize>,
yield_reg: Option<usize>,
) {
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

View File

@@ -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");

View File

@@ -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);