UNION ALL

core/translate/emitter.rs

@@ -152,9 +152,72 @@ pub fn emit_program(program: &mut ProgramBuilder, plan: Plan, syms: &SymbolTable
         Plan::Select(plan) => emit_program_for_select(program, plan, syms),
         Plan::Delete(plan) => emit_program_for_delete(program, plan, syms),
         Plan::Update(plan) => emit_program_for_update(program, plan, syms),
+        Plan::CompoundSelect { .. } => emit_program_for_compound_select(program, plan, syms),
     }
 }
 
+fn emit_program_for_compound_select(
+    program: &mut ProgramBuilder,
+    plan: Plan,
+    syms: &SymbolTable,
+) -> Result<()> {
+    let Plan::CompoundSelect {
+        mut first,
+        mut rest,
+        limit,
+        ..
+    } = plan
+    else {
+        crate::bail_parse_error!("expected compound select plan");
+    };
+
+    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(),
+    ));
+
+    for (select, _) in rest.iter() {
+        t_ctx_list.push(TranslateCtx::new(
+            program,
+            syms,
+            select.table_references.len(),
+            select.result_columns.len(),
+        ));
+    }
+
+    // Trivial exit on LIMIT 0
+    if let Some(limit) = limit {
+        if limit == 0 {
+            program.epilogue(TransactionMode::Read);
+            program.result_columns = first.result_columns;
+            program.table_references = first.table_references;
+            return Ok(());
+        }
+    }
+
+    let mut first_t_ctx = t_ctx_list.remove(0);
+    emit_query(program, &mut first, &mut first_t_ctx)?;
+
+    // TODO: add support for UNION, EXCEPT, INTERSECT
+    while !t_ctx_list.is_empty() {
+        let mut t_ctx = t_ctx_list.remove(0);
+        let (mut select, operator) = rest.remove(0);
+        if operator != ast::CompoundOperator::UnionAll {
+            crate::bail_parse_error!("unimplemented compound select operator: {:?}", operator);
+        }
+        emit_query(program, &mut select, &mut t_ctx)?;
+    }
+
+    program.epilogue(TransactionMode::Read);
+    program.result_columns = first.result_columns;
+    program.table_references = first.table_references;
+
+    Ok(())
+}
+
 fn emit_program_for_select(
     program: &mut ProgramBuilder,
     mut plan: SelectPlan,

core/translate/optimizer/mod.rs

@@ -37,6 +37,13 @@ pub fn optimize_plan(plan: &mut Plan, schema: &Schema) -> Result<()> {
         Plan::Select(plan) => optimize_select_plan(plan, schema),
         Plan::Delete(plan) => optimize_delete_plan(plan, schema),
         Plan::Update(plan) => optimize_update_plan(plan, schema),
+        Plan::CompoundSelect { first, rest, .. } => {
+            optimize_select_plan(first, schema)?;
+            for (plan, _) in rest {
+                optimize_select_plan(plan, schema)?;
+            }
+            Ok(())
+        }
     }
 }
 

core/translate/plan.rs

