Merge 'Support the OFFSET clause for Compound select' from meteorgan

Closes #2376
This commit is contained in:
Jussi Saurio
2025-08-01 10:18:13 +03:00
committed by GitHub
9 changed files with 137 additions and 38 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,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<usize>,
limit_ctx: Option<LimitCtx>,
offset_reg: Option<usize>,
yield_reg: Option<usize>,
) {
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

View File

@@ -266,7 +266,7 @@ pub fn emit_query<'a>(
t_ctx: &mut TranslateCtx<'a>,
) -> Result<usize> {
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);
}

View File

@@ -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,

View File

@@ -30,7 +30,7 @@ pub fn emit_select_result(
limit_ctx: Option<LimitCtx>,
) -> 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<usize>,
) -> Result<()> {
) {
match plan.offset {
Some(offset) if offset > 0 => {
program.add_comment(program.offset(), "OFFSET");
@@ -175,5 +175,4 @@ pub fn emit_offset(
}
_ => {}
}
Ok(())
}

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

@@ -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<LimitCtx>,
t_ctx: &TranslateCtx,
) -> Result<usize> {
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<LimitCtx>,
t_ctx: &TranslateCtx,
) -> Result<usize> {
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<LimitCtx>,
t_ctx: &TranslateCtx,
) -> Result<usize> {
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<LimitCtx>,
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,

View File

@@ -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!(