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    fn scan_generic_bounded(
388        &self,
389        subset: SubsetRef,
390        start: Offset,
391        n: usize,
392        cs: &[Constraint],
393        mut f: impl FnMut(RowId, &[Value]),
394    ) -> Option<Offset>
395    where
396        Self: Sized,
397    {
398        let Some((_low, hi)) = subset.bounds() else {
399            // Empty subset
400            return None;
401        };
402        assert!(
403            hi.index() <= self.data.data.len(),
404            "{} vs. {}",
405            hi.index(),
406            self.data.data.len()
407        );
408        if cs.is_empty() {
409            if self.data.stale_rows == 0 {
410                // Fast path: no stale rows, skip bounds check and is_stale check.
411                // SAFETY: all row IDs are in-bounds.
412                subset
413                    .iter_bounded(start.index(), start.index() + n, |row| {
414                        let entry = unsafe { self.data.data.get_row_unchecked(row) };
415                        f(row, entry);
416                    })
417                    .map(Offset::from_usize)
418            } else {
419                subset
420                    .iter_bounded(start.index(), start.index() + n, |row| {
421                        // SAFETY: all row IDs are in-bounds.
422                        let Some(entry) = (unsafe { self.data.get_row_unchecked(row) }) else {
423                            return;
424                        };
425                        f(row, entry);
426                    })
427                    .map(Offset::from_usize)
428            }
429        } else {
430            subset
431                .iter_bounded(start.index(), start.index() + n, |row| {
432                    // SAFETY: all row IDs are in-bounds.
433                    let Some(entry) = (unsafe { self.get_if_unchecked(cs, row) }) else {
434                        return;
435                    };
436                    f(row, entry);
437                })
438                .map(Offset::from_usize)
439        }
440    }
441
442    fn fast_subset(&self, constraint: &Constraint) -> Option<Subset> {
443        let sort_by = self.sort_by?;
444        match constraint {
445            Constraint::Eq { .. } => None,
446            Constraint::EqConst { col, val } => {
447                if col == &sort_by {
448                    match self.binary_search_sort_val(*val) {
449                        Ok((found, bound)) => Some(Subset::Dense(OffsetRange::new(found, bound))),
450                        Err(_) => Some(Subset::empty()),
451                    }
452                } else {
453                    None
454                }
455            }
456            Constraint::LtConst { col, val } => {
457                if col == &sort_by {
458                    match self.binary_search_sort_val(*val) {
459                        Ok((found, _)) => {
460                            Some(Subset::Dense(OffsetRange::new(RowId::new(0), found)))
461                        }
462                        Err(next) => Some(Subset::Dense(OffsetRange::new(RowId::new(0), next))),
463                    }
464                } else {
465                    None
466                }
467            }
468            Constraint::GtConst { col, val } => {
469                if col == &sort_by {
470                    match self.binary_search_sort_val(*val) {
471                        Ok((_, bound)) => {
472                            Some(Subset::Dense(OffsetRange::new(bound, self.data.next_row())))
473                        }
474                        Err(next) => {
475                            Some(Subset::Dense(OffsetRange::new(next, self.data.next_row())))
476                        }
477                    }
478                } else {
479                    None
480                }
481            }
482            Constraint::LeConst { col, val } => {
483                if col == &sort_by {
484                    match self.binary_search_sort_val(*val) {
485                        Ok((_, bound)) => {
486                            Some(Subset::Dense(OffsetRange::new(RowId::new(0), bound)))
487                        }
488                        Err(next) => Some(Subset::Dense(OffsetRange::new(RowId::new(0), next))),
489                    }
490                } else {
491                    None
492                }
493            }
494            Constraint::GeConst { col, val } => {
495                if col == &sort_by {
496                    match self.binary_search_sort_val(*val) {
497                        Ok((found, _)) => {
498                            Some(Subset::Dense(OffsetRange::new(found, self.data.next_row())))
499                        }
500                        Err(next) => {
501                            Some(Subset::Dense(OffsetRange::new(next, self.data.next_row())))
502                        }
503                    }
504                } else {
505                    None
506                }
507            }
508        }
509    }
510
511    fn refine_one(&self, mut subset: Subset, c: &Constraint) -> Subset {
512        // NB: we aren't using any of the `fast_subset` tricks here. We may want
513        // to if the higher-level implementations end up using it directly.
514        subset.retain(|row| self.eval(std::slice::from_ref(c), row));
515        subset
516    }
517
518    fn refine_ref(&self, subset: SubsetRef, cs: &[Constraint], _check_live: bool) -> Subset {
519        // Single fused pass: `eval` reads each row once, checking liveness
520        // (`get_row` skips stale rows) and all constraints together. A dense
521        // input whose rows all survive stays dense.
522        let n = subset.size();
523        let mut vec: Pooled<SortedOffsetVector> = with_pool_set(|ps| ps.get());
524        subset.offsets(|row| {
525            if self.eval(cs, row) {
526                // SAFETY: `offsets` visits rows in ascending order.
527                unsafe { vec.push_unchecked(row) }
528            }
529        });
530        if let SubsetRef::Dense(range) = subset
531            && vec.slice().inner().len() == n
532        {
533            return Subset::Dense(range);
534        }
535        Subset::Sparse(vec)
536    }
537
538    fn new_buffer(&self) -> Box<dyn MutationBuffer> {
539        let n_shards = self.hash.shard_data().n_shards();
540        Box::new(Buffer {
541            pending_rows: DenseIdMap::with_capacity(n_shards),
542            pending_removals: DenseIdMap::with_capacity(n_shards),
543            state: Arc::downgrade(&self.pending_state),
544            n_keys: u32::try_from(self.n_keys).expect("n_keys should fit in u32"),
545            n_cols: u32::try_from(self.n_columns).expect("n_columns should fit in u32"),
546            shard_data: self.hash.shard_data(),
547        })
548    }
549
550    fn merge(&mut self, exec_state: &mut ExecutionState) -> TableChange {
551        let removed = self.do_delete();
552        let added = self.do_insert(exec_state);
553        self.maybe_rehash();
554        TableChange { removed, added }
555    }
556
557    fn get_row(&self, key: &[Value]) -> Option<Row> {
558        let id = get_entry(key, self.n_keys, &self.hash, |row| {
559            &self.data.get_row(row).unwrap()[0..self.n_keys] == key
560        })?;
561        let mut vals = with_pool_set(|ps| ps.get::<Vec<Value>>());
562        vals.extend_from_slice(self.data.get_row(id).unwrap());
563        Some(Row { id, vals })
564    }
565
566    fn get_row_column(&self, key: &[Value], col: ColumnId) -> Option<Value> {
567        let id = get_entry(key, self.n_keys, &self.hash, |row| {
568            &self.data.get_row(row).unwrap()[0..self.n_keys] == key
569        })?;
570        Some(self.data.get_row(id).unwrap()[col.index()])
571    }
572}
573
574impl SortedWritesTable {
575    /// Create a new [`SortedWritesTable`] with the given number of keys,
576    /// columns, and an optional sort column.
577    ///
578    /// The `merge_fn` is used to evaluate conflicts when more than one row is
579    /// inserted with the same primary key. The old and new proposed values are
580    /// passed as the second and third arguments, respectively, with the
581    /// function filling the final argument with the contents of the new row.
582    /// The return value indicates whether or not the contents of the vector
583    /// should be used.
584    ///
585    /// Merge functions can access the database via [`ExecutionState`].
586    pub fn new(
587        n_keys: usize,
588        n_columns: usize,
589        sort_by: Option<ColumnId>,
590        to_rebuild: Vec<ColumnId>,
591        merge_fn: Box<MergeFn>,
592    ) -> Self {
593        let hash = ShardedHashTable::<TableEntry>::default();
594        let shard_data = hash.shard_data();
595        let rebuild_index = Index::new(to_rebuild.clone(), ColumnIndex::new());
596        SortedWritesTable {
597            generation: Generation::new(0),
598            data: Rows::new(RowBuffer::new(n_columns)),
599            hash,
600            n_keys,
601            n_columns,
602            sort_by,
603            offsets: Default::default(),
604            pending_state: Arc::new(PendingState::new(shard_data)),
605            merge: merge_fn.into(),
606            to_rebuild,
607            rebuild_index,
608            subset_tracker: Default::default(),
609        }
610    }
611
612    /// Flush all pending removals, in parallel.
613    fn parallel_delete(&mut self) -> bool {
614        let shard_data = self.hash.shard_data();
615        let pending_removals = &self.pending_state.pending_removals;
616        let data = &self.data.data;
617        let n_keys = self.n_keys;
618        let stale_delta: usize = parallel::map_mut(self.hash.mut_shards(), |shard_id, shard| {
619            let shard_id = ShardId::from_usize(shard_id);
620            if pending_removals[shard_id].is_empty() {
621                return 0;
622            }
623            let queue = &pending_removals[shard_id];
624            let mut marked_stale = 0;
625            while let Some(buf) = queue.pop() {
626                buf.for_each(|to_remove| {
627                    let (actual_shard, hc) = hash_code(shard_data, to_remove, n_keys);
628                    assert_eq!(actual_shard, shard_id);
629                    if let Ok(entry) = shard.find_entry(hc, |entry| {
630                        entry.hashcode == (hc as _)
631                            && &data.get_row(entry.row)[0..n_keys] == to_remove
632                    }) {
633                        let (ent, _) = entry.remove();
634                        // SAFETY: The safety requirements of
635                        // `set_stale_shared` are that there are no
636                        // concurrent accesses to `row`. No other threads
637                        // can access this row within this method because
638                        // different `shards` partition the space
639                        // (guaranteed by the assertion above), and we
640                        // launch at most one thread per shard.
641                        marked_stale += unsafe { !data.set_stale_shared(ent.row) } as usize;
642                    }
643                });
644            }
645            marked_stale
646        })
647        .into_iter()
648        .sum();
649        // Update the stale count with the total marked stale.
650        self.data.stale_rows += stale_delta;
651        stale_delta > 0
652    }
653    fn serial_delete(&mut self) -> bool {
654        let shard_data = self.hash.shard_data();
655        let mut changed = false;
656        self.hash
657            .mut_shards()
658            .iter_mut()
659            .enumerate()
660            .for_each(|(shard_id, shard)| {
661                let shard_id = ShardId::from_usize(shard_id);
662                let queue = &self.pending_state.pending_removals[shard_id];
663                while let Some(buf) = queue.pop() {
664                    buf.for_each(|to_remove| {
665                        let (actual_shard, hc) = hash_code(shard_data, to_remove, self.n_keys);
666                        assert_eq!(actual_shard, shard_id);
667                        if let Ok(entry) = shard.find_entry(hc, |entry| {
668                            entry.hashcode == (hc as _)
669                                && &self.data.get_row(entry.row).unwrap()[0..self.n_keys]
670                                    == to_remove
671                        }) {
672                            let (ent, _) = entry.remove();
673                            self.data.set_stale(ent.row);
674                            changed = true;
675                        }
676                    })
677                }
678            });
679        changed
680    }
681
682    fn do_delete(&mut self) -> bool {
683        let total = self.pending_state.total_removals.swap(0, Ordering::Relaxed);
684
685        if parallelize_table_op(total) {
686            self.parallel_delete()
687        } else {
688            self.serial_delete()
689        }
690    }
691
692    fn do_insert(&mut self, exec_state: &mut ExecutionState) -> bool {
693        let total = self.pending_state.total_rows.swap(0, Ordering::Relaxed);
694        self.data.data.reserve(total);
695        if parallelize_table_op(total) {
696            if let Some(col) = self.sort_by {
697                self.parallel_insert(
698                    exec_state,
699                    SortChecker {
700                        col,
701                        current: None,
702                        baseline: self.offsets.last().map(|(v, _)| *v),
703                    },
704                )
705            } else {
706                self.parallel_insert(exec_state, ())
707            }
708        } else {
709            self.serial_insert(exec_state)
710        }
711    }
712
713    fn serial_insert(&mut self, exec_state: &mut ExecutionState) -> bool {
714        let mut changed = false;
715        let n_keys = self.n_keys;
716        let mut scratch = with_pool_set(|ps| ps.get::<Vec<Value>>());
717        for (_outer_shard, queue) in self.pending_state.pending_rows.iter() {
718            if let Some(sort_by) = self.sort_by {
719                while let Some(buf) = queue.pop() {
720                    for query in buf.non_stale() {
721                        let key = &query[0..n_keys];
722                        let entry = get_entry_mut(query, n_keys, &mut self.hash, |row| {
723                            let Some(row) = self.data.get_row(row) else {
724                                return false;
725                            };
726                            &row[0..n_keys] == key
727                        });
728
729                        if let Some(row) = entry {
730                            // First case: overwriting an existing value. Apply merge
731                            // function. Insert new row and update hash table if merge
732                            // changes anything.
733                            let cur = self
734                                .data
735                                .get_row(*row)
736                                .expect("table should not point to stale entry");
737                            if (self.merge)(exec_state, cur, query, &mut scratch) {
738                                let sort_val = query[sort_by.index()];
739                                let new = self.data.add_row(&scratch);
740                                if let Some(largest) = self.offsets.last().map(|(v, _)| *v) {
741                                    assert!(
742                                        sort_val >= largest,
743                                        "inserting row that violates sort order ({sort_val:?} vs. {largest:?})"
744                                    );
745                                    if sort_val > largest {
746                                        self.offsets.push((sort_val, new));
747                                    }
748                                } else {
749                                    self.offsets.push((sort_val, new));
750                                }
751                                self.data.set_stale(*row);
752                                *row = new;
753                                changed = true;
754                            }
755                            scratch.clear();
756                        } else {
757                            let sort_val = query[sort_by.index()];
758                            // New value: update invariants.
759                            let new = self.data.add_row(query);
760                            if let Some(largest) = self.offsets.last().map(|(v, _)| *v) {
761                                assert!(
762                                    sort_val >= largest,
763                                    "inserting row that violates sort order {sort_val:?} vs. {largest:?}"
764                                );
765                                if sort_val > largest {
766                                    self.offsets.push((sort_val, new));
767                                }
768                            } else {
769                                self.offsets.push((sort_val, new));
770                            }
771                            let (shard, hc) = hash_code(self.hash.shard_data(), query, self.n_keys);
772                            debug_assert_eq!(shard, _outer_shard);
773                            self.hash.mut_shards()[shard.index()].insert_unique(
774                                hc as _,
775                                TableEntry {
776                                    hashcode: hc as _,
777                                    row: new,
778                                },
779                                TableEntry::hashcode,
780                            );
781                            changed = true;
782                        }
783                    }
784                }
785            } else {
786                // Simplified variant without the sorting constraint.
787                while let Some(buf) = queue.pop() {
788                    for query in buf.non_stale() {
789                        let key = &query[0..n_keys];
790                        let entry = get_entry_mut(query, n_keys, &mut self.hash, |row| {
791                            let Some(row) = self.data.get_row(row) else {
792                                return false;
793                            };
794                            &row[0..n_keys] == key
795                        });
796
797                        if let Some(row) = entry {
798                            let cur = self
799                                .data
800                                .get_row(*row)
801                                .expect("table should not point to stale entry");
802                            if (self.merge)(exec_state, cur, query, &mut scratch) {
803                                let new = self.data.add_row(&scratch);
804                                self.data.set_stale(*row);
805                                *row = new;
806                                changed = true;
807                            }
808                            scratch.clear();
809                        } else {
810                            // New value: update invariants.
811                            let new = self.data.add_row(query);
812                            let (shard, hc) = hash_code(self.hash.shard_data(), query, self.n_keys);
813                            debug_assert_eq!(shard, _outer_shard);
814                            self.hash.mut_shards()[shard.index()].insert_unique(
815                                hc as _,
816                                TableEntry {
817                                    hashcode: hc as _,
818                                    row: new,
819                                },
820                                TableEntry::hashcode,
821                            );
822                            changed = true;
823                        }
824                    }
825                }
826            };
827        }
828        changed
829    }
830
831    fn parallel_insert<C: OrderingChecker>(
832        &mut self,
833        exec_state: &ExecutionState,
834        checker: C,
835    ) -> bool {
836        const BATCH_SIZE: usize = 1 << 18;
837        // Parallel insert uses one giant parallel foreach. We have updates
838        // pre-sharded, and one logical thread can process updates for each
839        // shard independently. Updates happen in three phases, which comments
840        // describe below.
841        let shard_data = self.hash.shard_data();
842        let n_keys = self.n_keys;
843        let n_cols = self.n_columns;
844        let next_offset = RowId::from_usize(self.data.data.len());
845        let row_writer = self.data.data.parallel_writer();
846        let pending_rows = &self.pending_state.pending_rows;
847        let merge = self.merge.clone();
848        let pending_adds = parallel::map_mut(self.hash.mut_shards(), |shard_id, shard| {
849            let shard_id = ShardId::from_usize(shard_id);
850            let mut checker = checker.clone();
851            let mut exec_state = exec_state.clone();
852            let mut scratch = with_pool_set(|ps| ps.get::<Vec<Value>>());
853            let queue = &pending_rows[shard_id];
854            let mut marked_stale = 0usize;
855            let mut staged = StagedOutputs::new(n_keys, n_cols, BATCH_SIZE);
856            let mut changed = false;
857            // The core flush loop: We call once `staged` reaches `BATCH_SIZE` or
858            // when we're done.
859            macro_rules! flush_staged_outputs {
860                    () => {{
861                        // Phase 2: Write the staged rows to the row writer. This only
862                        // works due to the `ParallelRowBufWriter` machinery.
863                        let (start_row, stale) = staged.write_output(&row_writer);
864                        marked_stale += stale;
865                        // Phase 3: With the values buffered in the row buffer, we can
866                        // write them back to the shard, pointed to the correct rows.
867
868                        // In the serial implementation, we do phases 2 and 3 inline with
869                        // processing the incoming mutation, but separating them out
870                        // this way allows us to do a single write to the shared row
871                        // buffer, rather than one per row, which would cause
872                        // contention.
873                        let mut cur_row = start_row;
874                        let read_handle = row_writer.read_handle();
875                        for row in staged.rows() {
876                            if row.first().map(Value::is_stale).unwrap_or(false) {
877                                cur_row = cur_row.inc();
878                                continue;
879                            }
880                            use hashbrown::hash_table::Entry;
881                            checker.check_local(row);
882                            changed = true;
883                            let key = &row[0..n_keys];
884                            let (_actual_shard, hc) = hash_code(shard_data, row, n_keys);
885                            #[cfg(any(debug_assertions, test))]
886                            {
887                                unsafe {
888                                    // read the value we wrote at this row and
889                                    // check that it matches.
890                                    assert_eq!(read_handle.get_row_unchecked(cur_row), row);
891                                }
892                            }
893                            debug_assert_eq!(_actual_shard, shard_id);
894                            match shard.entry(
895                                hc,
896                                // SAFETY: `ent` must point to a valid row
897                                |ent| unsafe {
898                                    ent.hashcode == hc as HashCode
899                                        && &read_handle.get_row_unchecked(ent.row)[0..n_keys] == key
900                                },
901                                TableEntry::hashcode,
902                            ) {
903                                Entry::Occupied(mut occ) => {
904                                    // SAFETY: `occ` must point to a valid row: we only insert valid rows
905                                    // into the map.
906                                    let cur = unsafe { read_handle.get_row_unchecked(occ.get().row) };
907
908                                    // SAFETY: The safety requirements of
909                                    // `set_stale_shared` are that there are no
910                                    // concurrent accesses to `row`. We have
911                                    // exclusive access to any row whose hash matches this
912                                    // shard.
913                                    if (merge)(&mut exec_state, cur, row, &mut scratch) {
914                                        unsafe {
915                                            let _was_stale = read_handle.set_stale_shared(occ.get().row);
916                                            debug_assert!(!_was_stale);
917                                        }
918                                        occ.get_mut().row = cur_row;
919                                        changed = true;
920                                    } else {
921                                        // Mark the new row as stale: we didn't end up needing it.
922                                        unsafe {
923                                            let _was_stale = read_handle.set_stale_shared(cur_row);
924                                            debug_assert!(!_was_stale);
925                                        }
926                                    }
927                                    marked_stale += 1;
928                                    scratch.clear();
929                                }
930                                Entry::Vacant(v) => {
931                                    changed = true;
932                                    v.insert(TableEntry {
933                                        hashcode: hc as HashCode,
934                                        row: cur_row,
935                                    });
936                                }
937                            }
938
939                            cur_row = cur_row.inc();
940                        }
941                        staged.clear();
942                    }};
943                }
944            // Phase 1: process all incoming updates:
945            // * Add new values to `staged`
946            // * Removing entries in `shard` and mark them as stale in
947            // `data` if they will be overwritten.
948            while let Some(buf) = queue.pop() {
949                // We create a read_handle once per batch to avoid blocking
950                // too many threads if someone needs to resize the row
951                // writer.
952                for row in buf.non_stale() {
953                    staged.insert(row, |cur, new, out| (merge)(&mut exec_state, cur, new, out));
954                    if staged.len() >= BATCH_SIZE {
955                        flush_staged_outputs!();
956                    }
957                }
958            }
959            flush_staged_outputs!();
960            (checker, marked_stale, changed)
961        });
962        self.data.data = row_writer.finish();
963        // Now we just need to reset our invariants.
964
965        // Confirm none of the writes violated sort order and update the
966        // `offsets` vector.
967        let checker = C::check_global(pending_adds.iter().map(|(checker, _, _)| checker));
968        checker.update_offsets(next_offset, &mut self.offsets);
969
970        // Update the staleness counters.
971        self.data.stale_rows += pending_adds
972            .iter()
973            .map(|(_, stale, _)| *stale)
974            .sum::<usize>();
975
976        // Register any changes.
977        pending_adds.iter().any(|(_, _, changed)| *changed)
978    }
979
980    fn binary_search_sort_val(&self, val: Value) -> Result<(RowId, RowId), RowId> {
981        debug_assert!(
982            self.offsets.windows(2).all(|x| x[0].1 < x[1].1),
983            "{:?}",
984            self.offsets
985        );
986
987        debug_assert!(
988            self.offsets.windows(2).all(|x| x[0].0 < x[1].0),
989            "{:?}",
990            self.offsets
991        );
992        match self.offsets.binary_search_by_key(&val, |(v, _)| *v) {
993            Ok(got) => Ok((
994                self.offsets[got].1,
995                self.offsets
996                    .get(got + 1)
997                    .map(|(_, r)| *r)
998                    .unwrap_or(self.data.next_row()),
999            )),
1000            Err(next) => Err(self
1001                .offsets
1002                .get(next)
1003                .map(|(_, id)| *id)
1004                .unwrap_or(self.data.next_row())),
1005        }
1006    }
1007    fn eval(&self, cs: &[Constraint], row: RowId) -> bool {
1008        self.get_if(cs, row).is_some()
1009    }
1010
1011    fn eval_constraints(cs: &[Constraint], row: &[Value]) -> bool {
1012        cs.iter().all(|constraint| match constraint {
1013            Constraint::Eq { l_col, r_col } => row[l_col.index()] == row[r_col.index()],
1014            Constraint::EqConst { col, val } => row[col.index()] == *val,
1015            Constraint::LtConst { col, val } => row[col.index()] < *val,
1016            Constraint::GtConst { col, val } => row[col.index()] > *val,
1017            Constraint::LeConst { col, val } => row[col.index()] <= *val,
1018            Constraint::GeConst { col, val } => row[col.index()] >= *val,
1019        })
1020    }
1021
1022    unsafe fn get_if_unchecked(&self, cs: &[Constraint], row: RowId) -> Option<&[Value]> {
1023        let row = unsafe { self.data.data.get_row_unchecked(row) };
1024        if Self::eval_constraints(cs, row) {
1025            Some(row)
1026        } else {
1027            None
1028        }
1029    }
1030
1031    fn get_if(&self, cs: &[Constraint], row: RowId) -> Option<&[Value]> {
1032        let row = self.data.get_row(row)?;
1033        if Self::eval_constraints(cs, row) {
1034            Some(row)
1035        } else {
1036            None
1037        }
1038    }
1039
1040    fn maybe_rehash(&mut self) {
1041        if self.data.stale_rows <= cmp::max(16, self.data.data.len() / 2) {
1042            return;
1043        }
1044
1045        if parallelize_table_op(self.data.data.len()) {
1046            self.parallel_rehash();
1047        } else {
1048            self.rehash();
1049        }
1050    }
1051    fn parallel_rehash(&mut self) {
1052        // Parallel rehashes go "hash-first" rather than "rows-first".
1053        //
1054        // We iterate over each shard and then write out new contents to a fresh row, in parallel.
1055        let Some(sort_by) = self.sort_by else {
1056            // Just do a serial rehash for now. We currently do not have a use-case for parallel
1057            // compaction of unsorted tables.
1058            //
1059            // Implementing parallel compaction for an unsorted table is much easier: each shard
1060            // can write to a contiguous chunk of the `scratch` buffer, with the offsets being
1061            // pre-chunked based on the size of each shard.
1062            self.rehash();
1063            return;
1064        };
1065        self.generation = self.generation.inc();
1066        assert!(!self.offsets.is_empty());
1067        struct TimestampStats {
1068            value: Value,
1069            count: usize,
1070            histogram: Pooled<DenseIdMap<ShardId, usize>>,
1071        }
1072        impl Default for TimestampStats {
1073            fn default() -> TimestampStats {
1074                TimestampStats {
1075                    value: Value::stale(),
1076                    count: 0,
1077                    histogram: with_pool_set(|ps| ps.get()),
1078                }
1079            }
1080        }
1081        let mut results = Vec::<TimestampStats>::with_capacity(self.offsets.len());
1082        let offset_windows = self
1083            .offsets
1084            .windows(2)
1085            .map(|xs| {
1086                let [(start_val, start_row), (_, end_row)] = xs else {
1087                    unreachable!()
1088                };
1089                (*start_val, *start_row, *end_row)
1090            })
1091            .collect::<Vec<_>>();
1092        // Use a macro rather than a lambda to avoid borrow issues.
1093        macro_rules! compute_hist {
1094            ($start_val: expr, $start_row: expr, $end_row: expr) => {{
1095                let mut histogram: Pooled<DenseIdMap<ShardId, usize>> =
1096                    with_pool_set(|ps| ps.get());
1097                let mut cur_row = $start_row;
1098                let mut count = 0;
1099                while cur_row < $end_row {
1100                    if let Some(row) = self.data.get_row(cur_row) {
1101                        count += 1;
1102                        let (shard, _) = hash_code(self.hash.shard_data(), row, self.n_keys);
1103                        *histogram.get_or_default(shard) += 1;
1104                    }
1105                    cur_row = cur_row.inc();
1106                }
1107                TimestampStats {
1108                    value: $start_val,
1109                    count,
1110                    histogram,
1111                }
1112            }};
1113        }
1114        results.extend(parallel::map(
1115            &offset_windows,
1116            |_, (start_val, start_row, end_row)| compute_hist!(*start_val, *start_row, *end_row),
1117        ));
1118        // And here we handle the final one.
1119        let (start_val, start_row) = self.offsets.last().unwrap();
1120        let end_row = self.data.next_row();
1121        let last = compute_hist!(*start_val, *start_row, end_row);
1122        results.push(last);
1123        // Now we need to compute cumulative statistics on the row layouts here.
1124        // We do this serially a we currently don't have a ton of use for cases with thousands
1125        // of timestamps or more. There are well-known parallel algorithms for computing these
1126        // cumulative statistics in parallel, but they are not currently a good fit here.
1127        let mut prev_count = 0;
1128        self.offsets.clear();
1129        for stats in results.iter_mut() {
1130            if stats.count == 0 {
1131                continue;
1132            }
1133            self.offsets
1134                .push((stats.value, RowId::from_usize(prev_count)));
1135            let mut inner = prev_count;
1136            for (_, count) in stats.histogram.iter_mut() {
1137                // Each entry in the histogram now points to the start row for that shard's
1138                // rows for a given timestamp.
1139                let tmp = *count;
1140                *count = inner;
1141                inner += tmp;
1142            }
1143            prev_count += stats.count;
1144            debug_assert_eq!(inner, prev_count)
1145        }
1146
1147        // Now the part with some unsafe code.
1148        // We will iterate over each shard and use the statistics in `results` to guide where
1149        // each row will go.
1150        //
1151        // This involves doing unsynchronized writes to the table (ptr::copy_nonoverlapping)
1152        // followed by a set_len. The safety of these operations relies on the fact that:
1153        // * No one grabs a reference to the interior of `scratch` until these operations have
1154        //   finished.
1155        // * `scratch` does not overlap `data`.
1156        // * The sharding function completely partitions the set of objects in the table: one
1157        //   shard's writes will never stomp on those of another.
1158
1159        self.data.scratch.clear();
1160        self.data.scratch.reserve(prev_count);
1161        let scratch_ptr = self.data.scratch.raw_rows() as usize;
1162        let data = &self.data.data;
1163        let n_columns = self.n_columns;
1164        parallel::for_each_mut(self.hash.mut_shards(), |shard_id, shard| {
1165            let shard_id = ShardId::from_usize(shard_id);
1166            let scratch_ptr = scratch_ptr as *const Value;
1167            let mut progress = HashMap::<Value /* timestamp */, RowId /* next row */>::default();
1168            progress.reserve(results.len());
1169            for stats in &results {
1170                let Some(start) = stats.histogram.get(shard_id) else {
1171                    continue;
1172                };
1173                progress.insert(stats.value, RowId::from_usize(*start));
1174            }
1175            for TableEntry { row: row_id, .. } in shard.iter_mut() {
1176                let row = data.get_row(*row_id);
1177                debug_assert!(!row[0].is_stale(), "shard should not map to a stale value");
1178                let val = row[sort_by.index()];
1179                let next = progress[&val];
1180                // SAFETY: see above longer comment.
1181                unsafe {
1182                    std::ptr::copy_nonoverlapping(
1183                        row.as_ptr(),
1184                        scratch_ptr.add(next.index() * n_columns) as *mut Value,
1185                        n_columns,
1186                    )
1187                }
1188                *row_id = next;
1189                progress.insert(val, next.inc());
1190            }
1191        });
1192        // SAFETY: see above longer comment.
1193        unsafe { self.data.scratch.set_len(prev_count) };
1194        mem::swap(&mut self.data.data, &mut self.data.scratch);
1195        self.data.stale_rows = 0;
1196    }
1197    fn rehash_impl(
1198        sort_by: Option<ColumnId>,
1199        n_keys: usize,
1200        rows: &mut Rows,
1201        offsets: &mut Vec<(Value, RowId)>,
1202        hash: &mut ShardedHashTable<TableEntry>,
1203    ) {
1204        if let Some(sort_by) = sort_by {
1205            offsets.clear();
1206            rows.remove_stale(|row, old, new| {
1207                let stale_entry = get_entry_mut(row, n_keys, hash, |x| x == old)
1208                    .expect("non-stale entry not mapped in hash");
1209                *stale_entry = new;
1210                let sort_col = row[sort_by.index()];
1211                if let Some((max, _)) = offsets.last() {
1212                    if sort_col > *max {
1213                        offsets.push((sort_col, new));
1214                    }
1215                } else {
1216                    offsets.push((sort_col, new));
1217                }
1218            })
1219        } else {
1220            rows.remove_stale(|row, old, new| {
1221                let stale_entry = get_entry_mut(row, n_keys, hash, |x| x == old)
1222                    .expect("non-stale entry not mapped in hash");
1223                *stale_entry = new;
1224            })
1225        }
1226    }
1227
1228    fn rehash(&mut self) {
1229        self.generation = self.generation.inc();
1230        Self::rehash_impl(
1231            self.sort_by,
1232            self.n_keys,
1233            &mut self.data,
1234            &mut self.offsets,
1235            &mut self.hash,
1236        )
1237    }
1238}
1239
1240fn get_entry(
1241    row: &[Value],
1242    n_keys: usize,
1243    table: &ShardedHashTable<TableEntry>,
1244    test: impl Fn(RowId) -> bool,
1245) -> Option<RowId> {
1246    let (shard, hash) = hash_code(table.shard_data(), row, n_keys);
1247    table
1248        .get_shard(shard)
1249        .find(hash, |ent| {
1250            ent.hashcode == hash as HashCode && test(ent.row)
1251        })
1252        .map(|ent| ent.row)
1253}
1254
1255fn get_entry_mut<'a>(
1256    row: &[Value],
1257    n_keys: usize,
1258    table: &'a mut ShardedHashTable<TableEntry>,
1259    test: impl Fn(RowId) -> bool,
1260) -> Option<&'a mut RowId> {
1261    let (shard, hash) = hash_code(table.shard_data(), row, n_keys);
1262    table.mut_shards()[shard.index()]
1263        .find_mut(hash, |ent| {
1264            ent.hashcode == hash as HashCode && test(ent.row)
1265        })
1266        .map(|ent| &mut ent.row)
1267}
1268
1269fn hash_code(shard_data: ShardData, row: &[Value], n_keys: usize) -> (ShardId, u64) {
1270    let mut hasher = FxHasher::default();
1271    for val in &row[0..n_keys] {
1272        hasher.write_usize(val.index());
1273    }
1274    let full_code = hasher.finish();
1275    // We keep this cast here to allow for experimenting with HashCode=u32.
1276    #[allow(clippy::unnecessary_cast)]
1277    (shard_data.shard_id(full_code), full_code as HashCode as u64)
1278}
1279
1280/// A simple struct for packaging up pending mutations to a `SortedWritesTable`.
1281struct PendingState {
1282    pending_rows: DenseIdMap<ShardId, SegQueue<RowBuffer>>,
1283    pending_removals: DenseIdMap<ShardId, SegQueue<ArbitraryRowBuffer>>,
1284    total_removals: AtomicUsize,
1285    total_rows: AtomicUsize,
1286}
1287
1288impl PendingState {
1289    fn new(shard_data: ShardData) -> PendingState {
1290        let n_shards = shard_data.n_shards();
1291        let mut pending_rows = DenseIdMap::with_capacity(n_shards);
1292        let mut pending_removals = DenseIdMap::with_capacity(n_shards);
1293        for i in 0..n_shards {
1294            pending_rows.insert(ShardId::from_usize(i), SegQueue::default());
1295            pending_removals.insert(ShardId::from_usize(i), SegQueue::default());
1296        }
1297
1298        PendingState {
1299            pending_rows,
1300            pending_removals,
1301            total_removals: AtomicUsize::new(0),
1302            total_rows: AtomicUsize::new(0),
1303        }
1304    }
1305    fn clear(&self) {
1306        for (_, queue) in self.pending_rows.iter() {
1307            while queue.pop().is_some() {}
1308        }
1309
1310        for (_, queue) in self.pending_removals.iter() {
1311            while queue.pop().is_some() {}
1312        }
1313    }
1314
1315    /// This is only really used in debugging, but it's annoying enough to write
1316    /// that it may help to have around.
1317    ///
1318    /// We also, however, use it in the clone impl (which should only be called when pending state
1319    /// is empty).
1320    fn deep_copy(&self) -> PendingState {
1321        let mut pending_rows = DenseIdMap::new();
1322        let mut pending_removals = DenseIdMap::new();
1323        fn drain_queue<T>(queue: &SegQueue<T>) -> Vec<T> {
1324            let mut res = Vec::new();
1325            while let Some(x) = queue.pop() {
1326                res.push(x);
1327            }
1328            res
1329        }
1330        for (shard, queue) in self.pending_rows.iter() {
1331            let contents = drain_queue(queue);
1332            let new_queue = SegQueue::default();
1333            for x in contents {
1334                new_queue.push(x.clone());
1335                queue.push(x);
1336            }
1337            pending_rows.insert(shard, new_queue);
1338        }
1339
1340        for (shard, queue) in self.pending_removals.iter() {
1341            let contents = drain_queue(queue);
1342            let new_queue = SegQueue::default();
1343            for x in contents {
1344                new_queue.push(x.clone());
1345                queue.push(x);
1346            }
1347            pending_removals.insert(shard, new_queue);
1348        }
1349
1350        PendingState {
1351            pending_rows,
1352            pending_removals,
1353            total_removals: AtomicUsize::new(self.total_removals.load(Ordering::Acquire)),
1354            total_rows: AtomicUsize::new(self.total_rows.load(Ordering::Acquire)),
1355        }
1356    }
1357}
1358
1359/// A trait that encapsulates the logic of potentially checking that written
1360/// columns appear in sorted order.
1361///
1362/// For rows that are sorted by a column, an OrderingChecker asserts that all
1363/// new rows have the same value in that column, and that the column is greater
1364/// than or equal to the column value coming in. For rows not sorted, these
1365/// checks become no-ops.
1366trait OrderingChecker: Clone + Send + Sync {
1367    /// Check any invariants locally, updating the state of the checker when
1368    /// doing so.
1369    fn check_local(&mut self, row: &[Value]);
1370    /// Combine the states of multiple checkers, returning a new checker with
1371    /// all information assimilated. This is the checker that is suitable for
1372    /// calling `update_offsets` with.
1373    fn check_global<'a>(checkers: impl Iterator<Item = &'a Self>) -> Self
1374    where
1375        Self: 'a;
1376    /// Update the sorted offset vector with the current state of the checker.
1377    fn update_offsets(&self, start: RowId, offsets: &mut Vec<(Value, RowId)>);
1378}
1379
1380impl OrderingChecker for () {
1381    fn check_local(&mut self, _: &[Value]) {}
1382    fn check_global<'a>(_: impl Iterator<Item = &'a ()>) {}
1383    fn update_offsets(&self, _: RowId, _: &mut Vec<(Value, RowId)>) {}
1384}
1385
1386#[derive(Copy, Clone)]
1387struct SortChecker {
1388    col: ColumnId,
1389    baseline: Option<Value>,
1390    current: Option<Value>,
1391}
1392
1393impl OrderingChecker for SortChecker {
1394    fn check_local(&mut self, row: &[Value]) {
1395        let val = row[self.col.index()];
1396        if let Some(cur) = self.current {
1397            assert_eq!(
1398                cur, val,
1399                "concurrently inserting rows with different sort keys"
1400            );
1401        } else {
1402            self.current = Some(val);
1403            if let Some(baseline) = self.baseline {
1404                assert!(val >= baseline, "inserted row violates sort order");
1405            }
1406        }
1407    }
1408
1409    fn check_global<'a>(mut checkers: impl Iterator<Item = &'a Self>) -> Self {
1410        let Some(start) = checkers.next() else {
1411            return SortChecker {
1412                col: ColumnId::new(!0),
1413                baseline: None,
1414                current: None,
1415            };
1416        };
1417        let mut expected = start.current;
1418        for checker in checkers {
1419            assert_eq!(checker.baseline, start.baseline);
1420            match (&mut expected, checker.current) {
1421                (None, None) => {}
1422                (cur @ None, Some(x)) => {
1423                    *cur = Some(x);
1424                }
1425                (Some(_), None) => {}
1426                (Some(x), Some(y)) => {
1427                    assert_eq!(
1428                        *x, y,
1429                        "concurrently inserting rows with different sort keys"
1430                    );
1431                }
1432            }
1433        }
1434        SortChecker {
1435            col: start.col,
1436            baseline: start.baseline,
1437            current: expected,
1438        }
1439    }
1440
1441    fn update_offsets(&self, start: RowId, offsets: &mut Vec<(Value, RowId)>) {
1442        if let Some(cur) = self.current {
1443            if let Some((max, _)) = offsets.last() {
1444                if cur > *max {
1445                    offsets.push((cur, start));
1446                }
1447            } else {
1448                offsets.push((cur, start));
1449            }
1450        }
1451    }
1452}
1453
1454/// A type similar to a SortedWritesTable used to buffer outputs. The main thing
1455/// that StagedOutputs handles is running the merge function for a table on
1456/// multiple updates to the same key that show up in the same round of
1457/// insertions.
1458struct StagedOutputs {
1459    shard_data: ShardData,
1460    n_keys: usize,
1461    hash: Pooled<HashTable<TableEntry>>,
1462    rows: RowBuffer,
1463    n_stale: usize,
1464    scratch: Pooled<Vec<Value>>,
1465}
1466
1467impl StagedOutputs {
1468    fn rows(&self) -> impl Iterator<Item = &[Value]> {
1469        self.rows.iter()
1470    }
1471    fn new(n_keys: usize, n_cols: usize, capacity: usize) -> Self {
1472        let mut res = with_pool_set(|ps| StagedOutputs {
1473            shard_data: ShardData::new(1),
1474            n_keys,
1475            n_stale: 0,
1476            hash: ps.get(),
1477            rows: RowBuffer::new(n_cols),
1478            scratch: ps.get(),
1479        });
1480        res.hash.reserve(capacity, TableEntry::hashcode);
1481        res.rows.reserve(capacity);
1482        res
1483    }
1484    fn clear(&mut self) {
1485        self.hash.clear();
1486        self.rows.clear();
1487        self.n_stale = 0;
1488    }
1489    fn len(&self) -> usize {
1490        self.rows.len() - self.n_stale
1491    }
1492
1493    fn insert(
1494        &mut self,
1495        row: &[Value],
1496        mut merge_fn: impl FnMut(&[Value], &[Value], &mut Vec<Value>) -> bool,
1497    ) {
1498        if row[0].is_stale() {
1499            return;
1500        }
1501        use hashbrown::hash_table::Entry;
1502        let (_, hc) = hash_code(self.shard_data, row, self.n_keys);
1503        let entry = self.hash.entry(
1504            hc,
1505            |te| {
1506                te.hashcode() == hc
1507                    && self.rows.get_row(te.row)[0..self.n_keys] == row[0..self.n_keys]
1508            },
1509            TableEntry::hashcode,
1510        );
1511        match entry {
1512            Entry::Occupied(mut occupied_entry) => {
1513                let cur = self.rows.get_row(occupied_entry.get().row);
1514                if merge_fn(cur, row, &mut self.scratch) {
1515                    let new = self.rows.add_row(&self.scratch);
1516                    self.rows.set_stale(occupied_entry.get().row);
1517                    self.n_stale += 1;
1518                    occupied_entry.get_mut().row = new;
1519                }
1520                self.scratch.clear();
1521            }
1522            Entry::Vacant(vacant_entry) => {
1523                let next = self.rows.add_row(row);
1524                vacant_entry.insert(TableEntry {
1525                    hashcode: hc as _,
1526                    row: next,
1527                });
1528            }
1529        }
1530    }
1531
1532    /// Write the contents of the staged outputs to the given writer, returning the initial RowId
1533    /// of the new output. Returns the number of stale values in the buffer that was appended.
1534    fn write_output(&self, output: &ParallelRowBufWriter) -> (RowId, usize) {
1535        (output.append_contents(&self.rows), self.n_stale)
1536    }
1537}