Skip to main content

uqa_planner/
join_graph.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Join graph: nodes are individual relations (a single table or a
8//! materialised CTE / subquery output), edges are equijoin predicates.
9//! The DPccp enumerator walks this graph to pick a join order.
10//!
11//! The graph is dense by design (`u64` bitmask side sets) so DP cache
12//! lookups can use bitmasks directly. With 64 relations max this
13//! suffices for every shape the UQA SQL compiler can build.
14
15use thiserror::Error;
16
17#[derive(Debug, Error, Clone, PartialEq)]
18pub enum JoinGraphError {
19    #[error("join-order optimization requires at least one relation")]
20    EmptyGraph,
21    #[error("JoinGraph supports at most 64 relations; cannot add `{name}` at index {index}")]
22    TooManyRelations { name: String, index: usize },
23    #[error(
24        "join edge references relation index {index}, but the graph contains {relation_count} relations"
25    )]
26    UnknownRelation { index: usize, relation_count: usize },
27    #[error("relation `{name}` has invalid cardinality estimate {rows}; expected a finite non-negative value")]
28    InvalidCardinality { name: String, rows: f64 },
29    #[error(
30        "relation `{name}` has invalid access cost {cost}; expected a finite non-negative value"
31    )]
32    InvalidAccessCost { name: String, cost: f64 },
33    #[error("join selectivity must be finite and between 0 and 1, got {selectivity}")]
34    InvalidSelectivity { selectivity: f64 },
35    #[error("duplicate join relation alias `{alias}`")]
36    DuplicateAlias { alias: String },
37    #[error("join predicate references unknown relation alias `{alias}`")]
38    UnknownAlias { alias: String },
39    #[error("invalid join plan: {detail}")]
40    InvalidPlan { detail: String },
41}
42
43pub type JoinGraphResult<T> = Result<T, JoinGraphError>;
44
45#[derive(Debug, Clone)]
46pub struct JoinEdge {
47    /// Bitmask of relations on the left side of the equijoin.
48    pub left: u64,
49    /// Bitmask of relations on the right side of the equijoin.
50    pub right: u64,
51    /// Selectivity of the join predicate (1 / max(distinct_left,
52    /// distinct_right) for an equijoin).
53    pub selectivity: f64,
54}
55
56#[derive(Debug, Clone, Default)]
57pub struct JoinGraph {
58    /// Relation labels in declaration order. The bit index of a
59    /// relation is its position in this vector.
60    pub(crate) relations: Vec<String>,
61    /// Per-relation row count estimate.
62    pub(crate) cardinalities: Vec<f64>,
63    /// Cost of producing each base relation after its local access predicates.
64    pub(crate) access_costs: Vec<f64>,
65    /// Connecting edges. Order is irrelevant; the enumerator probes
66    /// them by bitmask membership.
67    pub(crate) edges: Vec<JoinEdge>,
68}
69
70impl JoinGraph {
71    pub fn new() -> Self {
72        Self::default()
73    }
74
75    pub fn add_relation(&mut self, name: impl Into<String>, rows: f64) -> JoinGraphResult<usize> {
76        self.add_relation_with_cost(name, rows, rows)
77    }
78
79    pub fn add_relation_with_cost(
80        &mut self,
81        name: impl Into<String>,
82        rows: f64,
83        access_cost: f64,
84    ) -> JoinGraphResult<usize> {
85        let idx = self.relations.len();
86        let name = name.into();
87        if idx >= 64 {
88            return Err(JoinGraphError::TooManyRelations { name, index: idx });
89        }
90        if !rows.is_finite() || rows < 0.0 {
91            return Err(JoinGraphError::InvalidCardinality { name, rows });
92        }
93        if !access_cost.is_finite() || access_cost < 0.0 {
94            return Err(JoinGraphError::InvalidAccessCost {
95                name,
96                cost: access_cost,
97            });
98        }
99        self.relations.push(name);
100        self.cardinalities.push(rows);
101        self.access_costs.push(access_cost);
102        Ok(idx)
103    }
104
105    pub fn add_edge(
106        &mut self,
107        left_idx: usize,
108        right_idx: usize,
109        selectivity: f64,
110    ) -> JoinGraphResult<()> {
111        for index in [left_idx, right_idx] {
112            if index >= self.relations.len() {
113                return Err(JoinGraphError::UnknownRelation {
114                    index,
115                    relation_count: self.relations.len(),
116                });
117            }
118        }
119        if !selectivity.is_finite() || !(0.0..=1.0).contains(&selectivity) {
120            return Err(JoinGraphError::InvalidSelectivity { selectivity });
121        }
122        self.edges.push(JoinEdge {
123            left: 1u64 << left_idx,
124            right: 1u64 << right_idx,
125            selectivity,
126        });
127        Ok(())
128    }
129
130    pub fn relation_count(&self) -> usize {
131        self.relations.len()
132    }
133
134    /// All edges connecting a left subset to a right subset of the
135    /// graph. An edge is "between" `(s1, s2)` when one endpoint is in
136    /// `s1` and the other is in `s2`.
137    pub fn edges_between(&self, s1: u64, s2: u64) -> Vec<&JoinEdge> {
138        self.edges
139            .iter()
140            .filter(|e| {
141                let l_in_1 = e.left & s1 != 0;
142                let r_in_2 = e.right & s2 != 0;
143                let l_in_2 = e.left & s2 != 0;
144                let r_in_1 = e.right & s1 != 0;
145                (l_in_1 && r_in_2) || (l_in_2 && r_in_1)
146            })
147            .collect()
148    }
149
150    pub fn full_set(&self) -> u64 {
151        match self.relations.len() {
152            0 => 0,
153            64 => u64::MAX,
154            count => (1u64 << count) - 1,
155        }
156    }
157
158    /// Indices of every relation directly connected to `node` by any edge.
159    /// Each call walks the edge list once: `O(|edges|)`.
160    pub fn neighbors(&self, node: usize) -> Vec<usize> {
161        if node >= self.relations.len() {
162            return Vec::new();
163        }
164        let mark = 1u64 << node;
165        let mut out: Vec<usize> = Vec::new();
166        let mut seen: u64 = 0;
167        for edge in &self.edges {
168            let other = if edge.left == mark {
169                edge.right
170            } else if edge.right == mark {
171                edge.left
172            } else {
173                continue;
174            };
175            if other == 0 {
176                continue;
177            }
178            // `other` is a singleton bitmask of the other side.
179            let Ok(idx) = usize::try_from(other.trailing_zeros()) else {
180                continue;
181            };
182            if seen & (1u64 << idx) != 0 {
183                continue;
184            }
185            seen |= 1u64 << idx;
186            out.push(idx);
187        }
188        out.sort_unstable();
189        out
190    }
191}
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196
197    #[test]
198    fn full_set_covers_every_relation() {
199        let mut g = JoinGraph::new();
200        for i in 0..3 {
201            g.add_relation(format!("t{i}"), 100.0).unwrap();
202        }
203        assert_eq!(g.full_set(), 0b111);
204    }
205
206    #[test]
207    fn edges_between_finds_predicates_across_partition() {
208        let mut g = JoinGraph::new();
209        let a = g.add_relation("a", 100.0).unwrap();
210        let b = g.add_relation("b", 100.0).unwrap();
211        let c = g.add_relation("c", 100.0).unwrap();
212        g.add_edge(a, b, 0.01).unwrap();
213        g.add_edge(b, c, 0.01).unwrap();
214        let s1 = 1u64 << a;
215        let s2 = (1u64 << b) | (1u64 << c);
216        let between = g.edges_between(s1, s2);
217        assert_eq!(between.len(), 1);
218    }
219
220    #[test]
221    fn capacity_and_edge_errors_are_returned_without_panicking() {
222        let mut graph = JoinGraph::new();
223        for index in 0..64 {
224            graph.add_relation(format!("t{index}"), 1.0).unwrap();
225        }
226        assert_eq!(graph.full_set(), u64::MAX);
227        assert!(matches!(
228            graph.add_relation("overflow", 1.0),
229            Err(JoinGraphError::TooManyRelations { index: 64, .. })
230        ));
231        assert!(matches!(
232            graph.add_edge(0, 64, 1.0),
233            Err(JoinGraphError::UnknownRelation { index: 64, .. })
234        ));
235        assert!(graph.neighbors(64).is_empty());
236    }
237
238    #[test]
239    fn rejects_non_finite_cost_inputs_before_enumeration() {
240        let mut graph = JoinGraph::new();
241        assert!(matches!(
242            graph.add_relation("bad", f64::NAN),
243            Err(JoinGraphError::InvalidCardinality { .. })
244        ));
245        assert!(matches!(
246            graph.add_relation_with_cost("bad_cost", 1.0, f64::INFINITY),
247            Err(JoinGraphError::InvalidAccessCost { .. })
248        ));
249        let left = graph.add_relation("left", 1.0).unwrap();
250        let right = graph.add_relation("right", 1.0).unwrap();
251        for selectivity in [f64::NAN, -0.1, 1.1] {
252            assert!(matches!(
253                graph.add_edge(left, right, selectivity),
254                Err(JoinGraphError::InvalidSelectivity { .. })
255            ));
256        }
257    }
258}