Skip to main content

siderite_orm/queryset/
combine.rs

1//! Set operations: `union`, `union_all`, `intersection`, `difference`.
2
3use super::QuerySet;
4use crate::error::QueryError;
5use crate::expr::Expr;
6use crate::model::Model;
7use crate::plan::{Compound, QueryPlan, SetOp};
8
9impl<M: Model> QuerySet<M> {
10    /// Rows in either queryset, duplicates removed (`UNION`).
11    ///
12    /// Both querysets must project the same columns (same names, same
13    /// order). Order and limit the *combined* result by calling `order_by` /
14    /// `limit` on the returned queryset; ordering inside `other` is dropped,
15    /// and `order_by` terms must name output columns.
16    ///
17    /// # Errors
18    /// [`QueryError::InvalidPlan`] if the projections differ or the querysets
19    /// are bound to different databases (see
20    /// [`Db::same_database`](crate::Db::same_database)).
21    pub fn union(self, other: QuerySet<M>) -> Result<Self, QueryError> {
22        self.combine(SetOp::Union, other)
23    }
24
25    /// Rows in either queryset, keeping duplicates (`UNION ALL`).
26    ///
27    /// # Errors
28    /// As [`union`](Self::union).
29    pub fn union_all(self, other: QuerySet<M>) -> Result<Self, QueryError> {
30        self.combine(SetOp::UnionAll, other)
31    }
32
33    /// Rows in both querysets (`INTERSECT`).
34    ///
35    /// # Errors
36    /// As [`union`](Self::union).
37    pub fn intersection(self, other: QuerySet<M>) -> Result<Self, QueryError> {
38        self.combine(SetOp::Intersect, other)
39    }
40
41    /// Rows of this queryset that are not in `other` (`EXCEPT`).
42    ///
43    /// # Errors
44    /// As [`union`](Self::union).
45    pub fn difference(self, other: QuerySet<M>) -> Result<Self, QueryError> {
46        self.combine(SetOp::Except, other)
47    }
48
49    fn combine(self, op: SetOp, other: QuerySet<M>) -> Result<Self, QueryError> {
50        self.ready()?;
51        other.ready()?;
52        if !self.db.same_database(&other.db) {
53            return Err(QueryError::InvalidPlan(
54                "querysets bound to different databases cannot be combined".into(),
55            ));
56        }
57        let (mut base, mut member) = (self.materialize_empty(), other.materialize_empty());
58        if shape::<M>(&base.plan) != shape::<M>(&member.plan) {
59            return Err(QueryError::InvalidPlan(
60                "set operation members must project the same columns".into(),
61            ));
62        }
63        member.plan.ordering.clear();
64        base.plan.compound.push(Compound {
65            op,
66            plan: member.plan,
67        });
68        Ok(base)
69    }
70
71    /// Replace the `none()` flag by an always-false filter, so the queryset
72    /// can take part in a set operation.
73    fn materialize_empty(mut self) -> Self {
74        if self.empty {
75            self.empty = false;
76            self.plan.filter = Some(Expr::Or(Vec::new()));
77        }
78        self
79    }
80}
81
82/// Output column names of `plan`.
83fn shape<M: Model>(plan: &QueryPlan) -> Vec<String> {
84    if plan.projection.is_empty() {
85        return M::META.fields.iter().map(|f| f.column.to_owned()).collect();
86    }
87    plan.projection
88        .iter()
89        .map(|s| match (&s.alias, &s.expr) {
90            (Some(alias), _) => alias.to_string(),
91            (None, Expr::Column(column)) => column.name.to_string(),
92            (None, _) => String::new(),
93        })
94        .collect()
95}