Merge 'UNION' from Jussi Saurio

```sql

limbo> select * from products where id between 1 and 3
UNION ALL
select * from products where id between 2 and 4;
┌───┬─────────┬──────┐
│ 1 │ hat     │ 79.0 │
├───┼─────────┼──────┤
│ 2 │ cap     │ 82.0 │
├───┼─────────┼──────┤
│ 3 │ shirt   │ 18.0 │
├───┼─────────┼──────┤
│ 2 │ cap     │ 82.0 │
├───┼─────────┼──────┤
│ 3 │ shirt   │ 18.0 │
├───┼─────────┼──────┤
│ 4 │ sweater │ 25.0 │
└───┴─────────┴──────┘
limbo> select * from products where id between 1 and 3
UNION
select * from products where id between 2 and 4;
┌───┬─────────┬──────┐
│ 1 │ hat     │ 79.0 │
├───┼─────────┼──────┤
│ 2 │ cap     │ 82.0 │
├───┼─────────┼──────┤
│ 3 │ shirt   │ 18.0 │
├───┼─────────┼──────┤
│ 4 │ sweater │ 25.0 │
└───┴─────────┴──────┘
limbo>
```
Similarly as UNION ALL (#1541 ), supports LIMIT but not OFFSET or ORDER
BY.
Augments `compound_select_fuzz()` to work with both UNION and UNION ALL

Closes #1545

core/translate/emitter.rs

@@ -4,7 +4,7 @@
 use std::rc::Rc;
 use std::sync::Arc;
 
-use limbo_sqlite3_parser::ast::{self};
+use limbo_sqlite3_parser::ast::{self, SortOrder};
 
 use super::aggregation::emit_ungrouped_aggregation;
 use super::expr::{translate_condition_expr, translate_expr, ConditionMetadata};
@@ -15,13 +15,15 @@ use super::main_loop::{
     close_loop, emit_loop, init_distinct, init_loop, open_loop, LeftJoinMetadata, LoopLabels,
 };
 use super::order_by::{emit_order_by, init_order_by, SortMetadata};
-use super::plan::{JoinOrderMember, Operation, SelectPlan, TableReference, UpdatePlan};
+use super::plan::{
+    JoinOrderMember, Operation, QueryDestination, SelectPlan, TableReference, UpdatePlan,
+};
 use super::schema::ParseSchema;
 use super::select::emit_simple_count;
 use super::subquery::emit_subqueries;
 use crate::error::SQLITE_CONSTRAINT_PRIMARYKEY;
 use crate::function::Func;
-use crate::schema::Index;
+use crate::schema::{Index, IndexColumn};
 use crate::translate::plan::{DeletePlan, Plan, Search};
 use crate::translate::values::emit_values;
 use crate::util::exprs_are_equivalent;
@@ -208,67 +210,134 @@ fn emit_program_for_compound_select(
 
     // Each subselect gets their own TranslateCtx, but they share the same limit_ctx
     // because the LIMIT applies to the entire compound select, not just a single subselect.
-    let mut t_ctx_list = Vec::with_capacity(rest.len() + 1);
-    let reg_limit = if let Some(limit) = limit {
+    // The way LIMIT works with compound selects is:
+    // - If a given subselect appears BEFORE any UNION, then do NOT count those rows towards the LIMIT,
+    //   because the rows from those subselects need to be deduplicated before they start being counted.
+    // - If a given subselect appears AFTER the last UNION, then count those rows towards the LIMIT immediately.
+    let limit_ctx = limit.map(|limit| {
         let reg = program.alloc_register();
         program.emit_insn(Insn::Integer {
             value: limit as i64,
             dest: reg,
         });
-        Some(reg)
-    } else {
-        None
-    };
-    let limit_ctx = if let Some(reg_limit) = reg_limit {
-        Some(LimitCtx::new_shared(reg_limit))
-    } else {
-        None
-    };
-    let mut t_ctx_first = TranslateCtx::new(
+        LimitCtx::new_shared(reg)
+    });
+
+    // Each subselect gets their own TranslateCtx.
+    let mut t_ctx_list = Vec::with_capacity(rest.len() + 1);
+    t_ctx_list.push(TranslateCtx::new(
         program,
         syms,
         first.table_references.len(),
         first.result_columns.len(),
-    );
-    t_ctx_first.limit_ctx = limit_ctx;
-    t_ctx_list.push(t_ctx_first);
-
-    for (select, _) in rest.iter() {
-        let mut t_ctx = TranslateCtx::new(
+    ));
+    rest.iter().for_each(|(select, _)| {
+        let t_ctx = TranslateCtx::new(
             program,
             syms,
             select.table_references.len(),
             select.result_columns.len(),
         );
-        t_ctx.limit_ctx = limit_ctx;
         t_ctx_list.push(t_ctx);
+    });
+
+    // Compound select operators have the same precedence and are left-associative.
+    // If there is any remaining UNION operator on the right side of a given sub-SELECT,
+    // all of the rows from the preceding UNION arms need to be deduplicated.
+    // This is done by creating an ephemeral index and inserting all the rows from the left side of
+    // the last UNION arm into it.
+    // Then, as soon as there are no more UNION operators left, all the deduplicated rows from the
+    // ephemeral index are emitted, and lastly the rows from the remaining sub-SELECTS are emitted
+    // as is, as they don't require deduplication.
+    let mut first_t_ctx = t_ctx_list.remove(0);
+    let requires_union_deduplication = rest
+        .iter()
+        .any(|(_, operator)| operator == &ast::CompoundOperator::Union);
+    if requires_union_deduplication {
+        // appears BEFORE a UNION operator, so do not count those rows towards the LIMIT.
+        first.limit = None;
+    } else {
+        // appears AFTER the last UNION operator, so count those rows towards the LIMIT.
+        first_t_ctx.limit_ctx = limit_ctx;
     }
 
-    let mut first_t_ctx = t_ctx_list.remove(0);
+    let mut union_dedupe_index = if requires_union_deduplication {
+        let dedupe_index = get_union_dedupe_index(program, &first);
+        first.query_destination = QueryDestination::EphemeralIndex {
+            cursor_id: dedupe_index.0,
+            index: dedupe_index.1.clone(),
+        };
+        Some(dedupe_index)
+    } else {
+        None
+    };
+
+    // Emit the first SELECT
     emit_query(program, &mut first, &mut first_t_ctx)?;
 
-    // TODO: add support for UNION, EXCEPT, INTERSECT
+    // Emit the remaining SELECTs. Any selects on the left side of a UNION must deduplicate their
+    // results with the ephemeral index created above.
     while !t_ctx_list.is_empty() {
         let label_next_select = program.allocate_label();
         // If the LIMIT is reached in any subselect, jump to either:
         // a) the IfNot of the next subselect, or
         // b) the end of the program
-        if let Some(reg_limit) = reg_limit {
+        if let Some(limit_ctx) = limit_ctx {
             program.emit_insn(Insn::IfNot {
-                reg: reg_limit,
+                reg: limit_ctx.reg_limit,
                 target_pc: label_next_select,
                 jump_if_null: true,
             });
         }
         let mut t_ctx = t_ctx_list.remove(0);
+        let requires_union_deduplication = rest
+            .iter()
+            .any(|(_, operator)| operator == &ast::CompoundOperator::Union);
         let (mut select, operator) = rest.remove(0);
-        if operator != ast::CompoundOperator::UnionAll {
+        if operator != ast::CompoundOperator::UnionAll && operator != ast::CompoundOperator::Union {
             crate::bail_parse_error!("unimplemented compound select operator: {:?}", operator);
         }
+
+        if requires_union_deduplication {
+            // Again: appears BEFORE a UNION operator, so do not count those rows towards the LIMIT.
+            select.limit = None;
+        } else {
+            // appears AFTER the last UNION operator, so count those rows towards the LIMIT.
+            t_ctx.limit_ctx = limit_ctx;
+        }
+
+        if requires_union_deduplication {
+            select.query_destination = QueryDestination::EphemeralIndex {
+                cursor_id: union_dedupe_index.as_ref().unwrap().0,
+                index: union_dedupe_index.as_ref().unwrap().1.clone(),
+            };
+        } else if let Some((dedupe_cursor_id, dedupe_index)) = union_dedupe_index.take() {
+            // When there are no more UNION operators left, all the deduplicated rows from the preceding union arms need to be emitted
+            // as result rows.
+            read_deduplicated_union_rows(
+                program,
+                dedupe_cursor_id,
+                dedupe_index.as_ref(),
+                limit_ctx,
+                label_next_select,
+            );
+        }
         emit_query(program, &mut select, &mut t_ctx)?;
         program.preassign_label_to_next_insn(label_next_select);
     }
 
+    if let Some((dedupe_cursor_id, dedupe_index)) = union_dedupe_index {
+        let label_jump_over_dedupe = program.allocate_label();
+        read_deduplicated_union_rows(
+            program,
+            dedupe_cursor_id,
+            dedupe_index.as_ref(),
+            limit_ctx,
+            label_jump_over_dedupe,
+        );
+        program.preassign_label_to_next_insn(label_jump_over_dedupe);
+    }
+
     program.epilogue(TransactionMode::Read);
     program.result_columns = first.result_columns;
     program.table_references = first.table_references;
@@ -276,6 +345,84 @@ fn emit_program_for_compound_select(
     Ok(())
 }
 
+/// Creates an ephemeral index that will be used to deduplicate the results of any sub-selects
+/// that appear before the last UNION operator.
+fn get_union_dedupe_index(
+    program: &mut ProgramBuilder,
+    first_select_in_compound: &SelectPlan,
+) -> (usize, Arc<Index>) {
+    let dedupe_index = Arc::new(Index {
+        columns: first_select_in_compound
+            .result_columns
+            .iter()
+            .map(|c| IndexColumn {
+                name: c
+                    .name(&first_select_in_compound.table_references)
+                    .map(|n| n.to_string())
+                    .unwrap_or_default(),
+                order: SortOrder::Asc,
+                pos_in_table: 0,
+                collation: None, // FIXME: this should be inferred
+            })
+            .collect(),
+        name: "union_dedupe".to_string(),
+        root_page: 0,
+        ephemeral: true,
+        table_name: String::new(),
+        unique: true,
+        has_rowid: false,
+    });
+    let cursor_id = program.alloc_cursor_id(
+        Some(dedupe_index.name.clone()),
+        CursorType::BTreeIndex(dedupe_index.clone()),
+    );
+    program.emit_insn(Insn::OpenEphemeral {
+        cursor_id,
+        is_table: false,
+    });
+    (cursor_id, dedupe_index.clone())
+}
+
+/// Emits the bytecode for reading deduplicated rows from the ephemeral index created for UNION operators.
+fn read_deduplicated_union_rows(
+    program: &mut ProgramBuilder,
+    dedupe_cursor_id: usize,
+    dedupe_index: &Index,
+    limit_ctx: Option<LimitCtx>,
+    label_limit_reached: BranchOffset,
+) {
+    let label_dedupe_next = program.allocate_label();
+    let label_dedupe_loop_start = program.allocate_label();
+    let dedupe_cols_start_reg = program.alloc_registers(dedupe_index.columns.len());
+    program.emit_insn(Insn::Rewind {
+        cursor_id: dedupe_cursor_id,
+        pc_if_empty: label_dedupe_next,
+    });
+    program.preassign_label_to_next_insn(label_dedupe_loop_start);
+    for col_idx in 0..dedupe_index.columns.len() {
+        program.emit_insn(Insn::Column {
+            cursor_id: dedupe_cursor_id,
+            column: col_idx,
+            dest: dedupe_cols_start_reg + col_idx,
+        });
+    }
+    program.emit_insn(Insn::ResultRow {
+        start_reg: dedupe_cols_start_reg,
+        count: dedupe_index.columns.len(),
+    });
+    if let Some(limit_ctx) = limit_ctx {
+        program.emit_insn(Insn::DecrJumpZero {
+            reg: limit_ctx.reg_limit,
+            target_pc: label_limit_reached,
+        })
+    }
+    program.preassign_label_to_next_insn(label_dedupe_next);
+    program.emit_insn(Insn::Next {
+        cursor_id: dedupe_cursor_id,
+        pc_if_next: label_dedupe_loop_start,
+    });
+}
+
 fn emit_program_for_select(
     program: &mut ProgramBuilder,
     mut plan: SelectPlan,

core/translate/main_loop.rs

@@ -30,7 +30,7 @@ use super::{
     order_by::{order_by_sorter_insert, sorter_insert},
     plan::{
         convert_where_to_vtab_constraint, Aggregate, GroupBy, IterationDirection, JoinOrderMember,
-        Operation, Search, SeekDef, SelectPlan, SelectQueryType, TableReference, WhereTerm,
+        Operation, QueryDestination, Search, SeekDef, SelectPlan, TableReference, WhereTerm,
     },
 };
 
@@ -466,8 +466,8 @@ pub fn open_loop(
                     }
                     Table::FromClauseSubquery(from_clause_subquery) => {
                         let (yield_reg, coroutine_implementation_start) =
-                            match &from_clause_subquery.plan.query_type {
-                                SelectQueryType::Subquery {
+                            match &from_clause_subquery.plan.query_destination {
+                                QueryDestination::CoroutineYield {
                                     yield_reg,
                                     coroutine_implementation_start,
                                 } => (*yield_reg, *coroutine_implementation_start),

core/translate/plan.rs

@@ -309,16 +309,30 @@ pub enum Plan {
     Update(UpdatePlan),
 }
 
-/// The type of the query, either top level or subquery
+/// The destination of the results of a query.
+/// Typically, the results of a query are returned to the caller.
+/// However, there are some cases where the results are not returned to the caller,
+/// but rather are yielded to a parent query via coroutine, or stored in a temp table,
+/// later used by the parent query.
 #[derive(Debug, Clone)]
-pub enum SelectQueryType {
-    TopLevel,
-    Subquery {
-        /// The register that holds the program offset that handles jumping to/from the subquery.
+pub enum QueryDestination {
+    /// The results of the query are returned to the caller.
+    ResultRows,
+    /// The results of the query are yielded to a parent query via coroutine.
+    CoroutineYield {
+        /// The register that holds the program offset that handles jumping to/from the coroutine.
         yield_reg: usize,
-        /// The index of the first instruction in the bytecode that implements the subquery.
+        /// The index of the first instruction in the bytecode that implements the coroutine.
         coroutine_implementation_start: BranchOffset,
     },
+    /// The results of the query are stored in an ephemeral index,
+    /// later used by the parent query.
+    EphemeralIndex {
+        /// The cursor ID of the ephemeral index that will be used to store the results.
+        cursor_id: CursorID,
+        /// The index that will be used to store the results.
+        index: Arc<Index>,
+    },
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -431,8 +445,8 @@ pub struct SelectPlan {
     pub offset: Option<isize>,
     /// query contains a constant condition that is always false
     pub contains_constant_false_condition: bool,
-    /// query type (top level or subquery)
-    pub query_type: SelectQueryType,
+    /// the destination of the resulting rows from this plan.
+    pub query_destination: QueryDestination,
     /// whether the query is DISTINCT
     pub distinctness: Distinctness,
     /// values: https://sqlite.org/syntax/select-core.html
@@ -479,7 +493,10 @@ impl SelectPlan {
     pub fn is_simple_count(&self) -> bool {
         if !self.where_clause.is_empty()
             || self.aggregates.len() != 1
-            || matches!(self.query_type, SelectQueryType::Subquery { .. })
+            || matches!(
+                self.query_destination,
+                QueryDestination::CoroutineYield { .. }
+            )
             || self.table_references.len() != 1
             || self.result_columns.len() != 1
             || self.group_by.is_some()

core/translate/planner.rs

@@ -2,7 +2,7 @@ use super::{
     expr::walk_expr,
     plan::{
         Aggregate, ColumnUsedMask, Distinctness, EvalAt, IterationDirection, JoinInfo,
-        JoinOrderMember, Operation, Plan, ResultSetColumn, SelectPlan, SelectQueryType,
+        JoinOrderMember, Operation, Plan, QueryDestination, ResultSetColumn, SelectPlan,
         TableReference, WhereTerm,
     },
     select::prepare_select_plan,
@@ -298,7 +298,7 @@ fn parse_from_clause_table<'a>(
             else {
                 crate::bail_parse_error!("Only non-compound SELECT queries are currently supported in FROM clause subqueries");
             };
-            subplan.query_type = SelectQueryType::Subquery {
+            subplan.query_destination = QueryDestination::CoroutineYield {
                 yield_reg: usize::MAX, // will be set later in bytecode emission
                 coroutine_implementation_start: BranchOffset::Placeholder, // will be set later in bytecode emission
             };
@@ -455,8 +455,7 @@ pub fn parse_from<'a>(
             let Plan::Select(mut cte_plan) = cte_plan else {
                 crate::bail_parse_error!("Only SELECT queries are currently supported in CTEs");
             };
-            // CTE can be rewritten as a subquery.
-            cte_plan.query_type = SelectQueryType::Subquery {
+            cte_plan.query_destination = QueryDestination::CoroutineYield {
                 yield_reg: usize::MAX, // will be set later in bytecode emission
                 coroutine_implementation_start: BranchOffset::Placeholder, // will be set later in bytecode emission
             };

core/translate/result_row.rs

@@ -1,12 +1,16 @@
 use crate::{
-    vdbe::{builder::ProgramBuilder, insn::Insn, BranchOffset},
+    vdbe::{
+        builder::ProgramBuilder,
+        insn::{IdxInsertFlags, Insn},
+        BranchOffset,
+    },
     Result,
 };
 
 use super::{
     emitter::{LimitCtx, Resolver},
     expr::translate_expr,
-    plan::{Distinctness, SelectPlan, SelectQueryType},
+    plan::{Distinctness, QueryDestination, SelectPlan},
 };
 
 /// Emits the bytecode for:
@@ -81,14 +85,33 @@ pub fn emit_result_row_and_limit(
     reg_limit_offset_sum: Option<usize>,
     label_on_limit_reached: Option<BranchOffset>,
 ) -> Result<()> {
-    match &plan.query_type {
-        SelectQueryType::TopLevel => {
+    match &plan.query_destination {
+        QueryDestination::ResultRows => {
             program.emit_insn(Insn::ResultRow {
                 start_reg: result_columns_start_reg,
                 count: plan.result_columns.len(),
             });
         }
-        SelectQueryType::Subquery { yield_reg, .. } => {
+        QueryDestination::EphemeralIndex {
+            cursor_id: index_cursor_id,
+            index: dedupe_index,
+        } => {
+            let record_reg = program.alloc_register();
+            program.emit_insn(Insn::MakeRecord {
+                start_reg: result_columns_start_reg,
+                count: plan.result_columns.len(),
+                dest_reg: record_reg,
+                index_name: Some(dedupe_index.name.clone()),
+            });
+            program.emit_insn(Insn::IdxInsert {
+                cursor_id: *index_cursor_id,
+                record_reg,
+                unpacked_start: None,
+                unpacked_count: None,
+                flags: IdxInsertFlags::new(),
+            });
+        }
+        QueryDestination::CoroutineYield { yield_reg, .. } => {
             program.emit_insn(Insn::Yield {
                 yield_reg: *yield_reg,
                 end_offset: BranchOffset::Offset(0),

core/translate/select.rs

@@ -1,5 +1,7 @@
 use super::emitter::{emit_program, TranslateCtx};
-use super::plan::{select_star, Distinctness, JoinOrderMember, Operation, Search, SelectQueryType};
+use super::plan::{
+    select_star, Distinctness, JoinOrderMember, Operation, QueryDestination, Search,
+};
 use super::planner::Scope;
 use crate::function::{AggFunc, ExtFunc, Func};
 use crate::schema::Table;
@@ -100,9 +102,13 @@ pub fn prepare_select_plan<'a>(
             )?;
             let mut rest = Vec::with_capacity(compounds.len());
             for CompoundSelect { select, operator } in compounds {
-                // TODO: add support for UNION, EXCEPT and INTERSECT
-                if operator != ast::CompoundOperator::UnionAll {
-                    crate::bail_parse_error!("only UNION ALL is supported for compound SELECTs");
+                // TODO: add support for EXCEPT and INTERSECT
+                if operator != ast::CompoundOperator::UnionAll
+                    && operator != ast::CompoundOperator::Union
+                {
+                    crate::bail_parse_error!(
+                        "only UNION ALL and UNION are supported for compound SELECTs"
+                    );
                 }
                 let plan = prepare_one_select_plan(
                     schema,
@@ -231,7 +237,7 @@ fn prepare_one_select_plan<'a>(
                 limit: None,
                 offset: None,
                 contains_constant_false_condition: false,
-                query_type: SelectQueryType::TopLevel,
+                query_destination: QueryDestination::ResultRows,
                 distinctness: Distinctness::from_ast(distinctness.as_ref()),
                 values: vec![],
             };
@@ -545,7 +551,7 @@ fn prepare_one_select_plan<'a>(
                 limit: None,
                 offset: None,
                 contains_constant_false_condition: false,
-                query_type: SelectQueryType::TopLevel,
+                query_destination: QueryDestination::ResultRows,
                 distinctness: Distinctness::NonDistinct,
                 values,
             };

core/translate/subquery.rs

@@ -7,7 +7,7 @@ use crate::{
 use super::{
     emitter::{emit_query, LimitCtx, Resolver, TranslateCtx},
     main_loop::LoopLabels,
-    plan::{SelectPlan, SelectQueryType, TableReference},
+    plan::{QueryDestination, SelectPlan, TableReference},
 };
 
 /// Emit the subqueries contained in the FROM clause.
@@ -51,8 +51,8 @@ pub fn emit_subquery<'a>(
 ) -> Result<usize> {
     let yield_reg = program.alloc_register();
     let coroutine_implementation_start_offset = program.allocate_label();
-    match &mut plan.query_type {
-        SelectQueryType::Subquery {
+    match &mut plan.query_destination {
+        QueryDestination::CoroutineYield {
             yield_reg: y,
             coroutine_implementation_start,
         } => {

core/translate/values.rs

@@ -1,6 +1,6 @@
 use crate::translate::emitter::Resolver;
 use crate::translate::expr::{translate_expr_no_constant_opt, NoConstantOptReason};
-use crate::translate::plan::{SelectPlan, SelectQueryType};
+use crate::translate::plan::{QueryDestination, SelectPlan};
 use crate::vdbe::builder::ProgramBuilder;
 use crate::vdbe::insn::Insn;
 use crate::vdbe::BranchOffset;
@@ -16,11 +16,12 @@ pub fn emit_values(
         return Ok(start_reg);
     }
 
-    let reg_result_cols_start = match plan.query_type {
-        SelectQueryType::TopLevel => emit_toplevel_values(program, plan, resolver)?,
-        SelectQueryType::Subquery { yield_reg, .. } => {
+    let reg_result_cols_start = match plan.query_destination {
+        QueryDestination::ResultRows => emit_toplevel_values(program, plan, resolver)?,
+        QueryDestination::CoroutineYield { yield_reg, .. } => {
             emit_values_in_subquery(program, plan, resolver, yield_reg)?
         }
+        QueryDestination::EphemeralIndex { .. } => unreachable!(),
     };
     Ok(reg_result_cols_start)
 }
@@ -43,19 +44,20 @@ fn emit_values_when_single_row(
             NoConstantOptReason::RegisterReuse,
         )?;
     }
-    match plan.query_type {
-        SelectQueryType::TopLevel => {
+    match plan.query_destination {
+        QueryDestination::ResultRows => {
             program.emit_insn(Insn::ResultRow {
                 start_reg,
                 count: row_len,
             });
         }
-        SelectQueryType::Subquery { yield_reg, .. } => {
+        QueryDestination::CoroutineYield { yield_reg, .. } => {
             program.emit_insn(Insn::Yield {
                 yield_reg,
                 end_offset: BranchOffset::Offset(0),
             });
         }
+        QueryDestination::EphemeralIndex { .. } => unreachable!(),
     }
     Ok(start_reg)
 }

testing/select.test

@@ -36,7 +36,7 @@ do_execsql_test select-blob-empty {
 } {}
 
 do_execsql_test select-blob-ascii {
-  SELECT x'6C696D626f';
+  SELECT x'6C696D626F';
 } {limbo}
 
 do_execsql_test select-blob-emoji {
@@ -285,3 +285,76 @@ do_execsql_test_on_specific_db {:memory:} select-union-all-with-filters {
 6
 10}
 
+do_execsql_test_on_specific_db {:memory:} select-union-1 {
+  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');
+
+  select * from t UNION select * from u;
+} {x|x
+y|y}
+
+do_execsql_test_on_specific_db {:memory:} select-union-all-union {
+  CREATE TABLE t(x TEXT, y TEXT);
+  CREATE TABLE u(x TEXT, y TEXT);
+  CREATE TABLE v(x TEXT, y TEXT);
+  INSERT INTO t VALUES('x','x'),('y','y');
+  INSERT INTO u VALUES('x','x'),('y','y');
+  INSERT INTO v VALUES('x','x'),('y','y');
+
+  select * from t UNION select * from u UNION ALL select * from v;
+} {x|x
+y|y
+x|x
+y|y}
+
+do_execsql_test_on_specific_db {:memory:} select-union-all-union-2 {
+  CREATE TABLE t(x TEXT, y TEXT);
+  CREATE TABLE u(x TEXT, y TEXT);
+  CREATE TABLE v(x TEXT, y TEXT);
+  INSERT INTO t VALUES('x','x'),('y','y');
+  INSERT INTO u VALUES('x','x'),('y','y');
+  INSERT INTO v VALUES('x','x'),('y','y');
+
+  select * from t UNION ALL select * from u UNION select * from v;
+} {x|x
+y|y}
+
+do_execsql_test_on_specific_db {:memory:} select-union-3 {
+  CREATE TABLE t(x TEXT, y TEXT);
+  CREATE TABLE u(x TEXT, y TEXT);
+  CREATE TABLE v(x TEXT, y TEXT);
+  INSERT INTO t VALUES('x','x'),('y','y');
+  INSERT INTO u VALUES('x','x'),('y','y');
+  INSERT INTO v VALUES('x','x'),('y','y');
+
+  select * from t UNION select * from u UNION select * from v;
+} {x|x
+y|y}
+
+do_execsql_test_on_specific_db {:memory:} select-union-4 {
+  CREATE TABLE t(x TEXT, y TEXT);
+  CREATE TABLE u(x TEXT, y TEXT);
+  CREATE TABLE v(x TEXT, y TEXT);
+  INSERT INTO t VALUES('x','x'),('y','y');
+  INSERT INTO u VALUES('x','x'),('y','y');
+  INSERT INTO v VALUES('x','x'),('y','y');
+
+  select * from t UNION select * from u UNION select * from v UNION select * from t;
+} {x|x
+y|y}
+
+do_execsql_test_on_specific_db {:memory:} select-union-all-union-3 {
+  CREATE TABLE t(x TEXT, y TEXT);
+  CREATE TABLE u(x TEXT, y TEXT);
+  CREATE TABLE v(x TEXT, y TEXT);
+  INSERT INTO t VALUES('x','x'),('y','y');
+  INSERT INTO u VALUES('x','x'),('y','y');
+  INSERT INTO v VALUES('x','x'),('y','y');
+
+  select * from t UNION select * from u UNION select * from v UNION ALL select * from t;
+} {x|x
+y|y
+x|x
+y|y}

tests/integration/fuzz/mod.rs

@@ -424,11 +424,11 @@ mod tests {
         log::info!("compound_select_fuzz seed: {}", seed);
 
         // Constants for fuzzing parameters
-        const MAX_TABLES: usize = 5;
+        const MAX_TABLES: usize = 7;
         const MIN_TABLES: usize = 1;
-        const MAX_ROWS_PER_TABLE: usize = 15;
+        const MAX_ROWS_PER_TABLE: usize = 40;
         const MIN_ROWS_PER_TABLE: usize = 5;
-        const NUM_FUZZ_ITERATIONS: usize = 1000;
+        const NUM_FUZZ_ITERATIONS: usize = 2000;
         // How many more SELECTs than tables can be in a UNION (e.g., if 2 tables, max 2+2=4 SELECTs)
         const MAX_SELECTS_IN_UNION_EXTRA: usize = 2;
         const MAX_LIMIT_VALUE: usize = 50;
@@ -440,12 +440,16 @@ mod tests {
         let mut table_names = Vec::new();
         let num_tables = rng.random_range(MIN_TABLES..=MAX_TABLES);
 
+        const COLS: [&str; 3] = ["c1", "c2", "c3"];
         for i in 0..num_tables {
             let table_name = format!("t{}", i);
-            // Schema: c1 INTEGER, c2 INTEGER, c3 INTEGER for simplicity and UNION ALL compatibility
             let create_table_sql = format!(
-                "CREATE TABLE {} (c1 INTEGER, c2 INTEGER, c3 INTEGER)",
-                table_name
+                "CREATE TABLE {} ({})",
+                table_name,
+                COLS.iter()
+                    .map(|c| format!("{} INTEGER", c))
+                    .collect::<Vec<_>>()
+                    .join(", ")
             );
 
             limbo_exec_rows(&db, &limbo_conn, &create_table_sql);
@@ -453,9 +457,9 @@ mod tests {
 
             let num_rows_to_insert = rng.random_range(MIN_ROWS_PER_TABLE..=MAX_ROWS_PER_TABLE);
             for _ in 0..num_rows_to_insert {
-                let c1_val: i64 = rng.random_range(-1000..1000);
-                let c2_val: i64 = rng.random_range(-1000..1000);
-                let c3_val: i64 = rng.random_range(-1000..1000);
+                let c1_val: i64 = rng.random_range(-3..3);
+                let c2_val: i64 = rng.random_range(-3..3);
+                let c3_val: i64 = rng.random_range(-3..3);
 
                 let insert_sql = format!(
                     "INSERT INTO {} VALUES ({}, {}, {})",
@@ -468,18 +472,37 @@ mod tests {
         }
 
         for iter_num in 0..NUM_FUZZ_ITERATIONS {
-            // Number of SELECT clauses to be UNION ALL'd
+            // Number of SELECT clauses
             let num_selects_in_union =
                 rng.random_range(1..=(table_names.len() + MAX_SELECTS_IN_UNION_EXTRA));
             let mut select_statements = Vec::new();
 
+            // Randomly pick a subset of columns to select from
+            let num_cols_to_select = rng.random_range(1..=COLS.len());
+            let cols_to_select = COLS
+                .choose_multiple(&mut rng, num_cols_to_select)
+                .map(|c| c.to_string())
+                .collect::<Vec<_>>();
+
             for _ in 0..num_selects_in_union {
                 // Randomly pick a table
                 let table_to_select_from = &table_names[rng.random_range(0..table_names.len())];
-                select_statements.push(format!("SELECT c1, c2, c3 FROM {}", table_to_select_from));
+                select_statements.push(format!(
+                    "SELECT {} FROM {}",
+                    cols_to_select.join(", "),
+                    table_to_select_from
+                ));
             }
 
-            let mut query = select_statements.join(" UNION ALL ");
+            const COMPOUND_OPERATORS: [&str; 2] = [" UNION ALL ", " UNION "];
+
+            let mut query = String::new();
+            for (i, select_statement) in select_statements.iter().enumerate() {
+                if i > 0 {
+                    query.push_str(COMPOUND_OPERATORS.choose(&mut rng).unwrap());
+                }
+                query.push_str(select_statement);
+            }
 
             if rng.random_bool(0.8) {
                 let limit_val = rng.random_range(0..=MAX_LIMIT_VALUE); // LIMIT 0 is valid
@@ -497,9 +520,15 @@ mod tests {
             let sqlite_results = sqlite_exec_rows(&sqlite_conn, &query);
 
             assert_eq!(
-                limbo_results, sqlite_results,
-                "query: {}, limbo: {:?}, sqlite: {:?}, seed: {}",
-                query, limbo_results, sqlite_results, seed
+                limbo_results,
+                sqlite_results,
+                "query: {}, limbo.len(): {}, sqlite.len(): {}, limbo: {:?}, sqlite: {:?}, seed: {}",
+                query,
+                limbo_results.len(),
+                sqlite_results.len(),
+                limbo_results,
+                sqlite_results,
+                seed
             );
         }
     }