egglog_core_relations/uf/
mod.rs

1//! A table implementation backed by a union-find.
2
3use std::{
4    any::Any,
5    mem::{self, ManuallyDrop},
6    sync::{Arc, Weak},
7};
8
9use crate::numeric_id::{DenseIdMap, NumericId};
10use crossbeam_queue::SegQueue;
11use smallvec::SmallVec;
12
13use crate::{
14    TableChange, TaggedRowBuffer,
15    action::ExecutionState,
16    common::{HashMap, Value},
17    offsets::{OffsetRange, RowId, Subset, SubsetRef},
18    pool::with_pool_set,
19    row_buffer::RowBuffer,
20    table_spec::{
21        ColumnId, Constraint, Generation, MutationBuffer, Offset, Rebuilder, Row, Table, TableSpec,
22        TableVersion, ValueRebuilder, WrappedTableRef,
23    },
24};
25
26#[cfg(test)]
27mod tests;
28
29type UnionFind = crate::union_find::UnionFind<Value>;
30
31/// A special table backed by a union-find used to efficiently implement
32/// egglog-style canonicaliztion.
33///
34/// To canonicalize columns, we need to efficiently discover values that have
35/// ceased to be canonical. To do that we keep a table of _displaced_ values:
36///
37/// This table has three columns:
38/// 1. (the only key): a value that is _no longer canonical_ in the equivalence relation.
39/// 2. The canonical value of the equivalence class.
40/// 3. The timestamp at which the key stopped being canonical.
41///
42/// We do not store the second value explicitly: instead, we compute it
43/// on-the-fly using a union-find data-structure.
44///
45/// This is related to the 'Leader' encoding in some versions of egglog:
46/// Displaced is a version of Leader that _only_ stores ids when they cease to
47/// be canonical. Rows are also "automatically updated" with the current leader,
48/// rather than requiring the DB to replay history or canonicalize redundant
49/// values in the table.
50///
51/// To union new ids `l`, and `r`, stage an update `Displaced(l, r, ts)` where
52/// `ts` is the current timestamp. Note that all tie-breaks and other encoding
53/// decisions are made internally, so there may not literally be a row added
54/// with this value.
55pub struct DisplacedTable {
56    uf: UnionFind,
57    displaced: Vec<(Value, Value)>,
58    changed: bool,
59    lookup_table: HashMap<Value, RowId>,
60    buffered_writes: Arc<SegQueue<RowBuffer>>,
61}
62
63struct Canonicalizer<'a> {
64    cols: Vec<ColumnId>,
65    table: &'a DisplacedTable,
66}
67
68impl ValueRebuilder for Canonicalizer<'_> {
69    fn rebuild_val(&self, val: Value) -> Value {
70        self.table.uf.find_naive(val)
71    }
72    // `rebuild_slice` uses the default (per-value `rebuild_val`).
73}
74
75impl Rebuilder for Canonicalizer<'_> {
76    fn hint_col(&self) -> Option<ColumnId> {
77        Some(ColumnId::new(0))
78    }
79    fn rebuild_buf(
80        &self,
81        buf: &RowBuffer,
82        start: RowId,
83        end: RowId,
84        out: &mut TaggedRowBuffer,
85        _exec_state: &mut ExecutionState,
86    ) {
87        if start >= end {
88            return;
89        }
90        assert!(end.index() <= buf.len());
91        let mut cur = start;
92        // SAFETY: `cur` is always in-bounds, guaranteed by the above assertion.
93        // Special-case small columns: this gives us a modest speedup on rebuilding-heavy
94        // workloads.
95        match self.cols.as_slice() {
96            [c] => {
97                while cur < end {
98                    let row = unsafe { buf.get_row_unchecked(cur) };
99                    let to_canon = row[c.index()];
100                    let canon = self.table.uf.find_naive(to_canon);
101                    if canon != to_canon {
102                        out.add_row_with(cur, row, |dst| dst[c.index()] = canon);
103                    }
104                    cur = cur.inc();
105                }
106            }
107            [c1, c2] => {
108                while cur < end {
109                    let row = unsafe { buf.get_row_unchecked(cur) };
110                    let v1 = row[c1.index()];
111                    let v2 = row[c2.index()];
112                    let ca1 = self.table.uf.find_naive(v1);
113                    let ca2 = self.table.uf.find_naive(v2);
114                    if ca1 != v1 || ca2 != v2 {
115                        out.add_row_with(cur, row, |dst| {
116                            dst[c1.index()] = ca1;
117                            dst[c2.index()] = ca2;
118                        });
119                    }
120                    cur = cur.inc();
121                }
122            }
123            [c1, c2, c3] => {
124                while cur < end {
125                    let row = unsafe { buf.get_row_unchecked(cur) };
126                    let v1 = row[c1.index()];
127                    let v2 = row[c2.index()];
128                    let v3 = row[c3.index()];
129                    let ca1 = self.table.uf.find_naive(v1);
130                    let ca2 = self.table.uf.find_naive(v2);
131                    let ca3 = self.table.uf.find_naive(v3);
132                    if ca1 != v1 || ca2 != v2 || ca3 != v3 {
133                        out.add_row_with(cur, row, |dst| {
134                            dst[c1.index()] = ca1;
135                            dst[c2.index()] = ca2;
136                            dst[c3.index()] = ca3;
137                        });
138                    }
139                    cur = cur.inc();
140                }
141            }
142            cs => {
143                let mut canons: SmallVec<[Value; 8]> = SmallVec::with_capacity(cs.len());
144                while cur < end {
145                    let row = unsafe { buf.get_row_unchecked(cur) };
146                    canons.clear();
147                    let mut changed = false;
148                    for c in cs {
149                        let to_canon = row[c.index()];
150                        let canon = self.table.uf.find_naive(to_canon);
151                        changed |= canon != to_canon;
152                        canons.push(canon);
153                    }
154                    if changed {
155                        out.add_row_with(cur, row, |dst| {
156                            for (c, canon) in cs.iter().zip(canons.iter()) {
157                                dst[c.index()] = *canon;
158                            }
159                        });
160                    }
161                    cur = cur.inc();
162                }
163            }
164        }
165    }
166    fn rebuild_subset(
167        &self,
168        other: WrappedTableRef,
169        subset: SubsetRef,
170        out: &mut TaggedRowBuffer,
171        _exec_state: &mut ExecutionState,
172    ) {
173        let old_len = u32::try_from(out.len()).expect("row buffer sizes should fit in a u32");
174        let _next = other.scan_bounded(subset, Offset::new(0), usize::MAX, out);
175        debug_assert!(_next.is_none());
176        for i in old_len..u32::try_from(out.len()).expect("row buffer sizes should fit in a u32") {
177            let i = RowId::new(i);
178            let (_id, row) = out.get_row_mut(i);
179            let mut changed = false;
180            for col in &self.cols {
181                let to_canon = row[col.index()];
182                let canon = self.table.uf.find_naive(to_canon);
183                changed |= canon != to_canon;
184                row[col.index()] = canon;
185            }
186            if !changed {
187                out.set_stale(i);
188            }
189        }
190    }
191}
192
193impl Default for DisplacedTable {
194    fn default() -> Self {
195        Self {
196            uf: UnionFind::default(),
197            displaced: Vec::new(),
198            changed: false,
199            lookup_table: HashMap::default(),
200            buffered_writes: Arc::new(SegQueue::new()),
201        }
202    }
203}
204
205impl Clone for DisplacedTable {
206    fn clone(&self) -> Self {
207        DisplacedTable {
208            uf: self.uf.clone(),
209            displaced: self.displaced.clone(),
210            changed: self.changed,
211            lookup_table: self.lookup_table.clone(),
212            buffered_writes: Default::default(),
213        }
214    }
215}
216
217struct UfBuffer {
218    to_insert: ManuallyDrop<RowBuffer>,
219    buffered_writes: Weak<SegQueue<RowBuffer>>,
220}
221
222impl Drop for UfBuffer {
223    fn drop(&mut self) {
224        let Some(buffered_writes) = self.buffered_writes.upgrade() else {
225            // SAFETY: If we can't write updates, manually drop to_insert
226            unsafe {
227                ManuallyDrop::drop(&mut self.to_insert);
228            }
229            return;
230        };
231        // SAFETY: self.to_insert will not be used again after this point.
232        //
233        // This avoids creating a fresh row buffer via `mem::take` or `mem::swap` and
234        // dropping it immediately.
235        let to_insert = unsafe { ManuallyDrop::take(&mut self.to_insert) };
236        buffered_writes.push(to_insert);
237    }
238}
239
240impl MutationBuffer for UfBuffer {
241    fn stage_insert(&mut self, row: &[Value]) {
242        self.to_insert.add_row(row);
243    }
244    fn stage_remove(&mut self, _: &[Value]) {
245        panic!("attempting to remove data from a DisplacedTable")
246    }
247    fn fresh_handle(&self) -> Box<dyn MutationBuffer> {
248        Box::new(UfBuffer {
249            to_insert: ManuallyDrop::new(RowBuffer::new(self.to_insert.arity())),
250            buffered_writes: self.buffered_writes.clone(),
251        })
252    }
253}
254
255impl Table for DisplacedTable {
256    fn dyn_clone(&self) -> Box<dyn Table> {
257        Box::new(self.clone())
258    }
259    fn as_any(&self) -> &dyn Any {
260        self
261    }
262    fn spec(&self) -> TableSpec {
263        let mut uncacheable_columns = DenseIdMap::default();
264        // The second column of this table is determined dynamically by the union-find.
265        uncacheable_columns.insert(ColumnId::new(1), true);
266        TableSpec {
267            n_keys: 1,
268            n_vals: 2,
269            uncacheable_columns,
270            allows_delete: false,
271        }
272    }
273
274    fn rebuilder<'a>(&'a self, cols: &[ColumnId]) -> Option<Box<dyn Rebuilder + 'a>> {
275        Some(Box::new(Canonicalizer {
276            cols: cols.to_vec(),
277            table: self,
278        }))
279    }
280
281    fn clear(&mut self) {
282        self.uf.reset();
283        self.displaced.clear();
284    }
285
286    fn all(&self) -> Subset {
287        Subset::Dense(OffsetRange::new(
288            RowId::new(0),
289            RowId::from_usize(self.displaced.len()),
290        ))
291    }
292
293    fn len(&self) -> usize {
294        self.displaced.len()
295    }
296
297    fn version(&self) -> TableVersion {
298        TableVersion {
299            major: Generation::new(0),
300            minor: Offset::from_usize(self.displaced.len()),
301        }
302    }
303
304    fn updates_since(&self, offset: Offset) -> Subset {
305        Subset::Dense(OffsetRange::new(
306            RowId::from_usize(offset.index()),
307            RowId::from_usize(self.displaced.len()),
308        ))
309    }
310
311    fn scan_generic_bounded(
312        &self,
313        subset: SubsetRef,
314        start: Offset,
315        n: usize,
316        cs: &[Constraint],
317        mut f: impl FnMut(RowId, &[Value]),
318    ) -> Option<Offset>
319    where
320        Self: Sized,
321    {
322        if cs.is_empty() {
323            let start = start.index();
324            subset
325                .iter_bounded(start, start + n, |row| {
326                    f(row, self.expand(row).as_slice());
327                })
328                .map(Offset::from_usize)
329        } else {
330            let start = start.index();
331            subset
332                .iter_bounded(start, start + n, |row| {
333                    if cs.iter().all(|c| self.eval(c, row)) {
334                        f(row, self.expand(row).as_slice());
335                    }
336                })
337                .map(Offset::from_usize)
338        }
339    }
340
341    fn refine_one(&self, mut subset: Subset, c: &Constraint) -> Subset {
342        subset.retain(|row| self.eval(c, row));
343        subset
344    }
345
346    fn fast_subset(&self, constraint: &Constraint) -> Option<Subset> {
347        let ts = ColumnId::new(2);
348        match constraint {
349            Constraint::Eq { .. } => None,
350            Constraint::EqConst { col, val } => {
351                if *col == ColumnId::new(1) {
352                    return None;
353                }
354                if *col == ColumnId::new(0) {
355                    return Some(match self.lookup_table.get(val) {
356                        Some(row) => Subset::Dense(OffsetRange::new(
357                            *row,
358                            RowId::from_usize(row.index() + 1),
359                        )),
360                        None => Subset::empty(),
361                    });
362                }
363                match self.timestamp_bounds(*val) {
364                    Ok((start, end)) => Some(Subset::Dense(OffsetRange::new(start, end))),
365                    Err(_) => None,
366                }
367            }
368            Constraint::LtConst { col, val } => {
369                if *col != ts {
370                    return None;
371                }
372                match self.timestamp_bounds(*val) {
373                    Err(bound) | Ok((bound, _)) => {
374                        Some(Subset::Dense(OffsetRange::new(RowId::new(0), bound)))
375                    }
376                }
377            }
378            Constraint::GtConst { col, val } => {
379                if *col != ts {
380                    return None;
381                }
382
383                match self.timestamp_bounds(*val) {
384                    Err(bound) | Ok((_, bound)) => Some(Subset::Dense(OffsetRange::new(
385                        bound,
386                        RowId::from_usize(self.displaced.len()),
387                    ))),
388                }
389            }
390            Constraint::LeConst { col, val } => {
391                if *col != ts {
392                    return None;
393                }
394
395                match self.timestamp_bounds(*val) {
396                    Err(bound) | Ok((_, bound)) => {
397                        Some(Subset::Dense(OffsetRange::new(RowId::new(0), bound)))
398                    }
399                }
400            }
401            Constraint::GeConst { col, val } => {
402                if *col != ts {
403                    return None;
404                }
405
406                match self.timestamp_bounds(*val) {
407                    Err(bound) | Ok((bound, _)) => Some(Subset::Dense(OffsetRange::new(
408                        bound,
409                        RowId::from_usize(self.displaced.len()),
410                    ))),
411                }
412            }
413        }
414    }
415
416    fn get_row(&self, key: &[Value]) -> Option<Row> {
417        assert_eq!(key.len(), 1, "attempt to lookup a row with the wrong key");
418        let row_id = *self.lookup_table.get(&key[0])?;
419        let mut vals = with_pool_set(|ps| ps.get::<Vec<Value>>());
420        vals.extend_from_slice(self.expand(row_id).as_slice());
421        Some(Row { id: row_id, vals })
422    }
423
424    fn get_row_column(&self, key: &[Value], col: ColumnId) -> Option<Value> {
425        assert_eq!(key.len(), 1, "attempt to lookup a row with the wrong key");
426        if col == ColumnId::new(1) {
427            Some(self.uf.find_naive(key[0]))
428        } else {
429            let row_id = *self.lookup_table.get(&key[0])?;
430            Some(self.expand(row_id)[col.index()])
431        }
432    }
433
434    fn new_buffer(&self) -> Box<dyn MutationBuffer> {
435        Box::new(UfBuffer {
436            to_insert: ManuallyDrop::new(RowBuffer::new(3)),
437            buffered_writes: Arc::downgrade(&self.buffered_writes),
438        })
439    }
440
441    fn merge(&mut self, _: &mut ExecutionState) -> TableChange {
442        while let Some(rowbuf) = self.buffered_writes.pop() {
443            for row in rowbuf.iter() {
444                self.changed |= self.insert_impl(row).is_some();
445            }
446        }
447        let changed = mem::take(&mut self.changed);
448        // UF table rows can be updated "in place", we count both added and removed as changed in
449        // this case.
450        TableChange {
451            added: changed,
452            removed: changed,
453        }
454    }
455}
456
457impl DisplacedTable {
458    pub fn underlying_uf(&self) -> &UnionFind {
459        &self.uf
460    }
461    fn expand(&self, row: RowId) -> [Value; 3] {
462        let (child, ts) = self.displaced[row.index()];
463        [child, self.uf.find_naive(child), ts]
464    }
465    fn timestamp_bounds(&self, val: Value) -> Result<(RowId, RowId), RowId> {
466        match self.displaced.binary_search_by_key(&val, |(_, ts)| *ts) {
467            Ok(mut off) => {
468                let mut next = off;
469                while off > 0 && self.displaced[off - 1].1 == val {
470                    off -= 1;
471                }
472                while next < self.displaced.len() && self.displaced[next].1 == val {
473                    next += 1;
474                }
475                Ok((RowId::from_usize(off), RowId::from_usize(next)))
476            }
477            Err(off) => Err(RowId::from_usize(off)),
478        }
479    }
480    fn eval(&self, constraint: &Constraint, row: RowId) -> bool {
481        let vals = self.expand(row);
482        eval_constraint(&vals, constraint)
483    }
484    fn insert_impl(&mut self, row: &[Value]) -> Option<(Value, Value)> {
485        assert_eq!(row.len(), 3, "attempt to insert a row with the wrong arity");
486        if self.uf.find(row[0]) == self.uf.find(row[1]) {
487            return None;
488        }
489        let (parent, child) = self.uf.union(row[0], row[1]);
490
491        // Compress paths somewhat, given that we perform naive finds everywhere else.
492        let _ = self.uf.find(parent);
493        let _ = self.uf.find(child);
494        let ts = row[2];
495        if let Some((_, highest)) = self.displaced.last() {
496            assert!(
497                *highest <= ts,
498                "must insert rows with increasing timestamps"
499            );
500        }
501        let next = RowId::from_usize(self.displaced.len());
502        self.displaced.push((child, ts));
503        self.lookup_table.insert(child, next);
504        Some((parent, child))
505    }
506}
507
508fn eval_constraint<const N: usize>(vals: &[Value; N], constraint: &Constraint) -> bool {
509    match constraint {
510        Constraint::Eq { l_col, r_col } => vals[l_col.index()] == vals[r_col.index()],
511        Constraint::EqConst { col, val } => vals[col.index()] == *val,
512        Constraint::LtConst { col, val } => vals[col.index()] < *val,
513        Constraint::GtConst { col, val } => vals[col.index()] > *val,
514        Constraint::LeConst { col, val } => vals[col.index()] <= *val,
515        Constraint::GeConst { col, val } => vals[col.index()] >= *val,
516    }
517}
egglog_core_relations/uf/mod.rs

egglog_core_relations/uf/
mod.rs
