Skip to main content

egglog_core_relations/table/
mod.rs

1//! A generic table implementation supporting sorted writes.
2//!
3//! The primary difference between this table and the `Function` implementation
4//! in egglog is that high level concepts like "timestamp" and "merge function"
5//! are abstracted away from the core functionality of the table.
6
7use std::{
8    any::Any,
9    cmp,
10    hash::Hasher,
11    mem,
12    sync::{
13        Arc, Weak,
14        atomic::{AtomicUsize, Ordering},
15    },
16};
17
18use crate::numeric_id::{DenseIdMap, NumericId};
19use crossbeam_queue::SegQueue;
20use hashbrown::HashTable;
21use rustc_hash::FxHasher;
22use sharded_hash_table::ShardedHashTable;
23
24use crate::{
25    Pooled, TableChange, TableId,
26    action::ExecutionState,
27    common::{HashMap, ShardData, ShardId, SubsetTracker, Value},
28    hash_index::{ColumnIndex, Index},
29    offsets::{OffsetRange, Offsets, RowId, SortedOffsetVector, Subset, SubsetRef},
30    parallel,
31    parallel_heuristics::parallelize_table_op,
32    pool::with_pool_set,
33    row_buffer::{ParallelRowBufWriter, RowBuffer},
34    table_spec::{
35        ColumnId, Constraint, Generation, MutationBuffer, Offset, Row, Table, TableSpec,
36        TableVersion,
37    },
38};
39
40mod rebuild;
41mod sharded_hash_table;
42#[cfg(test)]
43mod tests;
44
45// NB: Having this type def lets us switch between 64 and 32 bits of hashcode.
46//
47// We should consider just using u64 everywhere though. Hashbrown doesn't play nicely with 32-bit
48// hashcodes because it uses both the high and low bits of a 64-bit code.
49
50type HashCode = u64;
51
52/// A pointer to a row in the table.
53#[derive(Clone, Debug)]
54pub(crate) struct TableEntry {
55    hashcode: HashCode,
56    row: RowId,
57}
58
59impl TableEntry {
60    fn hashcode(&self) -> u64 {
61        // We keep the cast here to make it easy to switch to HashCode=u32.
62        #[allow(clippy::unnecessary_cast)]
63        {
64            self.hashcode as u64
65        }
66    }
67}
68
69/// The core data for a table.
70///
71/// This type is a thin wrapper around `RowBuffer`. The big difference is that
72/// it keeps track of how many stale rows are present.
73#[derive(Clone)]
74struct Rows {
75    data: RowBuffer,
76    scratch: RowBuffer,
77    stale_rows: usize,
78}
79
80impl Rows {
81    fn new(data: RowBuffer) -> Rows {
82        let arity = data.arity();
83        Rows {
84            data,
85            scratch: RowBuffer::new(arity),
86            stale_rows: 0,
87        }
88    }
89    fn clear(&mut self) {
90        self.data.clear();
91        self.stale_rows = 0;
92    }
93    fn next_row(&self) -> RowId {
94        RowId::from_usize(self.data.len())
95    }
96    fn set_stale(&mut self, row: RowId) {
97        if !self.data.set_stale(row) {
98            self.stale_rows += 1;
99        }
100    }
101
102    fn get_row(&self, row: RowId) -> Option<&[Value]> {
103        let row = self.data.get_row(row);
104        if row[0].is_stale() { None } else { Some(row) }
105    }
106
107    /// A variant of `get_row` without bounds-checking on `row`.
108    unsafe fn get_row_unchecked(&self, row: RowId) -> Option<&[Value]> {
109        let row = unsafe { self.data.get_row_unchecked(row) };
110        if row[0].is_stale() { None } else { Some(row) }
111    }
112
113    fn add_row(&mut self, row: &[Value]) -> RowId {
114        if row[0].is_stale() {
115            self.stale_rows += 1;
116        }
117        self.data.add_row(row)
118    }
119
120    fn remove_stale(&mut self, remap: impl FnMut(&[Value], RowId, RowId)) {
121        self.data.remove_stale(remap);
122        self.stale_rows = 0;
123    }
124}
125
126/// The type of closures that are used to merge values in a [`SortedWritesTable`].
127///
128/// The first argument grants access to database using an [`ExecutionState`], the second argument
129/// is the current value of the tuple. The third argument is the new, or "incoming" value of the
130/// tuple. The fourth argument is a mutable reference to a vector that will be used to store the
131/// output of the merge function _if_ it changes the value of the tuple. If it does not, then the
132/// merge function should return `false`.
133pub type MergeFn =
134    dyn Fn(&mut ExecutionState, &[Value], &[Value], &mut Vec<Value>) -> bool + Send + Sync;
135
136pub struct SortedWritesTable {
137    generation: Generation,
138    data: Rows,
139    hash: ShardedHashTable<TableEntry>,
140
141    n_keys: usize,
142    n_columns: usize,
143    sort_by: Option<ColumnId>,
144    offsets: Vec<(Value, RowId)>,
145
146    pending_state: Arc<PendingState>,
147    merge: Arc<MergeFn>,
148    to_rebuild: Vec<ColumnId>,
149    rebuild_index: Index<ColumnIndex>,
150    // Used to manage incremental rebuilds.
151    subset_tracker: SubsetTracker,
152}
153
154impl Clone for SortedWritesTable {
155    fn clone(&self) -> SortedWritesTable {
156        SortedWritesTable {
157            generation: self.generation,
158            data: self.data.clone(),
159            hash: self.hash.clone(),
160            n_keys: self.n_keys,
161            n_columns: self.n_columns,
162            sort_by: self.sort_by,
163            offsets: self.offsets.clone(),
164            pending_state: Arc::new(self.pending_state.deep_copy()),
165            merge: self.merge.clone(),
166            to_rebuild: self.to_rebuild.clone(),
167            rebuild_index: Index::new(self.to_rebuild.clone(), ColumnIndex::new()),
168            subset_tracker: Default::default(),
169        }
170    }
171}
172
173/// A variant of [`RowBuffer`] that can handle arity 0.
174///
175/// We use this to handle empty keys, where the deletion API needs to handle "row buffers of empty
176/// rows". The goal here is to keep most of the API RowBuffer-centric and avoid complicating the
177/// code too much: actual code that was optimized to handle arity 0 would look a bit different.
178#[derive(Clone)]
179enum ArbitraryRowBuffer {
180    NonEmpty(RowBuffer),
181    Empty { rows: usize },
182}
183
184impl ArbitraryRowBuffer {
185    fn new(arity: usize) -> ArbitraryRowBuffer {
186        if arity == 0 {
187            ArbitraryRowBuffer::Empty { rows: 0 }
188        } else {
189            ArbitraryRowBuffer::NonEmpty(RowBuffer::new(arity))
190        }
191    }
192
193    fn add_row(&mut self, row: &[Value]) {
194        match self {
195            ArbitraryRowBuffer::NonEmpty(buf) => {
196                buf.add_row(row);
197            }
198            ArbitraryRowBuffer::Empty { rows } => {
199                *rows += 1;
200            }
201        }
202    }
203
204    fn len(&self) -> usize {
205        match self {
206            ArbitraryRowBuffer::NonEmpty(buf) => buf.len(),
207            ArbitraryRowBuffer::Empty { rows } => *rows,
208        }
209    }
210
211    fn for_each(&self, mut f: impl FnMut(&[Value])) {
212        match self {
213            ArbitraryRowBuffer::NonEmpty(buf) => {
214                for row in buf.iter() {
215                    f(row);
216                }
217            }
218            ArbitraryRowBuffer::Empty { rows } => {
219                for _ in 0..*rows {
220                    f(&[]);
221                }
222            }
223        }
224    }
225}
226
227struct Buffer {
228    pending_rows: DenseIdMap<ShardId, RowBuffer>,
229    pending_removals: DenseIdMap<ShardId, ArbitraryRowBuffer>,
230    state: Weak<PendingState>,
231    n_cols: u32,
232    n_keys: u32,
233    shard_data: ShardData,
234}
235
236impl MutationBuffer for Buffer {
237    fn stage_insert(&mut self, row: &[Value]) {
238        let (shard, _) = hash_code(self.shard_data, row, self.n_keys as _);
239        self.pending_rows
240            .get_or_insert(shard, || RowBuffer::new(self.n_cols as _))
241            .add_row(row);
242    }
243    fn stage_remove(&mut self, key: &[Value]) {
244        let (shard, _) = hash_code(self.shard_data, key, self.n_keys as _);
245        self.pending_removals
246            .get_or_insert(shard, || ArbitraryRowBuffer::new(self.n_keys as _))
247            .add_row(key);
248    }
249    fn fresh_handle(&self) -> Box<dyn MutationBuffer> {
250        Box::new(Buffer {
251            pending_rows: Default::default(),
252            pending_removals: Default::default(),
253            state: self.state.clone(),
254            n_cols: self.n_cols,
255            n_keys: self.n_keys,
256            shard_data: self.shard_data,
257        })
258    }
259}
260
261impl Drop for Buffer {
262    fn drop(&mut self) {
263        if let Some(state) = self.state.upgrade() {
264            let mut rows = 0;
265            for shard_id in 0..self.pending_rows.n_ids() {
266                let shard = ShardId::from_usize(shard_id);
267                let Some(buf) = self.pending_rows.take(shard) else {
268                    continue;
269                };
270                rows += buf.len();
271                state.pending_rows[shard].push(buf);
272            }
273            state.total_rows.fetch_add(rows, Ordering::Relaxed);
274
275            let mut rows = 0;
276            for shard_id in 0..self.pending_removals.n_ids() {
277                let shard = ShardId::from_usize(shard_id);
278                let Some(buf) = self.pending_removals.take(shard) else {
279                    continue;
280                };
281                rows += buf.len();
282                state.pending_removals[shard].push(buf);
283            }
284            state.total_removals.fetch_add(rows, Ordering::Relaxed);
285        }
286    }
287}
288
289impl Table for SortedWritesTable {
290    fn dyn_clone(&self) -> Box<dyn Table> {
291        Box::new(self.clone())
292    }
293    fn as_any(&self) -> &dyn Any {
294        self
295    }
296    fn clear(&mut self) {
297        self.pending_state.clear();
298        if self.data.data.len() == 0 {
299            return;
300        }
301        self.offsets.clear();
302        self.data.clear();
303        self.hash.clear();
304        self.generation = Generation::from_usize(self.version().major.index() + 1);
305    }
306
307    fn spec(&self) -> TableSpec {
308        TableSpec {
309            n_keys: self.n_keys,
310            n_vals: self.n_columns - self.n_keys,
311            uncacheable_columns: Default::default(),
312            allows_delete: true,
313        }
314    }
315
316    fn apply_rebuild(
317        &mut self,
318        table_id: TableId,
319        table: &crate::WrappedTable,
320        next_ts: Value,
321        exec_state: &mut ExecutionState,
322    ) -> bool {
323        self.do_rebuild(table_id, table, next_ts, exec_state)
324    }
325
326    fn refresh_rows_for_values(&mut self, dirty_ids: &[Value], next_ts: Value) -> bool {
327        SortedWritesTable::refresh_rows_for_values(self, dirty_ids, next_ts)
328    }
329
330    fn version(&self) -> TableVersion {
331        TableVersion {
332            major: self.generation,
333            minor: Offset::from_usize(self.data.next_row().index()),
334        }
335    }
336
337    fn updates_since(&self, offset: Offset) -> Subset {
338        Subset::Dense(OffsetRange::new(
339            RowId::from_usize(offset.index()),
340            self.data.next_row(),
341        ))
342    }
343
344    fn all(&self) -> Subset {
345        Subset::Dense(OffsetRange::new(RowId::new(0), self.data.next_row()))
346    }
347
348    fn has_stale_rows(&self) -> bool {
349        self.data.stale_rows > 0
350    }
351
352    fn len(&self) -> usize {
353        self.data.data.len() - self.data.stale_rows
354    }
355
356    fn scan_generic(&self, subset: SubsetRef, mut f: impl FnMut(RowId, &[Value]))
357    where
358        Self: Sized,
359    {
360        let Some((_low, hi)) = subset.bounds() else {
361            // Empty subset
362            return;
363        };
364        assert!(
365            hi.index() <= self.data.data.len(),
366            "{} vs. {}",
367            hi.index(),
368            self.data.data.len()
369        );
370        if self.data.stale_rows == 0 {
371            // Fast path: no stale rows, skip is_stale check per row.
372            // SAFETY: subsets are sorted, low must be at most hi, and hi is less
373            // than the length of the table.
374            // TODO: provide a safe API for this `get_row_unchecked` usage since we have
375            // checked the full bounds above.
376            subset.offsets(|row| unsafe { f(row, self.data.data.get_row_unchecked(row)) })
377        } else {
378            // SAFETY: same as above.
379            subset.offsets(|row| unsafe {
380                if let Some(vals) = self.data.get_row_unchecked(row) {
381                    f(row, vals)
382                }
383            })
384        }
385    }
386
387    /// Stale rows are skipped, constrained or not, so a caller need not filter
388    /// them back out of the rows it receives. `tests/repro-stale-rows.egg`
389    /// guards the constrained path.
390    fn scan_generic_bounded(
391        &self,
392        subset: SubsetRef,
393        start: Offset,
394        n: usize,
395        cs: &[Constraint],
396        mut f: impl FnMut(RowId, &[Value]),
397    ) -> Option<Offset>
398    where
399        Self: Sized,
400    {
401        let Some((_low, hi)) = subset.bounds() else {
402            // Empty subset
403            return None;
404        };
405        assert!(
406            hi.index() <= self.data.data.len(),
407            "{} vs. {}",
408            hi.index(),
409            self.data.data.len()
410        );
411        if cs.is_empty() {
412            if self.data.stale_rows == 0 {
413                // Fast path: no stale rows, skip bounds check and is_stale check.
414                // SAFETY: all row IDs are in-bounds.
415                subset
416                    .iter_bounded(start.index(), start.index() + n, |row| {
417                        let entry = unsafe { self.data.data.get_row_unchecked(row) };
418                        f(row, entry);
419                    })
420                    .map(Offset::from_usize)
421            } else {
422                subset
423                    .iter_bounded(start.index(), start.index() + n, |row| {
424                        // SAFETY: all row IDs are in-bounds.
425                        let Some(entry) = (unsafe { self.data.get_row_unchecked(row) }) else {
426                            return;
427                        };
428                        f(row, entry);
429                    })
430                    .map(Offset::from_usize)
431            }
432        } else {
433            subset
434                .iter_bounded(start.index(), start.index() + n, |row| {
435                    // SAFETY: all row IDs are in-bounds.
436                    let Some(entry) = (unsafe { self.get_if_unchecked(cs, row) }) else {
437                        return;
438                    };
439                    f(row, entry);
440                })
441                .map(Offset::from_usize)
442        }
443    }
444
445    fn fast_subset(&self, constraint: &Constraint) -> Option<Subset> {
446        let sort_by = self.sort_by?;
447        match constraint {
448            Constraint::Eq { .. } => None,
449            Constraint::EqConst { col, val } => {
450                if col == &sort_by {
451                    match self.binary_search_sort_val(*val) {
452                        Ok((found, bound)) => Some(Subset::Dense(OffsetRange::new(found, bound))),
453                        Err(_) => Some(Subset::empty()),
454                    }
455                } else {
456                    None
457                }
458            }
459            Constraint::LtConst { col, val } => {
460                if col == &sort_by {
461                    match self.binary_search_sort_val(*val) {
462                        Ok((found, _)) => {
463                            Some(Subset::Dense(OffsetRange::new(RowId::new(0), found)))
464                        }
465                        Err(next) => Some(Subset::Dense(OffsetRange::new(RowId::new(0), next))),
466                    }
467                } else {
468                    None
469                }
470            }
471            Constraint::GtConst { col, val } => {
472                if col == &sort_by {
473                    match self.binary_search_sort_val(*val) {
474                        Ok((_, bound)) => {
475                            Some(Subset::Dense(OffsetRange::new(bound, self.data.next_row())))
476                        }
477                        Err(next) => {
478                            Some(Subset::Dense(OffsetRange::new(next, self.data.next_row())))
479                        }
480                    }
481                } else {
482                    None
483                }
484            }
485            Constraint::LeConst { col, val } => {
486                if col == &sort_by {
487                    match self.binary_search_sort_val(*val) {
488                        Ok((_, bound)) => {
489                            Some(Subset::Dense(OffsetRange::new(RowId::new(0), bound)))
490                        }
491                        Err(next) => Some(Subset::Dense(OffsetRange::new(RowId::new(0), next))),
492                    }
493                } else {
494                    None
495                }
496            }
497            Constraint::GeConst { col, val } => {
498                if col == &sort_by {
499                    match self.binary_search_sort_val(*val) {
500                        Ok((found, _)) => {
501                            Some(Subset::Dense(OffsetRange::new(found, self.data.next_row())))
502                        }
503                        Err(next) => {
504                            Some(Subset::Dense(OffsetRange::new(next, self.data.next_row())))
505                        }
506                    }
507                } else {
508                    None
509                }
510            }
511        }
512    }
513
514    fn refine_one(&self, mut subset: Subset, c: &Constraint) -> Subset {
515        // NB: we aren't using any of the `fast_subset` tricks here. We may want
516        // to if the higher-level implementations end up using it directly.
517        subset.retain(|row| self.eval(std::slice::from_ref(c), row));
518        subset
519    }
520
521    fn refine_ref(&self, subset: SubsetRef, cs: &[Constraint], _check_live: bool) -> Subset {
522        // Single fused pass: `eval` reads each row once, checking liveness
523        // (`get_row` skips stale rows) and all constraints together. A dense
524        // input whose rows all survive stays dense.
525        let n = subset.size();
526        let mut vec: Pooled<SortedOffsetVector> = with_pool_set(|ps| ps.get());
527        subset.offsets(|row| {
528            if self.eval(cs, row) {
529                // SAFETY: `offsets` visits rows in ascending order.
530                unsafe { vec.push_unchecked(row) }
531            }
532        });
533        if let SubsetRef::Dense(range) = subset
534            && vec.slice().inner().len() == n
535        {
536            return Subset::Dense(range);
537        }
538        Subset::Sparse(vec)
539    }
540
541    fn new_buffer(&self) -> Box<dyn MutationBuffer> {
542        let n_shards = self.hash.shard_data().n_shards();
543        Box::new(Buffer {
544            pending_rows: DenseIdMap::with_capacity(n_shards),
545            pending_removals: DenseIdMap::with_capacity(n_shards),
546            state: Arc::downgrade(&self.pending_state),
547            n_keys: u32::try_from(self.n_keys).expect("n_keys should fit in u32"),
548            n_cols: u32::try_from(self.n_columns).expect("n_columns should fit in u32"),
549            shard_data: self.hash.shard_data(),
550        })
551    }
552
553    fn merge(&mut self, exec_state: &mut ExecutionState) -> TableChange {
554        let removed = self.do_delete();
555        let added = self.do_insert(exec_state);
556        self.maybe_rehash();
557        TableChange { removed, added }
558    }
559
560    fn get_row(&self, key: &[Value]) -> Option<Row> {
561        let id = get_entry(key, self.n_keys, &self.hash, |row| {
562            &self.data.get_row(row).unwrap()[0..self.n_keys] == key
563        })?;
564        let mut vals = with_pool_set(|ps| ps.get::<Vec<Value>>());
565        vals.extend_from_slice(self.data.get_row(id).unwrap());
566        Some(Row { id, vals })
567    }
568
569    fn get_row_column(&self, key: &[Value], col: ColumnId) -> Option<Value> {
570        let id = get_entry(key, self.n_keys, &self.hash, |row| {
571            &self.data.get_row(row).unwrap()[0..self.n_keys] == key
572        })?;
573        Some(self.data.get_row(id).unwrap()[col.index()])
574    }
575}
576
577impl SortedWritesTable {
578    /// Create a new [`SortedWritesTable`] with the given number of keys,
579    /// columns, and an optional sort column.
580    ///
581    /// The `merge_fn` is used to evaluate conflicts when more than one row is
582    /// inserted with the same primary key. The old and new proposed values are
583    /// passed as the second and third arguments, respectively, with the
584    /// function filling the final argument with the contents of the new row.
585    /// The return value indicates whether or not the contents of the vector
586    /// should be used.
587    ///
588    /// Merge functions can access the database via [`ExecutionState`].
589    pub fn new(
590        n_keys: usize,
591        n_columns: usize,
592        sort_by: Option<ColumnId>,
593        to_rebuild: Vec<ColumnId>,
594        merge_fn: Box<MergeFn>,
595    ) -> Self {
596        let hash = ShardedHashTable::<TableEntry>::default();
597        let shard_data = hash.shard_data();
598        let rebuild_index = Index::new(to_rebuild.clone(), ColumnIndex::new());
599        SortedWritesTable {
600            generation: Generation::new(0),
601            data: Rows::new(RowBuffer::new(n_columns)),
602            hash,
603            n_keys,
604            n_columns,
605            sort_by,
606            offsets: Default::default(),
607            pending_state: Arc::new(PendingState::new(shard_data)),
608            merge: merge_fn.into(),
609            to_rebuild,
610            rebuild_index,
611            subset_tracker: Default::default(),
612        }
613    }
614
615    /// Flush all pending removals, in parallel.
616    fn parallel_delete(&mut self) -> bool {
617        let shard_data = self.hash.shard_data();
618        let pending_removals = &self.pending_state.pending_removals;
619        let data = &self.data.data;
620        let n_keys = self.n_keys;
621        let stale_delta: usize = parallel::map_mut(self.hash.mut_shards(), |shard_id, shard| {
622            let shard_id = ShardId::from_usize(shard_id);
623            if pending_removals[shard_id].is_empty() {
624                return 0;
625            }
626            let queue = &pending_removals[shard_id];
627            let mut marked_stale = 0;
628            while let Some(buf) = queue.pop() {
629                buf.for_each(|to_remove| {
630                    let (actual_shard, hc) = hash_code(shard_data, to_remove, n_keys);
631                    assert_eq!(actual_shard, shard_id);
632                    if let Ok(entry) = shard.find_entry(hc, |entry| {
633                        entry.hashcode == (hc as _)
634                            && &data.get_row(entry.row)[0..n_keys] == to_remove
635                    }) {
636                        let (ent, _) = entry.remove();
637                        // SAFETY: The safety requirements of
638                        // `set_stale_shared` are that there are no
639                        // concurrent accesses to `row`. No other threads
640                        // can access this row within this method because
641                        // different `shards` partition the space
642                        // (guaranteed by the assertion above), and we
643                        // launch at most one thread per shard.
644                        marked_stale += unsafe { !data.set_stale_shared(ent.row) } as usize;
645                    }
646                });
647            }
648            marked_stale
649        })
650        .into_iter()
651        .sum();
652        // Update the stale count with the total marked stale.
653        self.data.stale_rows += stale_delta;
654        stale_delta > 0
655    }
656    fn serial_delete(&mut self) -> bool {
657        let shard_data = self.hash.shard_data();
658        let mut changed = false;
659        self.hash
660            .mut_shards()
661            .iter_mut()
662            .enumerate()
663            .for_each(|(shard_id, shard)| {
664                let shard_id = ShardId::from_usize(shard_id);
665                let queue = &self.pending_state.pending_removals[shard_id];
666                while let Some(buf) = queue.pop() {
667                    buf.for_each(|to_remove| {
668                        let (actual_shard, hc) = hash_code(shard_data, to_remove, self.n_keys);
669                        assert_eq!(actual_shard, shard_id);
670                        if let Ok(entry) = shard.find_entry(hc, |entry| {
671                            entry.hashcode == (hc as _)
672                                && &self.data.get_row(entry.row).unwrap()[0..self.n_keys]
673                                    == to_remove
674                        }) {
675                            let (ent, _) = entry.remove();
676                            self.data.set_stale(ent.row);
677                            changed = true;
678                        }
679                    })
680                }
681            });
682        changed
683    }
684
685    fn do_delete(&mut self) -> bool {
686        let total = self.pending_state.total_removals.swap(0, Ordering::Relaxed);
687
688        if parallelize_table_op(total) {
689            self.parallel_delete()
690        } else {
691            self.serial_delete()
692        }
693    }
694
695    fn do_insert(&mut self, exec_state: &mut ExecutionState) -> bool {
696        let total = self.pending_state.total_rows.swap(0, Ordering::Relaxed);
697        self.data.data.reserve(total);
698        if parallelize_table_op(total) {
699            if let Some(col) = self.sort_by {
700                self.parallel_insert(
701                    exec_state,
702                    SortChecker {
703                        col,
704                        current: None,
705                        baseline: self.offsets.last().map(|(v, _)| *v),
706                    },
707                )
708            } else {
709                self.parallel_insert(exec_state, ())
710            }
711        } else {
712            self.serial_insert(exec_state)
713        }
714    }
715
716    fn serial_insert(&mut self, exec_state: &mut ExecutionState) -> bool {
717        let mut changed = false;
718        let n_keys = self.n_keys;
719        let mut scratch = with_pool_set(|ps| ps.get::<Vec<Value>>());
720        for (_outer_shard, queue) in self.pending_state.pending_rows.iter() {
721            if let Some(sort_by) = self.sort_by {
722                while let Some(buf) = queue.pop() {
723                    for query in buf.non_stale() {
724                        let key = &query[0..n_keys];
725                        let entry = get_entry_mut(query, n_keys, &mut self.hash, |row| {
726                            let Some(row) = self.data.get_row(row) else {
727                                return false;
728                            };
729                            &row[0..n_keys] == key
730                        });
731
732                        if let Some(row) = entry {
733                            // First case: overwriting an existing value. Apply merge
734                            // function. Insert new row and update hash table if merge
735                            // changes anything.
736                            let cur = self
737                                .data
738                                .get_row(*row)
739                                .expect("table should not point to stale entry");
740                            if (self.merge)(exec_state, cur, query, &mut scratch) {
741                                let sort_val = query[sort_by.index()];
742                                let new = self.data.add_row(&scratch);
743                                if let Some(largest) = self.offsets.last().map(|(v, _)| *v) {
744                                    assert!(
745                                        sort_val >= largest,
746                                        "inserting row that violates sort order ({sort_val:?} vs. {largest:?})"
747                                    );
748                                    if sort_val > largest {
749                                        self.offsets.push((sort_val, new));
750                                    }
751                                } else {
752                                    self.offsets.push((sort_val, new));
753                                }
754                                self.data.set_stale(*row);
755                                *row = new;
756                                changed = true;
757                            }
758                            scratch.clear();
759                        } else {
760                            let sort_val = query[sort_by.index()];
761                            // New value: update invariants.
762                            let new = self.data.add_row(query);
763                            if let Some(largest) = self.offsets.last().map(|(v, _)| *v) {
764                                assert!(
765                                    sort_val >= largest,
766                                    "inserting row that violates sort order {sort_val:?} vs. {largest:?}"
767                                );
768                                if sort_val > largest {
769                                    self.offsets.push((sort_val, new));
770                                }
771                            } else {
772                                self.offsets.push((sort_val, new));
773                            }
774                            let (shard, hc) = hash_code(self.hash.shard_data(), query, self.n_keys);
775                            debug_assert_eq!(shard, _outer_shard);
776                            self.hash.mut_shards()[shard.index()].insert_unique(
777                                hc as _,
778                                TableEntry {
779                                    hashcode: hc as _,
780                                    row: new,
781                                },
782                                TableEntry::hashcode,
783                            );
784                            changed = true;
785                        }
786                    }
787                }
788            } else {
789                // Simplified variant without the sorting constraint.
790                while let Some(buf) = queue.pop() {
791                    for query in buf.non_stale() {
792                        let key = &query[0..n_keys];
793                        let entry = get_entry_mut(query, n_keys, &mut self.hash, |row| {
794                            let Some(row) = self.data.get_row(row) else {
795                                return false;
796                            };
797                            &row[0..n_keys] == key
798                        });
799
800                        if let Some(row) = entry {
801                            let cur = self
802                                .data
803                                .get_row(*row)
804                                .expect("table should not point to stale entry");
805                            if (self.merge)(exec_state, cur, query, &mut scratch) {
806                                let new = self.data.add_row(&scratch);
807                                self.data.set_stale(*row);
808                                *row = new;
809                                changed = true;
810                            }
811                            scratch.clear();
812                        } else {
813                            // New value: update invariants.
814                            let new = self.data.add_row(query);
815                            let (shard, hc) = hash_code(self.hash.shard_data(), query, self.n_keys);
816                            debug_assert_eq!(shard, _outer_shard);
817                            self.hash.mut_shards()[shard.index()].insert_unique(
818                                hc as _,
819                                TableEntry {
820                                    hashcode: hc as _,
821                                    row: new,
822                                },
823                                TableEntry::hashcode,
824                            );
825                            changed = true;
826                        }
827                    }
828                }
829            };
830        }
831        changed
832    }
833
834    fn parallel_insert<C: OrderingChecker>(
835        &mut self,
836        exec_state: &ExecutionState,
837        checker: C,
838    ) -> bool {
839        const BATCH_SIZE: usize = 1 << 18;
840        // Parallel insert uses one giant parallel foreach. We have updates
841        // pre-sharded, and one logical thread can process updates for each
842        // shard independently. Updates happen in three phases, which comments
843        // describe below.
844        let shard_data = self.hash.shard_data();
845        let n_keys = self.n_keys;
846        let n_cols = self.n_columns;
847        let next_offset = RowId::from_usize(self.data.data.len());
848        let row_writer = self.data.data.parallel_writer();
849        let pending_rows = &self.pending_state.pending_rows;
850        let merge = self.merge.clone();
851        let pending_adds = parallel::map_mut(self.hash.mut_shards(), |shard_id, shard| {
852            let shard_id = ShardId::from_usize(shard_id);
853            let mut checker = checker.clone();
854            let mut exec_state = exec_state.clone();
855            let mut scratch = with_pool_set(|ps| ps.get::<Vec<Value>>());
856            let queue = &pending_rows[shard_id];
857            let mut marked_stale = 0usize;
858            let mut staged = StagedOutputs::new(n_keys, n_cols, BATCH_SIZE);
859            let mut changed = false;
860            // The core flush loop: We call once `staged` reaches `BATCH_SIZE` or
861            // when we're done.
862            macro_rules! flush_staged_outputs {
863                    () => {{
864                        // Phase 2: Write the staged rows to the row writer. This only
865                        // works due to the `ParallelRowBufWriter` machinery.
866                        let (start_row, stale) = staged.write_output(&row_writer);
867                        marked_stale += stale;
868                        // Phase 3: With the values buffered in the row buffer, we can
869                        // write them back to the shard, pointed to the correct rows.
870
871                        // In the serial implementation, we do phases 2 and 3 inline with
872                        // processing the incoming mutation, but separating them out
873                        // this way allows us to do a single write to the shared row
874                        // buffer, rather than one per row, which would cause
875                        // contention.
876                        let mut cur_row = start_row;
877                        let read_handle = row_writer.read_handle();
878                        for row in staged.rows() {
879                            if row.first().map(Value::is_stale).unwrap_or(false) {
880                                cur_row = cur_row.inc();
881                                continue;
882                            }
883                            use hashbrown::hash_table::Entry;
884                            checker.check_local(row);
885                            changed = true;
886                            let key = &row[0..n_keys];
887                            let (_actual_shard, hc) = hash_code(shard_data, row, n_keys);
888                            #[cfg(any(debug_assertions, test))]
889                            {
890                                unsafe {
891                                    // read the value we wrote at this row and
892                                    // check that it matches.
893                                    assert_eq!(read_handle.get_row_unchecked(cur_row), row);
894                                }
895                            }
896                            debug_assert_eq!(_actual_shard, shard_id);
897                            match shard.entry(
898                                hc,
899                                // SAFETY: `ent` must point to a valid row
900                                |ent| unsafe {
901                                    ent.hashcode == hc as HashCode
902                                        && &read_handle.get_row_unchecked(ent.row)[0..n_keys] == key
903                                },
904                                TableEntry::hashcode,
905                            ) {
906                                Entry::Occupied(mut occ) => {
907                                    // SAFETY: `occ` must point to a valid row: we only insert valid rows
908                                    // into the map.
909                                    let cur = unsafe { read_handle.get_row_unchecked(occ.get().row) };
910
911                                    // SAFETY: The safety requirements of
912                                    // `set_stale_shared` are that there are no
913                                    // concurrent accesses to `row`. We have
914                                    // exclusive access to any row whose hash matches this
915                                    // shard.
916                                    if (merge)(&mut exec_state, cur, row, &mut scratch) {
917                                        unsafe {
918                                            let _was_stale = read_handle.set_stale_shared(occ.get().row);
919                                            debug_assert!(!_was_stale);
920                                        }
921                                        occ.get_mut().row = cur_row;
922                                        changed = true;
923                                    } else {
924                                        // Mark the new row as stale: we didn't end up needing it.
925                                        unsafe {
926                                            let _was_stale = read_handle.set_stale_shared(cur_row);
927                                            debug_assert!(!_was_stale);
928                                        }
929                                    }
930                                    marked_stale += 1;
931                                    scratch.clear();
932                                }
933                                Entry::Vacant(v) => {
934                                    changed = true;
935                                    v.insert(TableEntry {
936                                        hashcode: hc as HashCode,
937                                        row: cur_row,
938                                    });
939                                }
940                            }
941
942                            cur_row = cur_row.inc();
943                        }
944                        staged.clear();
945                    }};
946                }
947            // Phase 1: process all incoming updates:
948            // * Add new values to `staged`
949            // * Removing entries in `shard` and mark them as stale in
950            // `data` if they will be overwritten.
951            while let Some(buf) = queue.pop() {
952                // We create a read_handle once per batch to avoid blocking
953                // too many threads if someone needs to resize the row
954                // writer.
955                for row in buf.non_stale() {
956                    staged.insert(row, |cur, new, out| (merge)(&mut exec_state, cur, new, out));
957                    if staged.len() >= BATCH_SIZE {
958                        flush_staged_outputs!();
959                    }
960                }
961            }
962            flush_staged_outputs!();
963            (checker, marked_stale, changed)
964        });
965        self.data.data = row_writer.finish();
966        // Now we just need to reset our invariants.
967
968        // Confirm none of the writes violated sort order and update the
969        // `offsets` vector.
970        let checker = C::check_global(pending_adds.iter().map(|(checker, _, _)| checker));
971        checker.update_offsets(next_offset, &mut self.offsets);
972
973        // Update the staleness counters.
974        self.data.stale_rows += pending_adds
975            .iter()
976            .map(|(_, stale, _)| *stale)
977            .sum::<usize>();
978
979        // Register any changes.
980        pending_adds.iter().any(|(_, _, changed)| *changed)
981    }
982
983    fn binary_search_sort_val(&self, val: Value) -> Result<(RowId, RowId), RowId> {
984        debug_assert!(
985            self.offsets.windows(2).all(|x| x[0].1 < x[1].1),
986            "{:?}",
987            self.offsets
988        );
989
990        debug_assert!(
991            self.offsets.windows(2).all(|x| x[0].0 < x[1].0),
992            "{:?}",
993            self.offsets
994        );
995        match self.offsets.binary_search_by_key(&val, |(v, _)| *v) {
996            Ok(got) => Ok((
997                self.offsets[got].1,
998                self.offsets
999                    .get(got + 1)
1000                    .map(|(_, r)| *r)
1001                    .unwrap_or(self.data.next_row()),
1002            )),
1003            Err(next) => Err(self
1004                .offsets
1005                .get(next)
1006                .map(|(_, id)| *id)
1007                .unwrap_or(self.data.next_row())),
1008        }
1009    }
1010    fn eval(&self, cs: &[Constraint], row: RowId) -> bool {
1011        self.get_if(cs, row).is_some()
1012    }
1013
1014    fn eval_constraints(cs: &[Constraint], row: &[Value]) -> bool {
1015        cs.iter().all(|constraint| match constraint {
1016            Constraint::Eq { l_col, r_col } => row[l_col.index()] == row[r_col.index()],
1017            Constraint::EqConst { col, val } => row[col.index()] == *val,
1018            Constraint::LtConst { col, val } => row[col.index()] < *val,
1019            Constraint::GtConst { col, val } => row[col.index()] > *val,
1020            Constraint::LeConst { col, val } => row[col.index()] <= *val,
1021            Constraint::GeConst { col, val } => row[col.index()] >= *val,
1022        })
1023    }
1024
1025    unsafe fn get_if_unchecked(&self, cs: &[Constraint], row: RowId) -> Option<&[Value]> {
1026        let row = unsafe { self.data.get_row_unchecked(row) }?;
1027        if Self::eval_constraints(cs, row) {
1028            Some(row)
1029        } else {
1030            None
1031        }
1032    }
1033
1034    fn get_if(&self, cs: &[Constraint], row: RowId) -> Option<&[Value]> {
1035        let row = self.data.get_row(row)?;
1036        if Self::eval_constraints(cs, row) {
1037            Some(row)
1038        } else {
1039            None
1040        }
1041    }
1042
1043    fn maybe_rehash(&mut self) {
1044        if self.data.stale_rows <= cmp::max(16, self.data.data.len() / 2) {
1045            return;
1046        }
1047
1048        if parallelize_table_op(self.data.data.len()) {
1049            self.parallel_rehash();
1050        } else {
1051            self.rehash();
1052        }
1053    }
1054    fn parallel_rehash(&mut self) {
1055        // Parallel rehashes go "hash-first" rather than "rows-first".
1056        //
1057        // We iterate over each shard and then write out new contents to a fresh row, in parallel.
1058        let Some(sort_by) = self.sort_by else {
1059            // Just do a serial rehash for now. We currently do not have a use-case for parallel
1060            // compaction of unsorted tables.
1061            //
1062            // Implementing parallel compaction for an unsorted table is much easier: each shard
1063            // can write to a contiguous chunk of the `scratch` buffer, with the offsets being
1064            // pre-chunked based on the size of each shard.
1065            self.rehash();
1066            return;
1067        };
1068        self.generation = self.generation.inc();
1069        assert!(!self.offsets.is_empty());
1070        struct TimestampStats {
1071            value: Value,
1072            count: usize,
1073            histogram: Pooled<DenseIdMap<ShardId, usize>>,
1074        }
1075        impl Default for TimestampStats {
1076            fn default() -> TimestampStats {
1077                TimestampStats {
1078                    value: Value::stale(),
1079                    count: 0,
1080                    histogram: with_pool_set(|ps| ps.get()),
1081                }
1082            }
1083        }
1084        let mut results = Vec::<TimestampStats>::with_capacity(self.offsets.len());
1085        let offset_windows = self
1086            .offsets
1087            .windows(2)
1088            .map(|xs| {
1089                let [(start_val, start_row), (_, end_row)] = xs else {
1090                    unreachable!()
1091                };
1092                (*start_val, *start_row, *end_row)
1093            })
1094            .collect::<Vec<_>>();
1095        // Use a macro rather than a lambda to avoid borrow issues.
1096        macro_rules! compute_hist {
1097            ($start_val: expr, $start_row: expr, $end_row: expr) => {{
1098                let mut histogram: Pooled<DenseIdMap<ShardId, usize>> =
1099                    with_pool_set(|ps| ps.get());
1100                let mut cur_row = $start_row;
1101                let mut count = 0;
1102                while cur_row < $end_row {
1103                    if let Some(row) = self.data.get_row(cur_row) {
1104                        count += 1;
1105                        let (shard, _) = hash_code(self.hash.shard_data(), row, self.n_keys);
1106                        *histogram.get_or_default(shard) += 1;
1107                    }
1108                    cur_row = cur_row.inc();
1109                }
1110                TimestampStats {
1111                    value: $start_val,
1112                    count,
1113                    histogram,
1114                }
1115            }};
1116        }
1117        results.extend(parallel::map(
1118            &offset_windows,
1119            |_, (start_val, start_row, end_row)| compute_hist!(*start_val, *start_row, *end_row),
1120        ));
1121        // And here we handle the final one.
1122        let (start_val, start_row) = self.offsets.last().unwrap();
1123        let end_row = self.data.next_row();
1124        let last = compute_hist!(*start_val, *start_row, end_row);
1125        results.push(last);
1126        // Now we need to compute cumulative statistics on the row layouts here.
1127        // We do this serially a we currently don't have a ton of use for cases with thousands
1128        // of timestamps or more. There are well-known parallel algorithms for computing these
1129        // cumulative statistics in parallel, but they are not currently a good fit here.
1130        let mut prev_count = 0;
1131        self.offsets.clear();
1132        for stats in results.iter_mut() {
1133            if stats.count == 0 {
1134                continue;
1135            }
1136            self.offsets
1137                .push((stats.value, RowId::from_usize(prev_count)));
1138            let mut inner = prev_count;
1139            for (_, count) in stats.histogram.iter_mut() {
1140                // Each entry in the histogram now points to the start row for that shard's
1141                // rows for a given timestamp.
1142                let tmp = *count;
1143                *count = inner;
1144                inner += tmp;
1145            }
1146            prev_count += stats.count;
1147            debug_assert_eq!(inner, prev_count)
1148        }
1149
1150        // Now the part with some unsafe code.
1151        // We will iterate over each shard and use the statistics in `results` to guide where
1152        // each row will go.
1153        //
1154        // This involves doing unsynchronized writes to the table (ptr::copy_nonoverlapping)
1155        // followed by a set_len. The safety of these operations relies on the fact that:
1156        // * No one grabs a reference to the interior of `scratch` until these operations have
1157        //   finished.
1158        // * `scratch` does not overlap `data`.
1159        // * The sharding function completely partitions the set of objects in the table: one
1160        //   shard's writes will never stomp on those of another.
1161
1162        self.data.scratch.clear();
1163        self.data.scratch.reserve(prev_count);
1164        let scratch_ptr = self.data.scratch.raw_rows() as usize;
1165        let data = &self.data.data;
1166        let n_columns = self.n_columns;
1167        parallel::for_each_mut(self.hash.mut_shards(), |shard_id, shard| {
1168            let shard_id = ShardId::from_usize(shard_id);
1169            let scratch_ptr = scratch_ptr as *const Value;
1170            let mut progress = HashMap::<Value /* timestamp */, RowId /* next row */>::default();
1171            progress.reserve(results.len());
1172            for stats in &results {
1173                let Some(start) = stats.histogram.get(shard_id) else {
1174                    continue;
1175                };
1176                progress.insert(stats.value, RowId::from_usize(*start));
1177            }
1178            for TableEntry { row: row_id, .. } in shard.iter_mut() {
1179                let row = data.get_row(*row_id);
1180                debug_assert!(!row[0].is_stale(), "shard should not map to a stale value");
1181                let val = row[sort_by.index()];
1182                let next = progress[&val];
1183                // SAFETY: see above longer comment.
1184                unsafe {
1185                    std::ptr::copy_nonoverlapping(
1186                        row.as_ptr(),
1187                        scratch_ptr.add(next.index() * n_columns) as *mut Value,
1188                        n_columns,
1189                    )
1190                }
1191                *row_id = next;
1192                progress.insert(val, next.inc());
1193            }
1194        });
1195        // SAFETY: see above longer comment.
1196        unsafe { self.data.scratch.set_len(prev_count) };
1197        mem::swap(&mut self.data.data, &mut self.data.scratch);
1198        self.data.stale_rows = 0;
1199    }
1200    fn rehash_impl(
1201        sort_by: Option<ColumnId>,
1202        n_keys: usize,
1203        rows: &mut Rows,
1204        offsets: &mut Vec<(Value, RowId)>,
1205        hash: &mut ShardedHashTable<TableEntry>,
1206    ) {
1207        if let Some(sort_by) = sort_by {
1208            offsets.clear();
1209            rows.remove_stale(|row, old, new| {
1210                let stale_entry = get_entry_mut(row, n_keys, hash, |x| x == old)
1211                    .expect("non-stale entry not mapped in hash");
1212                *stale_entry = new;
1213                let sort_col = row[sort_by.index()];
1214                if let Some((max, _)) = offsets.last() {
1215                    if sort_col > *max {
1216                        offsets.push((sort_col, new));
1217                    }
1218                } else {
1219                    offsets.push((sort_col, new));
1220                }
1221            })
1222        } else {
1223            rows.remove_stale(|row, old, new| {
1224                let stale_entry = get_entry_mut(row, n_keys, hash, |x| x == old)
1225                    .expect("non-stale entry not mapped in hash");
1226                *stale_entry = new;
1227            })
1228        }
1229    }
1230
1231    fn rehash(&mut self) {
1232        self.generation = self.generation.inc();
1233        Self::rehash_impl(
1234            self.sort_by,
1235            self.n_keys,
1236            &mut self.data,
1237            &mut self.offsets,
1238            &mut self.hash,
1239        )
1240    }
1241}
1242
1243fn get_entry(
1244    row: &[Value],
1245    n_keys: usize,
1246    table: &ShardedHashTable<TableEntry>,
1247    test: impl Fn(RowId) -> bool,
1248) -> Option<RowId> {
1249    let (shard, hash) = hash_code(table.shard_data(), row, n_keys);
1250    table
1251        .get_shard(shard)
1252        .find(hash, |ent| {
1253            ent.hashcode == hash as HashCode && test(ent.row)
1254        })
1255        .map(|ent| ent.row)
1256}
1257
1258fn get_entry_mut<'a>(
1259    row: &[Value],
1260    n_keys: usize,
1261    table: &'a mut ShardedHashTable<TableEntry>,
1262    test: impl Fn(RowId) -> bool,
1263) -> Option<&'a mut RowId> {
1264    let (shard, hash) = hash_code(table.shard_data(), row, n_keys);
1265    table.mut_shards()[shard.index()]
1266        .find_mut(hash, |ent| {
1267            ent.hashcode == hash as HashCode && test(ent.row)
1268        })
1269        .map(|ent| &mut ent.row)
1270}
1271
1272fn hash_code(shard_data: ShardData, row: &[Value], n_keys: usize) -> (ShardId, u64) {
1273    let mut hasher = FxHasher::default();
1274    for val in &row[0..n_keys] {
1275        hasher.write_usize(val.index());
1276    }
1277    let full_code = hasher.finish();
1278    // We keep this cast here to allow for experimenting with HashCode=u32.
1279    #[allow(clippy::unnecessary_cast)]
1280    (shard_data.shard_id(full_code), full_code as HashCode as u64)
1281}
1282
1283/// A simple struct for packaging up pending mutations to a `SortedWritesTable`.
1284struct PendingState {
1285    pending_rows: DenseIdMap<ShardId, SegQueue<RowBuffer>>,
1286    pending_removals: DenseIdMap<ShardId, SegQueue<ArbitraryRowBuffer>>,
1287    total_removals: AtomicUsize,
1288    total_rows: AtomicUsize,
1289}
1290
1291impl PendingState {
1292    fn new(shard_data: ShardData) -> PendingState {
1293        let n_shards = shard_data.n_shards();
1294        let mut pending_rows = DenseIdMap::with_capacity(n_shards);
1295        let mut pending_removals = DenseIdMap::with_capacity(n_shards);
1296        for i in 0..n_shards {
1297            pending_rows.insert(ShardId::from_usize(i), SegQueue::default());
1298            pending_removals.insert(ShardId::from_usize(i), SegQueue::default());
1299        }
1300
1301        PendingState {
1302            pending_rows,
1303            pending_removals,
1304            total_removals: AtomicUsize::new(0),
1305            total_rows: AtomicUsize::new(0),
1306        }
1307    }
1308    fn clear(&self) {
1309        for (_, queue) in self.pending_rows.iter() {
1310            while queue.pop().is_some() {}
1311        }
1312
1313        for (_, queue) in self.pending_removals.iter() {
1314            while queue.pop().is_some() {}
1315        }
1316    }
1317
1318    /// This is only really used in debugging, but it's annoying enough to write
1319    /// that it may help to have around.
1320    ///
1321    /// We also, however, use it in the clone impl (which should only be called when pending state
1322    /// is empty).
1323    fn deep_copy(&self) -> PendingState {
1324        let mut pending_rows = DenseIdMap::new();
1325        let mut pending_removals = DenseIdMap::new();
1326        fn drain_queue<T>(queue: &SegQueue<T>) -> Vec<T> {
1327            let mut res = Vec::new();
1328            while let Some(x) = queue.pop() {
1329                res.push(x);
1330            }
1331            res
1332        }
1333        for (shard, queue) in self.pending_rows.iter() {
1334            let contents = drain_queue(queue);
1335            let new_queue = SegQueue::default();
1336            for x in contents {
1337                new_queue.push(x.clone());
1338                queue.push(x);
1339            }
1340            pending_rows.insert(shard, new_queue);
1341        }
1342
1343        for (shard, queue) in self.pending_removals.iter() {
1344            let contents = drain_queue(queue);
1345            let new_queue = SegQueue::default();
1346            for x in contents {
1347                new_queue.push(x.clone());
1348                queue.push(x);
1349            }
1350            pending_removals.insert(shard, new_queue);
1351        }
1352
1353        PendingState {
1354            pending_rows,
1355            pending_removals,
1356            total_removals: AtomicUsize::new(self.total_removals.load(Ordering::Acquire)),
1357            total_rows: AtomicUsize::new(self.total_rows.load(Ordering::Acquire)),
1358        }
1359    }
1360}
1361
1362/// A trait that encapsulates the logic of potentially checking that written
1363/// columns appear in sorted order.
1364///
1365/// For rows that are sorted by a column, an OrderingChecker asserts that all
1366/// new rows have the same value in that column, and that the column is greater
1367/// than or equal to the column value coming in. For rows not sorted, these
1368/// checks become no-ops.
1369trait OrderingChecker: Clone + Send + Sync {
1370    /// Check any invariants locally, updating the state of the checker when
1371    /// doing so.
1372    fn check_local(&mut self, row: &[Value]);
1373    /// Combine the states of multiple checkers, returning a new checker with
1374    /// all information assimilated. This is the checker that is suitable for
1375    /// calling `update_offsets` with.
1376    fn check_global<'a>(checkers: impl Iterator<Item = &'a Self>) -> Self
1377    where
1378        Self: 'a;
1379    /// Update the sorted offset vector with the current state of the checker.
1380    fn update_offsets(&self, start: RowId, offsets: &mut Vec<(Value, RowId)>);
1381}
1382
1383impl OrderingChecker for () {
1384    fn check_local(&mut self, _: &[Value]) {}
1385    fn check_global<'a>(_: impl Iterator<Item = &'a ()>) {}
1386    fn update_offsets(&self, _: RowId, _: &mut Vec<(Value, RowId)>) {}
1387}
1388
1389#[derive(Copy, Clone)]
1390struct SortChecker {
1391    col: ColumnId,
1392    baseline: Option<Value>,
1393    current: Option<Value>,
1394}
1395
1396impl OrderingChecker for SortChecker {
1397    fn check_local(&mut self, row: &[Value]) {
1398        let val = row[self.col.index()];
1399        if let Some(cur) = self.current {
1400            assert_eq!(
1401                cur, val,
1402                "concurrently inserting rows with different sort keys"
1403            );
1404        } else {
1405            self.current = Some(val);
1406            if let Some(baseline) = self.baseline {
1407                assert!(val >= baseline, "inserted row violates sort order");
1408            }
1409        }
1410    }
1411
1412    fn check_global<'a>(mut checkers: impl Iterator<Item = &'a Self>) -> Self {
1413        let Some(start) = checkers.next() else {
1414            return SortChecker {
1415                col: ColumnId::new(!0),
1416                baseline: None,
1417                current: None,
1418            };
1419        };
1420        let mut expected = start.current;
1421        for checker in checkers {
1422            assert_eq!(checker.baseline, start.baseline);
1423            match (&mut expected, checker.current) {
1424                (None, None) => {}
1425                (cur @ None, Some(x)) => {
1426                    *cur = Some(x);
1427                }
1428                (Some(_), None) => {}
1429                (Some(x), Some(y)) => {
1430                    assert_eq!(
1431                        *x, y,
1432                        "concurrently inserting rows with different sort keys"
1433                    );
1434                }
1435            }
1436        }
1437        SortChecker {
1438            col: start.col,
1439            baseline: start.baseline,
1440            current: expected,
1441        }
1442    }
1443
1444    fn update_offsets(&self, start: RowId, offsets: &mut Vec<(Value, RowId)>) {
1445        if let Some(cur) = self.current {
1446            if let Some((max, _)) = offsets.last() {
1447                if cur > *max {
1448                    offsets.push((cur, start));
1449                }
1450            } else {
1451                offsets.push((cur, start));
1452            }
1453        }
1454    }
1455}
1456
1457/// A type similar to a SortedWritesTable used to buffer outputs. The main thing
1458/// that StagedOutputs handles is running the merge function for a table on
1459/// multiple updates to the same key that show up in the same round of
1460/// insertions.
1461struct StagedOutputs {
1462    shard_data: ShardData,
1463    n_keys: usize,
1464    hash: Pooled<HashTable<TableEntry>>,
1465    rows: RowBuffer,
1466    n_stale: usize,
1467    scratch: Pooled<Vec<Value>>,
1468}
1469
1470impl StagedOutputs {
1471    fn rows(&self) -> impl Iterator<Item = &[Value]> {
1472        self.rows.iter()
1473    }
1474    fn new(n_keys: usize, n_cols: usize, capacity: usize) -> Self {
1475        let mut res = with_pool_set(|ps| StagedOutputs {
1476            shard_data: ShardData::new(1),
1477            n_keys,
1478            n_stale: 0,
1479            hash: ps.get(),
1480            rows: RowBuffer::new(n_cols),
1481            scratch: ps.get(),
1482        });
1483        res.hash.reserve(capacity, TableEntry::hashcode);
1484        res.rows.reserve(capacity);
1485        res
1486    }
1487    fn clear(&mut self) {
1488        self.hash.clear();
1489        self.rows.clear();
1490        self.n_stale = 0;
1491    }
1492    fn len(&self) -> usize {
1493        self.rows.len() - self.n_stale
1494    }
1495
1496    fn insert(
1497        &mut self,
1498        row: &[Value],
1499        mut merge_fn: impl FnMut(&[Value], &[Value], &mut Vec<Value>) -> bool,
1500    ) {
1501        if row[0].is_stale() {
1502            return;
1503        }
1504        use hashbrown::hash_table::Entry;
1505        let (_, hc) = hash_code(self.shard_data, row, self.n_keys);
1506        let entry = self.hash.entry(
1507            hc,
1508            |te| {
1509                te.hashcode() == hc
1510                    && self.rows.get_row(te.row)[0..self.n_keys] == row[0..self.n_keys]
1511            },
1512            TableEntry::hashcode,
1513        );
1514        match entry {
1515            Entry::Occupied(mut occupied_entry) => {
1516                let cur = self.rows.get_row(occupied_entry.get().row);
1517                if merge_fn(cur, row, &mut self.scratch) {
1518                    let new = self.rows.add_row(&self.scratch);
1519                    self.rows.set_stale(occupied_entry.get().row);
1520                    self.n_stale += 1;
1521                    occupied_entry.get_mut().row = new;
1522                }
1523                self.scratch.clear();
1524            }
1525            Entry::Vacant(vacant_entry) => {
1526                let next = self.rows.add_row(row);
1527                vacant_entry.insert(TableEntry {
1528                    hashcode: hc as _,
1529                    row: next,
1530                });
1531            }
1532        }
1533    }
1534
1535    /// Write the contents of the staged outputs to the given writer, returning the initial RowId
1536    /// of the new output. Returns the number of stale values in the buffer that was appended.
1537    fn write_output(&self, output: &ParallelRowBufWriter) -> (RowId, usize) {
1538        (output.append_contents(&self.rows), self.n_stale)
1539    }
1540}