@@ -264,6 +264,13 @@ impl Ord for EvalAt {
 #[derive(Debug, Clone)]
 pub enum Plan {
     Select(SelectPlan),
+    CompoundSelect {
+        first: SelectPlan,
+        rest: Vec<(SelectPlan, ast::CompoundOperator)>,
+        limit: Option<isize>,
+        offset: Option<isize>,
+        order_by: Option<Vec<(ast::Expr, SortOrder)>>,
+    },
     Delete(DeletePlan),
     Update(UpdatePlan),
 }
@@ -909,6 +916,41 @@ impl Display for Plan {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         match self {
             Self::Select(select_plan) => select_plan.fmt(f),
+            Self::CompoundSelect {
+                first,
+                rest,
+                limit,
+                offset,
+                order_by,
+            } => {
+                first.fmt(f)?;
+                for (plan, operator) in rest {
+                    writeln!(f, "{}", operator)?;
+                    plan.fmt(f)?;
+                }
+                if let Some(limit) = limit {
+                    writeln!(f, "LIMIT: {}", limit)?;
+                }
+                if let Some(offset) = offset {
+                    writeln!(f, "OFFSET: {}", offset)?;
+                }
+                if let Some(order_by) = order_by {
+                    writeln!(f, "ORDER BY:")?;
+                    for (expr, dir) in order_by {
+                        writeln!(
+                            f,
+                            "  - {} {}",
+                            expr,
+                            if *dir == SortOrder::Asc {
+                                "ASC"
+                            } else {
+                                "DESC"
+                            }
+                        )?;
+                    }
+                }
+                Ok(())
+            }
             Self::Delete(delete_plan) => delete_plan.fmt(f),
             Self::Update(update_plan) => update_plan.fmt(f),
         }

core/translate/select.rs

@@ -14,7 +14,7 @@ use crate::vdbe::builder::{ProgramBuilderOpts, QueryMode};
 use crate::vdbe::insn::Insn;
 use crate::SymbolTable;
 use crate::{schema::Schema, vdbe::builder::ProgramBuilder, Result};
-use limbo_sqlite3_parser::ast::{self, SortOrder};
+use limbo_sqlite3_parser::ast::{self, CompoundSelect, SortOrder};
 use limbo_sqlite3_parser::ast::{ResultColumn, SelectInner};
 
 pub fn translate_select(
@@ -26,16 +26,34 @@ pub fn translate_select(
 ) -> Result<ProgramBuilder> {
     let mut select_plan = prepare_select_plan(schema, select, syms, None)?;
     optimize_plan(&mut select_plan, schema)?;
-    let Plan::Select(ref select) = select_plan else {
-        panic!("select_plan is not a SelectPlan");
+    let opts = match &select_plan {
+        Plan::Select(select) => ProgramBuilderOpts {
+            query_mode,
+            num_cursors: count_plan_required_cursors(select),
+            approx_num_insns: estimate_num_instructions(select),
+            approx_num_labels: estimate_num_labels(select),
+        },
+        Plan::CompoundSelect { first, rest, .. } => ProgramBuilderOpts {
+            query_mode,
+            num_cursors: count_plan_required_cursors(first)
+                + rest
+                    .iter()
+                    .map(|(plan, _)| count_plan_required_cursors(plan))
+                    .sum::<usize>(),
+            approx_num_insns: estimate_num_instructions(first)
+                + rest
+                    .iter()
+                    .map(|(plan, _)| estimate_num_instructions(plan))
+                    .sum::<usize>(),
+            approx_num_labels: estimate_num_labels(first)
+                + rest
+                    .iter()
+                    .map(|(plan, _)| estimate_num_labels(plan))
+                    .sum::<usize>(),
+        },
+        other => panic!("plan is not a SelectPlan: {:?}", other),
     };
 
-    let opts = ProgramBuilderOpts {
-        query_mode,
-        num_cursors: count_plan_required_cursors(select),
-        approx_num_insns: estimate_num_instructions(select),
-        approx_num_labels: estimate_num_labels(select),
-    };
     program.extend(&opts);
     emit_program(&mut program, select_plan, syms)?;
     Ok(program)
@@ -43,11 +61,91 @@ pub fn translate_select(
 
 pub fn prepare_select_plan<'a>(
     schema: &Schema,
-    select: ast::Select,
+    mut select: ast::Select,
     syms: &SymbolTable,
     outer_scope: Option<&'a Scope<'a>>,
 ) -> Result<Plan> {
-    match *select.body.select {
+    let compounds = select.body.compounds.take();
+    match compounds {
+        None => {
+            let limit = select.limit.take();
+            Ok(Plan::Select(prepare_one_select_plan(
+                schema,
+                *select.body.select,
+                limit.as_deref(),
+                select.order_by.take(),
+                select.with.take(),
+                syms,
+                outer_scope,
+            )?))
+        }
+        Some(compounds) => {
+            let mut first = prepare_one_select_plan(
+                schema,
+                *select.body.select,
+                None,
+                None,
+                None,
+                syms,
+                outer_scope,
+            )?;
+            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");
+                }
+                let plan =
+                    prepare_one_select_plan(schema, *select, None, None, None, syms, outer_scope)?;
+                rest.push((plan, operator));
+            }
+            // Ensure all subplans have same number of result columns
+            let first_num_result_columns = first.result_columns.len();
+            for (plan, operator) in rest.iter() {
+                if plan.result_columns.len() != first_num_result_columns {
+                    crate::bail_parse_error!("SELECTs to the left and right of {} do not have the same number of result columns", operator);
+                }
+            }
+            let (limit, offset) = select.limit.map_or(Ok((None, None)), |l| parse_limit(&l))?;
+
+            first.limit = limit.clone();
+            for (plan, _) in rest.iter_mut() {
+                plan.limit = limit.clone();
+            }
+
+            // FIXME: handle OFFSET for compound selects
+            if offset.is_some() {
+                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");
+            }
+            // FIXME: handle WITH for compound selects
+            if select.with.is_some() {
+                crate::bail_parse_error!("WITH is not supported for compound SELECTs yet");
+            }
+            Ok(Plan::CompoundSelect {
+                first,
+                rest,
+                limit,
+                offset,
+                order_by: None,
+            })
+        }
+    }
+}
+
+fn prepare_one_select_plan<'a>(
+    schema: &Schema,
+    select: ast::OneSelect,
+    limit: Option<&ast::Limit>,
+    order_by: Option<Vec<ast::SortedColumn>>,
+    with: Option<ast::With>,
+    syms: &SymbolTable,
+    outer_scope: Option<&'a Scope<'a>>,
+) -> Result<SelectPlan> {
+    match select {
         ast::OneSelect::Select(select_inner) => {
             let SelectInner {
                 mut columns,
@@ -64,8 +162,6 @@ pub fn prepare_select_plan<'a>(
 
             let mut where_predicates = vec![];
 
-            let with = select.with;
-
             // Parse the FROM clause into a vec of TableReferences. Fold all the join conditions expressions into the WHERE clause.
             let table_references =
                 parse_from(schema, from, syms, with, &mut where_predicates, outer_scope)?;
@@ -375,7 +471,7 @@ pub fn prepare_select_plan<'a>(
             plan.aggregates = aggregate_expressions;
 
             // Parse the ORDER BY clause
-            if let Some(order_by) = select.order_by {
+            if let Some(order_by) = order_by {
                 let mut key = Vec::new();
 
                 for mut o in order_by {
@@ -397,11 +493,10 @@ pub fn prepare_select_plan<'a>(
             }
 
             // Parse the LIMIT/OFFSET clause
-            (plan.limit, plan.offset) =
-                select.limit.map_or(Ok((None, None)), |l| parse_limit(&l))?;
+            (plan.limit, plan.offset) = limit.map_or(Ok((None, None)), |l| parse_limit(l))?;
 
             // Return the unoptimized query plan
-            Ok(Plan::Select(plan))
+            Ok(plan)
         }
         ast::OneSelect::Values(values) => {
             let len = values[0].len();
@@ -430,7 +525,7 @@ pub fn prepare_select_plan<'a>(
                 values,
             };
 
-            Ok(Plan::Select(plan))
+            Ok(plan)
         }
     }
 }

testing/select.test

@@ -240,3 +240,48 @@ do_execsql_test select-invalid-numeric-text {
 do_execsql_test select-invalid-numeric-text {
   select -'E';
 } {0}
+
+do_execsql_test_on_specific_db {:memory:} select-union-all-1 {
+  CREATE TABLE t1(x INTEGER);
+  CREATE TABLE t2(x INTEGER); 
+  CREATE TABLE t3(x INTEGER);
+  
+  INSERT INTO t1 VALUES(1),(2),(3);
+  INSERT INTO t2 VALUES(4),(5),(6);
+  INSERT INTO t3 VALUES(7),(8),(9);
+  
+  SELECT x FROM t1
+  UNION ALL 
+  SELECT x FROM t2
+  UNION ALL
+  SELECT x FROM t3;
+} {1
+2
+3
+4
+5
+6
+7
+8
+9}
+
+do_execsql_test_on_specific_db {:memory:} select-union-all-with-filters {
+  CREATE TABLE t4(x INTEGER);
+  CREATE TABLE t5(x INTEGER);
+  CREATE TABLE t6(x INTEGER);
+  
+  INSERT INTO t4 VALUES(1),(2),(3),(4);
+  INSERT INTO t5 VALUES(5),(6),(7),(8);
+  INSERT INTO t6 VALUES(9),(10),(11),(12);
+
+  SELECT x FROM t4 WHERE x > 2
+  UNION ALL
+  SELECT x FROM t5 WHERE x < 7 
+  UNION ALL
+  SELECT x FROM t6 WHERE x = 10;
+} {3
+4
+5
+6
+10}
+