rupg_analyze/
setop.rs

1//! `transformSetOperationStmt` of `analyze.c`: `UNION`, `INTERSECT` and `EXCEPT`.
2//!
3//! Each `SELECT` of the tree is a subquery of the query, in the order of the tree from the left. The columns of the query are the columns of the leftmost `SELECT`, with the types that the operations give. `ORDER BY` can use only the names and the numbers of the columns.
4
5use rupg_common::{Error, Result, SqlState};
6use rupg_sql::nodes::{LockClauseStrength, SelectStmt, SetOperation};
7use rupg_types::oid;
8
9use crate::Analyzer;
10use crate::agg::Kind;
11use crate::coerce::{AtOpt, Context};
12use crate::expr::{Expr, ExprKind, Var};
13use crate::from::{Column, Relation};
14use crate::select::{Query, Target};
15use crate::sort::{SortGroup, set_group};
16use crate::typename::place;
17
18/// The kinds of set operation.
19#[derive(Clone, Copy, Debug, PartialEq, Eq)]
20pub enum SetKind {
21    Union,
22    Intersect,
23    Except,
24}
25
26/// A part of the tree of a set operation.
27#[derive(Clone, Debug, PartialEq)]
28pub enum SetTree {
29    /// A `SELECT`, as the index of its subquery in [`Query::relations`].
30    Leaf(usize),
31    /// An operation on two parts.
32    Op(Box<SetOp>),
33}
34
35/// An operation of the tree, as `SetOperationStmt`.
36#[derive(Clone, Debug, PartialEq)]
37pub struct SetOp {
38    pub kind: SetKind,
39    /// `ALL`: the operation keeps the rows that are equal.
40    pub all: bool,
41    pub left: SetTree,
42    pub right: SetTree,
43    /// The type and the typmod of each column of the result.
44    pub columns: Vec<(u32, i32)>,
45    /// The operators that find the equal rows, one for each column, as `groupClauses`. `UNION ALL` has none.
46    pub groups: Vec<SortGroup>,
47    /// For the `UNION` of a recursive query of `WITH`, as `RecursiveUnion`, the number of the work table that the reference of the recursive term reads. The right side runs again with the rows of the last step in the work table until it gives no new rows.
48    pub work: Option<usize>,
49    /// The casts of the columns of the operation to the types of the operation above it, which the planner fills for `EXPLAIN`: a `Result` above the plan of the operation computes them, as `recurse_set_operations` adds it. Each cast reads the column of the leftmost `SELECT`. Empty when the types are the same.
50    pub project: Vec<Expr>,
51}
52
53impl SetTree {
54    /// The index of the leftmost `SELECT`.
55    pub fn leftmost(&self) -> usize {
56        match self {
57            SetTree::Leaf(index) => *index,
58            SetTree::Op(op) => op.left.leftmost(),
59        }
60    }
61
62    /// Calls `f` for the index of each `SELECT` of the part.
63    pub fn each_leaf(&self, f: &mut impl FnMut(usize)) {
64        match self {
65            SetTree::Leaf(index) => f(*index),
66            SetTree::Op(op) => {
67                op.left.each_leaf(f);
68                op.right.each_leaf(f);
69            }
70        }
71    }
72}
73
74/// `LCS_asString`.
75pub fn strength_name(strength: LockClauseStrength) -> &'static str {
76    match strength {
77        LockClauseStrength::LCS_FORKEYSHARE => "FOR KEY SHARE",
78        LockClauseStrength::LCS_FORSHARE => "FOR SHARE",
79        LockClauseStrength::LCS_FORNOKEYUPDATE => "FOR NO KEY UPDATE",
80        _ => "FOR UPDATE",
81    }
82}
83
84/// The error of `FOR UPDATE` and the other locking clauses in a set operation.
85fn locking_error(s: &SelectStmt) -> Result<()> {
86    match s.lockingClause.first() {
87        Some(Some(rupg_sql::nodes::Node::LockingClause(lock))) => Err(Error::new(
88            SqlState::FEATURE_NOT_SUPPORTED,
89            format!("{} is not allowed with UNION/INTERSECT/EXCEPT", strength_name(lock.strength)),
90        )),
91        _ => Ok(()),
92    }
93}
94
95/// `select_common_type` gives the first expression with the type that it selects, and the error points at it.
96fn best(exprs: [&Expr; 2], ty: u32) -> Option<usize> {
97    let found = exprs.iter().find(|e| e.ty == ty);
98    found.unwrap_or(&exprs[0]).place()
99}
100
101impl Analyzer<'_> {
102    /// `transformSetOperationStmt`: the query tree of a set operation.
103    pub(crate) fn set_operation(&mut self, s: &SelectStmt, parent: Option<usize>) -> Result<Query> {
104        let mut leftmost = s;
105        while leftmost.op != SetOperation::SETOP_NONE
106            && let Some(left) = leftmost.larg.as_deref()
107        {
108            leftmost = left;
109        }
110        if leftmost.intoClause.is_some() {
111            return Err(Error::new(
112                SqlState::FEATURE_NOT_SUPPORTED,
113                "SELECT INTO is not supported yet",
114            ));
115        }
116        let (tree, _) = self.set_tree(s, true, std::ptr::from_ref(leftmost), parent)?;
117        let SetTree::Op(op) = tree else {
118            return Err(Error::internal("a set operation with no operation"));
119        };
120        let left = op.left.leftmost();
121        let Some(query) = self.scope.relations[left].subquery.as_deref() else {
122            return Err(Error::internal("a set operation with no leftmost SELECT"));
123        };
124        let mut targets = Vec::with_capacity(op.columns.len());
125        for (i, (target, &(ty, typmod))) in
126            query.targets.iter().filter(|t| !t.junk).zip(&op.columns).enumerate()
127        {
128            let attnum = i16::try_from(i + 1).unwrap_or(i16::MAX);
129            let var = Var { relation: left, attnum, levels_up: 0 };
130            let expr = Expr { kind: ExprKind::Var(var), ty, typmod, location: target.expr.place() };
131            targets.push(Target { name: target.name.clone(), expr, origin: None, junk: false });
132        }
133        // The names of the columns are visible to `ORDER BY`, which cannot add a junk column.
134        let columns = targets.iter().map(|t| (t.name.clone(), t.expr.clone())).collect();
135        let saved = self.add_columns(columns);
136        let width = targets.len();
137        let sort = self.sort_clause(&s.sortClause, &mut targets);
138        self.remove_columns(saved);
139        let sort = sort?;
140        if let Some(junk) = targets.get(width) {
141            return Err(Error::new(
142                SqlState::FEATURE_NOT_SUPPORTED,
143                "invalid UNION/INTERSECT/EXCEPT ORDER BY clause",
144            )
145            .with_detail("Only result column names can be used, not expressions or functions.")
146            .with_hint(
147                "Add the expression/function to every SELECT, or move the UNION into a FROM clause.",
148            )
149            .at_opt(junk.expr.place()));
150        }
151        let with_ties = s.limitOption == rupg_sql::nodes::LimitOption::LIMIT_OPTION_WITH_TIES;
152        let offset = self.with_kind(Kind::Offset, |a| {
153            a.limit_clause(s.limitOffset.as_ref(), "OFFSET", with_ties)
154        })?;
155        let limit = self.with_kind(Kind::Limit, |a| {
156            a.limit_clause(s.limitCount.as_ref(), "LIMIT", with_ties)
157        })?;
158        locking_error(s)?;
159        let relations = std::mem::take(&mut self.scope.relations);
160        Ok(Query {
161            relations,
162            from: Vec::new(),
163            targets,
164            filter: None,
165            grouped: false,
166            group: Vec::new(),
167            grouping_sets: Vec::new(),
168            group_distinct: false,
169            having: None,
170            sort,
171            distinct: Vec::new(),
172            distinct_on: false,
173            offset,
174            limit,
175            with_ties,
176            params: Vec::new(),
177            notices: Vec::new(),
178            set_op: Some(*op),
179            windows: Vec::new(),
180            ctes: Vec::new(),
181            recursive: false,
182            row_marks: Vec::new(),
183            write: None,
184            out_args: Vec::new(),
185            pulled: 0,
186            pushed: false,
187            barrier: Vec::new(),
188            minmax: None,
189            constraint_deps: Vec::new(),
190            shown: None,
191            perms: Vec::new(),
192            funcs: Vec::new(),
193        })
194    }
195
196    /// `transformSetOperationTree`: a part of the tree and the expression of each of its columns. A part with `ORDER BY`, `LIMIT`, `OFFSET` or `WITH` under the top is a subquery of its own. `leftmost` is the leftmost `SELECT`, which can have `INTO`.
197    fn set_tree(
198        &mut self,
199        s: &SelectStmt,
200        top: bool,
201        leftmost: *const SelectStmt,
202        parent: Option<usize>,
203    ) -> Result<(SetTree, Vec<Expr>)> {
204        if let Some(into) = &s.intoClause
205            && !std::ptr::eq(s, leftmost)
206        {
207            let at = into.rel.as_ref().and_then(|r| place(r.location));
208            return Err(Error::new(
209                SqlState::SYNTAX_ERROR,
210                "INTO is only allowed on first SELECT of UNION/INTERSECT/EXCEPT",
211            )
212            .at_opt(at));
213        }
214        if !top {
215            locking_error(s)?;
216        }
217        let leaf = s.op == SetOperation::SETOP_NONE
218            || (!top
219                && (!s.sortClause.is_empty()
220                    || s.limitOffset.is_some()
221                    || s.limitCount.is_some()
222                    || s.withClause.is_some()));
223        if leaf {
224            self.keep_unknowns = true;
225            let query = self.sub_select(s, false)?;
226            let shown = query.targets.iter().filter(|t| !t.junk);
227            let exprs: Vec<Expr> = shown.clone().map(|t| t.expr.clone()).collect();
228            let columns = shown
229                .map(|t| Column {
230                    name: t.name.clone(),
231                    ty: t.expr.ty,
232                    typmod: t.expr.typmod,
233                    not_null: false,
234                })
235                .collect();
236            let index = self.scope.relations.len();
237            let subquery = Some(Box::new(query));
238            let mut relation = Relation::new(0, columns, subquery, None, None);
239            relation.name = format!("*SELECT* {}", index + 1);
240            self.scope.relations.push(relation);
241            return Ok((SetTree::Leaf(index), exprs));
242        }
243        let (kind, context) = match s.op {
244            SetOperation::SETOP_UNION => (SetKind::Union, "UNION"),
245            SetOperation::SETOP_INTERSECT => (SetKind::Intersect, "INTERSECT"),
246            _ => (SetKind::Except, "EXCEPT"),
247        };
248        let (Some(larg), Some(rarg)) = (s.larg.as_deref(), s.rarg.as_deref()) else {
249            return Err(Error::internal("a set operation with no arguments"));
250        };
251        let (left, lexprs) = self.set_tree(larg, false, leftmost, parent)?;
252        // `determineRecursiveColTypes`: the recursive term reads the columns that the non-recursive term gives.
253        let recursive = self.recursive_parent(parent).is_some();
254        if top
255            && recursive
256            && let Some(slot) = parent
257        {
258            self.recursive_columns(slot, left.leftmost(), &lexprs)?;
259        }
260        let (right, rexprs) = self.set_tree(rarg, false, leftmost, parent)?;
261        if lexprs.len() != rexprs.len() {
262            return Err(Error::new(
263                SqlState::SYNTAX_ERROR,
264                format!("each {context} query must have the same number of columns"),
265            )
266            .at_opt(rexprs.iter().find_map(Expr::place)));
267        }
268        let mut columns = Vec::with_capacity(lexprs.len());
269        let mut groups = Vec::new();
270        let mut exprs = Vec::with_capacity(lexprs.len());
271        for (i, (l, r)) in lexprs.into_iter().zip(rexprs).enumerate() {
272            let ty = self.common_type(&[l.clone(), r.clone()], context)?;
273            let at = best([&l, &r], ty);
274            let l = self.set_column(&left, i, l, ty, context)?;
275            let r = self.set_column(&right, i, r, ty, context)?;
276            // select_common_typmod: the typmod of the columns when they have the type and the same typmod.
277            let typmod =
278                if l.ty == ty && r.ty == ty && l.typmod == r.typmod { l.typmod } else { -1 };
279            if kind != SetKind::Union || !s.all {
280                let mut group = set_group(i, ty, at)?;
281                // A recursive `UNION` must hash the rows, and a record can hash when each of its fields can, which the executor finds out.
282                if recursive && matches!(ty, oid::RECORD | oid::RECORD_ARRAY) {
283                    group.hashable = true;
284                }
285                groups.push(group);
286            }
287            columns.push((ty, typmod));
288            exprs.push(Expr { kind: ExprKind::CaseTest, ty, typmod, location: at });
289        }
290        let work = if top { self.work_table(parent) } else { None };
291        let op =
292            SetOp { kind, all: s.all, left, right, columns, groups, work, project: Vec::new() };
293        Ok((SetTree::Op(Box::new(op)), exprs))
294    }
295
296    /// `coerce_to_common_type` for a column of a part of a set operation. A constant or a parameter of type `unknown` in a `SELECT` gets the type, so that an error of its input points at it. The other columns of each `SELECT` of the part get a cast to the type, as `generate_setop_tlist` adds it.
297    fn set_column(
298        &mut self,
299        part: &SetTree,
300        column: usize,
301        expr: Expr,
302        ty: u32,
303        context: &str,
304    ) -> Result<Expr> {
305        let unknown_leaf = expr.ty == oid::UNKNOWN
306            && matches!(part, SetTree::Leaf(_))
307            && matches!(expr.kind, ExprKind::Const(_) | ExprKind::Param(_));
308        let coerced = if expr.ty != oid::UNKNOWN || unknown_leaf {
309            self.coerce_to_common(expr.clone(), ty, context)?
310        } else {
311            expr.clone()
312        };
313        let mut leaves = Vec::new();
314        part.each_leaf(&mut |index| leaves.push(index));
315        for index in leaves {
316            let relation = &mut self.scope.relations[index];
317            let Some(query) = relation.subquery.as_deref_mut() else { continue };
318            let Some(target) = query.targets.iter_mut().filter(|t| !t.junk).nth(column) else {
319                continue;
320            };
321            if target.expr.ty == ty {
322                continue;
323            }
324            let old = std::mem::replace(&mut target.expr, Expr::new(ExprKind::CaseTest, 0));
325            let new = if unknown_leaf { coerced.clone() } else { self.coerce_leaf(old, ty)? };
326            let Some(query) = self.scope.relations[index].subquery.as_deref_mut() else { continue };
327            if let Some(target) = query.targets.iter_mut().filter(|t| !t.junk).nth(column) {
328                target.expr = new;
329            }
330            if let Some(slot) = self.scope.relations[index].columns.get_mut(column) {
331                slot.ty = ty;
332                slot.typmod = -1;
333            }
334        }
335        Ok(if unknown_leaf { coerced } else { Expr { ty, typmod: coerced.typmod, ..expr } })
336    }
337
338    /// The cast of a column of a `SELECT` to the type of the set operation.
339    fn coerce_leaf(&mut self, expr: Expr, ty: u32) -> Result<Expr> {
340        self.coerce(expr, ty, -1, Context::Implicit, None)
341    }
342}
rupg_analyze/setop.rs

rupg_analyze/
setop.rs
