Skip to main content

nodedb_sql/planner/
union.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! UNION / UNION ALL planning.
4
5use sqlparser::ast::{self, SetExpr, SetOperator, SetQuantifier};
6
7use crate::error::{Result, SqlError};
8use crate::functions::registry::FunctionRegistry;
9use crate::types::*;
10
11/// Plan a UNION / UNION ALL / INTERSECT / EXCEPT operation.
12pub fn plan_set_operation(
13    op: &SetOperator,
14    left: &SetExpr,
15    right: &SetExpr,
16    quantifier: &SetQuantifier,
17    catalog: &dyn SqlCatalog,
18    functions: &FunctionRegistry,
19    temporal: crate::TemporalScope,
20) -> Result<SqlPlan> {
21    let left_plan = plan_set_expr(left, catalog, functions, temporal)?;
22    let right_plan = plan_set_expr(right, catalog, functions, temporal)?;
23
24    match op {
25        SetOperator::Union => {
26            let distinct = matches!(quantifier, SetQuantifier::Distinct | SetQuantifier::None);
27            Ok(SqlPlan::Union {
28                inputs: vec![left_plan, right_plan],
29                distinct,
30            })
31        }
32        SetOperator::Intersect => {
33            let all = matches!(quantifier, SetQuantifier::All);
34            Ok(SqlPlan::Intersect {
35                left: Box::new(left_plan),
36                right: Box::new(right_plan),
37                all,
38            })
39        }
40        SetOperator::Except => {
41            let all = matches!(quantifier, SetQuantifier::All);
42            Ok(SqlPlan::Except {
43                left: Box::new(left_plan),
44                right: Box::new(right_plan),
45                all,
46            })
47        }
48        _ => Err(SqlError::Unsupported {
49            detail: format!("set operation: {op}"),
50        }),
51    }
52}
53
54fn plan_set_expr(
55    expr: &SetExpr,
56    catalog: &dyn SqlCatalog,
57    functions: &FunctionRegistry,
58    temporal: crate::TemporalScope,
59) -> Result<SqlPlan> {
60    match expr {
61        SetExpr::Select(select) => {
62            // Wrap in a dummy Query to reuse plan_query.
63            let query = ast::Query {
64                with: None,
65                body: Box::new(SetExpr::Select(select.clone())),
66                order_by: None,
67                limit_clause: None,
68                fetch: None,
69                locks: Vec::new(),
70                for_clause: None,
71                settings: None,
72                format_clause: None,
73                pipe_operators: Vec::new(),
74            };
75            super::select::plan_query(&query, catalog, functions, temporal)
76        }
77        SetExpr::SetOperation {
78            op,
79            left,
80            right,
81            set_quantifier,
82        } => plan_set_operation(
83            op,
84            left,
85            right,
86            set_quantifier,
87            catalog,
88            functions,
89            temporal,
90        ),
91        _ => Err(SqlError::Unsupported {
92            detail: format!("set expression: {expr}"),
93        }),
94    }
95}