From fba3c266b7ce4db0bf86905ef6b6feb90017eb4e Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Sat, 25 Jul 2026 00:02:28 +0200 Subject: [PATCH 01/12] Store predicted rows contiguously --- core-relations/src/action/mod.rs | 227 +++++++++++++++++------------ core-relations/src/action/tests.rs | 87 ++++++++++- 2 files changed, 219 insertions(+), 95 deletions(-) diff --git a/core-relations/src/action/mod.rs b/core-relations/src/action/mod.rs index 6395ef61b..5ace8e2ee 100644 --- a/core-relations/src/action/mod.rs +++ b/core-relations/src/action/mod.rs @@ -3,7 +3,7 @@ //! This allows us to execute the "right-hand-side" of a rule. The //! implementation here is optimized to execute on a batch of rows at a time. use std::{ - ops::Deref, + hash::Hasher, sync::{ Arc, atomic::{AtomicBool, Ordering}, @@ -11,12 +11,12 @@ use std::{ }; use crate::{ - common::HashMap, free_join::{invoke_batch, invoke_batch_assign}, numeric_id::{DenseIdMap, NumericId}, }; use egglog_concurrency::NotificationList; -use smallvec::SmallVec; +use hashbrown::HashTable; +use rustc_hash::FxHasher; use crate::{ BaseValues, ContainerValues, ExternalFunctionId, WrappedTable, @@ -289,36 +289,89 @@ pub(crate) struct ExtractedBinding { pub(crate) vals: Pooled>, } +#[derive(Clone, Copy)] +struct PredictedEntry { + hash: u64, + index: usize, + table: TableId, + key_arity: u32, +} + #[derive(Default)] pub(crate) struct PredictedVals { - #[allow(clippy::type_complexity)] - data: HashMap<(TableId, SmallVec<[Value; 3]>), Pooled>>, + index: HashTable, + values: Vec, } impl Clear for PredictedVals { fn reuse(&self) -> bool { - self.data.capacity() > 0 + self.index.capacity() > 0 || self.values.capacity() > 0 } fn clear(&mut self) { - self.data.clear() + self.index.clear(); + self.values.clear(); } fn bytes(&self) -> usize { - self.data.capacity() - * (std::mem::size_of::<(TableId, SmallVec<[Value; 3]>)>() - + std::mem::size_of::>>()) + self.index.capacity() * std::mem::size_of::() + + self.values.capacity() * std::mem::size_of::() } } impl PredictedVals { - pub(crate) fn get_val( + fn hash(table: TableId, key: &[Value]) -> u64 { + let mut hasher = FxHasher::default(); + hasher.write_u32(table.rep()); + hasher.write_usize(key.len()); + for value in key { + hasher.write_u32(value.rep()); + } + hasher.finish() + } + + fn row(&self, entry: PredictedEntry, row_arity: usize) -> &[Value] { + &self.values[entry.index..entry.index + row_arity] + } + + fn matches(&self, entry: &PredictedEntry, hash: u64, table: TableId, key: &[Value]) -> bool { + if entry.hash != hash || entry.table != table || entry.key_arity as usize != key.len() { + return false; + } + self.values[entry.index..entry.index + key.len()] == *key + } + + pub(crate) fn get_or_insert_with( &mut self, table: TableId, key: &[Value], - default: impl FnOnce() -> Pooled>, - ) -> impl Deref>> + '_ { - self.data - .entry((table, SmallVec::from_slice(key))) - .or_insert_with(default) + row_arity: usize, + default: impl FnOnce(&mut Vec, usize), + ) -> (&[Value], bool) { + let hash = Self::hash(table, key); + if let Some(entry) = self + .index + .find(hash, |entry| self.matches(entry, hash, table, key)) + .copied() + { + return (self.row(entry, row_arity), false); + } + + let entry = PredictedEntry { + hash, + index: self.values.len(), + table, + key_arity: u32::try_from(key.len()).expect("predicted key arity must fit in u32"), + }; + self.values.reserve(row_arity); + let row_start = self.values.len(); + self.values.extend_from_slice(key); + default(&mut self.values, row_start); + assert_eq!( + self.values.len() - row_start, + row_arity, + "predicted row builder produced the wrong arity" + ); + self.index.insert_unique(hash, entry, |entry| entry.hash); + (self.row(entry, row_arity), true) } } @@ -550,44 +603,28 @@ impl<'a> ExecutionState<'a> { if let Some(row) = self.db.table_info[table].table.get_row(key) { return row.vals; } - Pooled::cloned( + let row_arity = key.len() + vals.len(); + let counters = self.db.counters; + let (row, inserted) = self.predicted - .get_val(table, key, || { - Self::construct_new_row( - &self.db, - &mut self.buffers, - &mut self.changed, - table, - key, - vals, - ) - }) - .deref(), - ) - } - - fn construct_new_row( - db: &DbView, - buffers: &mut MutationBuffers, - changed: &mut bool, - table: TableId, - key: &[Value], - vals: impl ExactSizeIterator, - ) -> Pooled> { + .get_or_insert_with(table, key, row_arity, |values, _| { + for val in vals { + values.push(match val { + MergeVal::Counter(ctr) => Value::from_usize(counters.inc(ctr)), + MergeVal::Constant(c) => c, + }); + } + }); + if inserted { + self.buffers + .lazy_init(table, || self.db.table_info[table].table.new_buffer()); + self.buffers.stage_insert(table, row); + self.changed = true; + } with_pool_set(|ps| { - let mut new = ps.get::>(); - new.reserve(key.len() + vals.len()); - new.extend_from_slice(key); - for val in vals { - new.push(match val { - MergeVal::Counter(ctr) => Value::from_usize(db.counters.inc(ctr)), - MergeVal::Constant(c) => c, - }) - } - buffers.lazy_init(table, || db.table_info[table].table.new_buffer()); - buffers.stage_insert(table, &new); - *changed = true; - new + let mut result = ps.get::>(); + result.extend_from_slice(row); + result }) } @@ -603,16 +640,25 @@ impl<'a> ExecutionState<'a> { if let Some(val) = self.db.table_info[table].table.get_row_column(key, col) { return val; } - self.predicted.get_val(table, key, || { - Self::construct_new_row( - &self.db, - &mut self.buffers, - &mut self.changed, - table, - key, - vals, - ) - })[col.index()] + let row_arity = key.len() + vals.len(); + let counters = self.db.counters; + let (row, inserted) = + self.predicted + .get_or_insert_with(table, key, row_arity, |values, _| { + for val in vals { + values.push(match val { + MergeVal::Counter(ctr) => Value::from_usize(counters.inc(ctr)), + MergeVal::Constant(c) => c, + }); + } + }); + if inserted { + self.buffers + .lazy_init(table, || self.db.table_info[table].table.new_buffer()); + self.buffers.stage_insert(table, row); + self.changed = true; + } + row[col.index()] } /// Trigger early stopping by setting the stop_match flag. @@ -683,7 +729,6 @@ impl ExecutionState<'_> { dst_col, dst_var, } => { - let pool = with_pool_set(|ps| ps.get_pool::>().clone()); self.buffers.lazy_init(*table_id, || { self.db.table_info[*table_id].table.new_buffer() }); @@ -707,42 +752,36 @@ impl ExecutionState<'_> { // // We avoid doing this more than once by using the // `predicted` map. - let prediction_key = ( - *table_id, - SmallVec::<[Value; 3]>::from_slice(key.as_slice()), - ); let buffers = &mut self.buffers; // Bind some mutable references because the closure passed // to or_insert_with is `move`. let ctrs = &self.db.counters; let bindings = &bindings; - let pool = pool.clone(); - let row = - self.predicted - .data - .entry(prediction_key) - .or_insert_with(move || { - let mut row = pool.get(); - row.extend_from_slice(key.as_slice()); - // Extend the key with the default values. - row.reserve(default.len()); - for val in default { - let val = match val { - WriteVal::QueryEntry(QueryEntry::Const(c)) => *c, - WriteVal::QueryEntry(QueryEntry::Var(v)) => { - bindings[*v][offset] - } - WriteVal::IncCounter(ctr) => { - Value::from_usize(ctrs.inc(*ctr)) - } - WriteVal::CurrentVal(ix) => row[*ix], - }; - row.push(val) - } - // Insert it into the table. - buffers.stage_insert(*table_id, &row); - row - }); + let row_arity = key.as_slice().len() + default.len(); + let (row, inserted) = self.predicted.get_or_insert_with( + *table_id, + key.as_slice(), + row_arity, + |values, row_start| { + // Extend the key with the default values. + for val in default { + let val = match val { + WriteVal::QueryEntry(QueryEntry::Const(c)) => *c, + WriteVal::QueryEntry(QueryEntry::Var(v)) => { + bindings[*v][offset] + } + WriteVal::IncCounter(ctr) => { + Value::from_usize(ctrs.inc(*ctr)) + } + WriteVal::CurrentVal(ix) => values[row_start + *ix], + }; + values.push(val) + } + }, + ); + if inserted { + buffers.stage_insert(*table_id, row); + } row[dst_col.index()] }); }); diff --git a/core-relations/src/action/tests.rs b/core-relations/src/action/tests.rs index 4f1c0484b..723b7ad6e 100644 --- a/core-relations/src/action/tests.rs +++ b/core-relations/src/action/tests.rs @@ -1,9 +1,94 @@ use crate::{ + TableId, Value, action::mask::{IterResult, ValueSource}, + numeric_id::NumericId, + pool::Clear, pool::{PoolSet, with_pool_set}, }; -use super::mask::{Mask, MaskIter}; +use super::{ + PredictedEntry, PredictedVals, + mask::{Mask, MaskIter}, +}; + +#[test] +fn predicted_vals_store_rows_contiguously() { + let mut predicted = PredictedVals::default(); + let table = TableId::from_usize(3); + let other_table = TableId::from_usize(4); + let key = [ + Value::from_usize(10), + Value::from_usize(11), + Value::from_usize(12), + Value::from_usize(13), + ]; + let row = [key[0], key[1], key[2], key[3], Value::from_usize(99)]; + let mut default_calls = 0; + + let (got, inserted) = predicted.get_or_insert_with(table, &key, row.len(), |values, _| { + default_calls += 1; + values.extend_from_slice(&row[key.len()..]); + }); + assert_eq!(got, row); + assert!(inserted); + + let (got, inserted) = predicted.get_or_insert_with(table, &key, row.len(), |_, _| { + default_calls += 1; + }); + assert_eq!(got, row); + assert!(!inserted); + assert_eq!(default_calls, 1); + + let other_row = [key[0], key[1], key[2], key[3], Value::from_usize(100)]; + let (got, inserted) = + predicted.get_or_insert_with(other_table, &key, other_row.len(), |values, _| { + default_calls += 1; + values.extend_from_slice(&other_row[key.len()..]); + }); + assert_eq!(got, other_row); + assert!(inserted); + assert_eq!(predicted.index.len(), 2); + assert_eq!(predicted.values.len(), 2 * row.len()); + + predicted.clear(); + assert!(predicted.index.is_empty()); + assert!(predicted.values.is_empty()); +} + +#[test] +fn predicted_vals_resolve_hash_collisions_with_the_backing_rows() { + let mut predicted = PredictedVals::default(); + let table = TableId::from_usize(3); + let collision_key = [Value::from_usize(10), Value::from_usize(11)]; + let key = [Value::from_usize(20), Value::from_usize(21)]; + let collision_row = [collision_key[0], collision_key[1], Value::from_usize(30)]; + let row = [key[0], key[1], Value::from_usize(31)]; + let hash = PredictedVals::hash(table, &key); + + predicted.values.extend_from_slice(&collision_row); + predicted.index.insert_unique( + hash, + PredictedEntry { + hash, + index: 0, + table, + key_arity: collision_key.len() as u32, + }, + |entry| entry.hash, + ); + + let (got, inserted) = predicted.get_or_insert_with(table, &key, row.len(), |values, _| { + values.extend_from_slice(&row[key.len()..]); + }); + assert_eq!(got, row); + assert!(inserted); + + let (got, inserted) = predicted.get_or_insert_with(table, &key, row.len(), |_, _| { + unreachable!("the inserted row should be found through the collision") + }); + assert_eq!(got, row); + assert!(!inserted); +} #[test] fn mask_iter() { From 100f879811ed4ec86a3a2bc6e6c14187609042cd Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Sat, 25 Jul 2026 01:03:25 +0200 Subject: [PATCH 02/12] Reserve fresh IDs per execution state --- core-relations/src/action/mod.rs | 40 +++++++++-- core-relations/src/free_join/mod.rs | 104 +++++++++++++++++++++++++--- core-relations/src/tests.rs | 89 ++++++++++++++++++++++++ egglog-bridge/src/lib.rs | 4 +- 4 files changed, 220 insertions(+), 17 deletions(-) diff --git a/core-relations/src/action/mod.rs b/core-relations/src/action/mod.rs index 5ace8e2ee..8289cfbb6 100644 --- a/core-relations/src/action/mod.rs +++ b/core-relations/src/action/mod.rs @@ -21,7 +21,9 @@ use rustc_hash::FxHasher; use crate::{ BaseValues, ContainerValues, ExternalFunctionId, WrappedTable, common::Value, - free_join::{CounterId, Counters, ExternalFunctions, TableId, TableInfo, Variable}, + free_join::{ + CounterId, CounterReservation, Counters, ExternalFunctions, TableId, TableInfo, Variable, + }, pool::{Clear, Pooled, with_pool_set}, table_spec::{ColumnId, MutationBuffer}, }; @@ -413,6 +415,7 @@ pub(crate) struct DbView<'a> { pub struct ExecutionState<'a> { pub(crate) predicted: PredictedVals, pub(crate) db: DbView<'a>, + counter_reservations: CounterReservations, buffers: MutationBuffers<'a>, /// Whether any mutations have been staged via this ExecutionState. pub(crate) changed: bool, @@ -421,6 +424,19 @@ pub struct ExecutionState<'a> { stop_match: Arc, } +#[derive(Default)] +struct CounterReservations { + ranges: DenseIdMap, +} + +impl CounterReservations { + fn next(&mut self, counters: &Counters, ctr: CounterId) -> usize { + self.ranges + .get_or_insert(ctr, || counters.take_reservation(ctr)) + .next() + } +} + /// Copyable query-side view of an [`ExecutionState`]. Join tasks only need /// immutable database access and the shared early-stop flag; an owned state is /// materialized lazily when an action batch actually executes. @@ -435,6 +451,7 @@ impl<'db> ExecutionStateSeed<'db, '_> { ExecutionState { predicted: Default::default(), db: self.db, + counter_reservations: Default::default(), buffers: MutationBuffers::new(self.db.notification_list, Default::default()), changed: false, stop_match: Arc::clone(self.stop_match), @@ -492,6 +509,7 @@ impl Clone for ExecutionState<'_> { ExecutionState { predicted: Default::default(), db: self.db, + counter_reservations: Default::default(), buffers: self.buffers.clone(), changed: false, stop_match: Arc::clone(&self.stop_match), @@ -507,6 +525,7 @@ impl<'a> ExecutionState<'a> { ExecutionState { predicted: Default::default(), db, + counter_reservations: Default::default(), buffers: MutationBuffers::new(db.notification_list, buffers), changed: false, stop_match: Arc::new(AtomicBool::new(false)), @@ -553,8 +572,8 @@ impl<'a> ExecutionState<'a> { self.db.external_funcs[func].invoke(self, args) } - pub fn inc_counter(&self, ctr: CounterId) -> usize { - self.db.counters.inc(ctr) + pub fn inc_counter(&mut self, ctr: CounterId) -> usize { + self.counter_reservations.next(self.db.counters, ctr) } pub fn read_counter(&self, ctr: CounterId) -> usize { @@ -605,12 +624,15 @@ impl<'a> ExecutionState<'a> { } let row_arity = key.len() + vals.len(); let counters = self.db.counters; + let counter_reservations = &mut self.counter_reservations; let (row, inserted) = self.predicted .get_or_insert_with(table, key, row_arity, |values, _| { for val in vals { values.push(match val { - MergeVal::Counter(ctr) => Value::from_usize(counters.inc(ctr)), + MergeVal::Counter(ctr) => { + Value::from_usize(counter_reservations.next(counters, ctr)) + } MergeVal::Constant(c) => c, }); } @@ -642,12 +664,15 @@ impl<'a> ExecutionState<'a> { } let row_arity = key.len() + vals.len(); let counters = self.db.counters; + let counter_reservations = &mut self.counter_reservations; let (row, inserted) = self.predicted .get_or_insert_with(table, key, row_arity, |values, _| { for val in vals { values.push(match val { - MergeVal::Counter(ctr) => Value::from_usize(counters.inc(ctr)), + MergeVal::Counter(ctr) => { + Value::from_usize(counter_reservations.next(counters, ctr)) + } MergeVal::Constant(c) => c, }); } @@ -755,7 +780,8 @@ impl ExecutionState<'_> { let buffers = &mut self.buffers; // Bind some mutable references because the closure passed // to or_insert_with is `move`. - let ctrs = &self.db.counters; + let ctrs = self.db.counters; + let counter_reservations = &mut self.counter_reservations; let bindings = &bindings; let row_arity = key.as_slice().len() + default.len(); let (row, inserted) = self.predicted.get_or_insert_with( @@ -771,7 +797,7 @@ impl ExecutionState<'_> { bindings[*v][offset] } WriteVal::IncCounter(ctr) => { - Value::from_usize(ctrs.inc(*ctr)) + Value::from_usize(counter_reservations.next(ctrs, *ctr)) } WriteVal::CurrentVal(ix) => values[row_start + *ix], }; diff --git a/core-relations/src/free_join/mod.rs b/core-relations/src/free_join/mod.rs index 82dbe92ee..5d9e571b3 100644 --- a/core-relations/src/free_join/mod.rs +++ b/core-relations/src/free_join/mod.rs @@ -12,6 +12,7 @@ use crate::{ hash_index::IndexCatalog, numeric_id::{DenseIdMap, DenseIdMapWithReuse, NumericId, define_id}, }; +use crossbeam_queue::SegQueue; use egglog_concurrency::{NotificationList, ResettableOnceLock}; use smallvec::SmallVec; @@ -261,15 +262,82 @@ dyn_clone::clone_trait_object!(ExternalFunction); pub(crate) type ExternalFunctions = DenseIdMapWithReuse>; +struct Counter { + next: AtomicUsize, + reservation_size: usize, + recycled: SegQueue>, +} + +impl Counter { + fn new(reservation_size: usize) -> Self { + assert!(reservation_size > 0); + Self { + next: AtomicUsize::new(0), + reservation_size, + recycled: SegQueue::new(), + } + } + + fn take_reservation(&self) -> std::ops::Range { + if self.reservation_size == 1 { + let start = self.next.fetch_add(1, Ordering::Release); + return start..start + 1; + } + self.recycled.pop().unwrap_or_else(|| { + let start = self + .next + .fetch_add(self.reservation_size, Ordering::Release); + start..start + self.reservation_size + }) + } +} + +pub(crate) struct CounterReservation { + counter: Arc, + range: std::ops::Range, +} + +impl CounterReservation { + fn new(counter: Arc) -> Self { + let range = counter.take_reservation(); + Self { counter, range } + } + + pub(crate) fn next(&mut self) -> usize { + if self.range.start == self.range.end { + self.range = self.counter.take_reservation(); + } + let result = self.range.start; + self.range.start += 1; + result + } +} + +impl Drop for CounterReservation { + fn drop(&mut self) { + if !self.range.is_empty() { + self.counter.recycled.push(self.range.clone()); + } + } +} + #[derive(Default)] -pub(crate) struct Counters(DenseIdMap); +pub(crate) struct Counters(DenseIdMap>); impl Clone for Counters { fn clone(&self) -> Counters { let mut map = DenseIdMap::new(); for (k, v) in self.0.iter() { // NB: we may want to experiment with Ordering::Relaxed here. - map.insert(k, AtomicUsize::new(v.load(Ordering::SeqCst))); + let cloned = Counter { + next: AtomicUsize::new(v.next.load(Ordering::SeqCst)), + reservation_size: v.reservation_size, + // The high-water mark already includes every recycled range, + // so omitting the free list is safe and avoids sharing + // reservations between independent database snapshots. + recycled: SegQueue::new(), + }; + map.insert(k, Arc::new(cloned)); } Counters(map) } @@ -277,12 +345,15 @@ impl Clone for Counters { impl Counters { pub(crate) fn read(&self, ctr: CounterId) -> usize { - self.0[ctr].load(Ordering::Acquire) + self.0[ctr].next.load(Ordering::Acquire) } pub(crate) fn inc(&self, ctr: CounterId) -> usize { // We synchronize with `read_counter` but not with other increments. // NB: we may want to experiment with Ordering::Relaxed here. - self.0[ctr].fetch_add(1, Ordering::Release) + self.0[ctr].next.fetch_add(1, Ordering::Release) + } + pub(crate) fn take_reservation(&self, ctr: CounterId) -> CounterReservation { + CounterReservation::new(Arc::clone(&self.0[ctr])) } } @@ -294,10 +365,8 @@ pub struct Database { // NB: some fields are pub(crate) to allow some internal modules to avoid // borrowing the whole table. pub(crate) tables: DenseIdMap, - // TODO: having a single AtomicUsize per counter can lead to contention. We - // should look into prefetching counters when creating a new ExecutionState - // and incrementing locally. Note that the batch size shouldn't be too big - // because we keep an array per id in the UF. + // Reservable counters amortize shared atomic increments across an + // ExecutionState. Exact counters retain one atomic increment per value. pub(crate) counters: Counters, pub(crate) external_functions: ExternalFunctions, container_values: ContainerValues, @@ -500,7 +569,24 @@ impl Database { /// /// These counters can be used to generate unique ids as part of an action. pub fn add_counter(&mut self) -> CounterId { - self.counters.0.push(AtomicUsize::new(0)) + self.counters.0.push(Arc::new(Counter::new(1))) + } + + /// Create a counter whose increments from one [`ExecutionState`] are + /// allocated in local reservations. + /// + /// This is intended for fresh identifiers, where uniqueness matters but a + /// concurrent read need not equal the number of identifiers already + /// returned. Ordinary counters created by [`Database::add_counter`] retain + /// exact increment/read behavior. + /// + /// # Panics + /// + /// Panics if `reservation_size` is zero. + pub fn add_reservable_counter(&mut self, reservation_size: usize) -> CounterId { + self.counters + .0 + .push(Arc::new(Counter::new(reservation_size))) } /// Increment the given counter and return its previous value. diff --git a/core-relations/src/tests.rs b/core-relations/src/tests.rs index a6a9ca497..21577edc1 100644 --- a/core-relations/src/tests.rs +++ b/core-relations/src/tests.rs @@ -69,6 +69,95 @@ fn table_rows(db: &Database, table: TableId) -> Vec> { rows } +#[test] +fn reservable_counter_recycles_execution_state_tails() { + let mut db = Database::new(); + let counter = db.add_reservable_counter(4); + + let first = + db.with_execution_state(|state| [state.inc_counter(counter), state.inc_counter(counter)]); + assert_eq!(first, [0, 1]); + assert_eq!(db.read_counter(counter), 4); + + let second = db.with_execution_state(|state| { + [ + state.inc_counter(counter), + state.inc_counter(counter), + state.inc_counter(counter), + ] + }); + assert_eq!(second, [2, 3, 4]); + assert_eq!(db.read_counter(counter), 8); + + let third = db.with_execution_state(|state| { + [ + state.inc_counter(counter), + state.inc_counter(counter), + state.inc_counter(counter), + ] + }); + assert_eq!(third, [5, 6, 7]); + assert_eq!(db.read_counter(counter), 8); +} + +#[test] +fn ordinary_counter_remains_exact_in_execution_state() { + let mut db = Database::new(); + let counter = db.add_counter(); + + db.with_execution_state(|state| { + assert_eq!(state.inc_counter(counter), 0); + assert_eq!(state.read_counter(counter), 1); + assert_eq!(state.inc_counter(counter), 1); + assert_eq!(state.read_counter(counter), 2); + }); + assert_eq!(db.read_counter(counter), 2); +} + +#[test] +fn reservable_counter_is_unique_across_execution_states() { + const THREADS: usize = 8; + const IDS_PER_THREAD: usize = 1_000; + + let mut db = Database::new(); + let counter = db.add_reservable_counter(64); + let mut ids = std::thread::scope(|scope| { + let handles = (0..THREADS) + .map(|_| { + let db = &db; + scope.spawn(move || { + db.with_execution_state(|state| { + (0..IDS_PER_THREAD) + .map(|_| state.inc_counter(counter)) + .collect::>() + }) + }) + }) + .collect::>(); + handles + .into_iter() + .flat_map(|handle| handle.join().unwrap()) + .collect::>() + }); + + assert_eq!(ids.len(), THREADS * IDS_PER_THREAD); + ids.sort_unstable(); + ids.dedup(); + assert_eq!(ids.len(), THREADS * IDS_PER_THREAD); + let high_water = db.read_counter(counter); + assert!(high_water >= ids.len()); + + let recycled = db.with_execution_state(|state| { + (ids.len()..high_water) + .map(|_| state.inc_counter(counter)) + .collect::>() + }); + ids.extend(recycled); + ids.sort_unstable(); + assert_eq!(ids, (0..high_water).collect::>()); + assert_eq!(db.read_counter(counter), high_water); +} + #[test] fn basic_query() { run_serial_and_parallel(basic_query_inner); diff --git a/egglog-bridge/src/lib.rs b/egglog-bridge/src/lib.rs index bc890c007..2f341046d 100644 --- a/egglog-bridge/src/lib.rs +++ b/egglog-bridge/src/lib.rs @@ -34,6 +34,8 @@ use once_cell::sync::Lazy; use smallvec::SmallVec; use web_time::{Duration, Instant}; +const FRESH_ID_RESERVATION_SIZE: usize = 256; + pub mod macros; pub(crate) mod rule; #[cfg(test)] @@ -204,7 +206,7 @@ impl EGraph { iter::empty(), iter::empty(), ); - let id_counter = db.add_counter(); + let id_counter = db.add_reservable_counter(FRESH_ID_RESERVATION_SIZE); let ts_counter = db.add_counter(); // Start the timestamp counter at 1. db.inc_counter(ts_counter); From 2eeaaec5cd9026bee3ea5c991615aa630904112c Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Sat, 25 Jul 2026 10:42:57 +0200 Subject: [PATCH 03/12] Explain counter reservation lifecycle --- core-relations/src/free_join/mod.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/core-relations/src/free_join/mod.rs b/core-relations/src/free_join/mod.rs index 5d9e571b3..f39f4f86b 100644 --- a/core-relations/src/free_join/mod.rs +++ b/core-relations/src/free_join/mod.rs @@ -262,6 +262,12 @@ dyn_clone::clone_trait_object!(ExternalFunction); pub(crate) type ExternalFunctions = DenseIdMapWithReuse>; +// Reservable counters give each execution state a disjoint range of values. +// Values are then handed out by advancing the state's local `range`, avoiding +// a shared atomic operation per value. Dropping the reservation returns its +// unused suffix to `recycled`, and future states reuse those suffixes before +// advancing the atomic high-water mark. A reservation size of one retains the +// exact increment/read behavior needed by observable counters. struct Counter { next: AtomicUsize, reservation_size: usize, From 70dc650e5bff84f3469cf30c82c1891154a58a99 Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Sat, 25 Jul 2026 11:49:20 +0200 Subject: [PATCH 04/12] Build column indexes from parallel sorted runs --- core-relations/src/hash_index/mod.rs | 370 +++++++++++++++++++++---- core-relations/src/hash_index/tests.rs | 59 +++- 2 files changed, 379 insertions(+), 50 deletions(-) diff --git a/core-relations/src/hash_index/mod.rs b/core-relations/src/hash_index/mod.rs index 052db80c5..5057a927f 100644 --- a/core-relations/src/hash_index/mod.rs +++ b/core-relations/src/hash_index/mod.rs @@ -5,7 +5,7 @@ use std::{ cmp, hash::{Hash, Hasher}, mem, - sync::Mutex, + sync::{Mutex, OnceLock}, }; use crate::{ @@ -26,7 +26,7 @@ use crate::{ parallel, parallel_heuristics::parallelize_index_construction, pool::{Pooled, with_pool_set}, - row_buffer::{RowBuffer, TaggedRowBuffer}, + row_buffer::{RowBuffer, RowSink, TaggedRowBuffer}, table_spec::{ColumnId, Generation, Offset, TableVersion, WrappedTableRef}, }; @@ -302,6 +302,114 @@ struct ColumnIndexShard { subsets: SubsetBuffer, } +struct ColumnIndexPairSink { + pairs: Vec<(Value, RowId)>, +} + +impl ColumnIndexPairSink { + fn with_capacity(capacity: usize) -> Self { + Self { + pairs: Vec::with_capacity(capacity), + } + } +} + +impl RowSink for ColumnIndexPairSink { + fn add_row(&mut self, row_id: RowId, keys: &[Value]) { + for (i, key) in keys.iter().enumerate() { + // A value repeated across covered columns maps the row only once. + if !keys[..i].contains(key) { + self.pairs.push((*key, row_id)); + } + } + } +} + +struct SortedIndexChunk { + /// Ordered by (physical shard, value, row id) and de-duplicated. + pairs: Vec<(Value, RowId)>, + /// Each physical shard owns one contiguous range in `pairs`. + shard_ranges: IdVec>, +} + +fn value_shard_id(shard_data: ShardData, value: Value) -> ShardId { + let mut hasher = FxHasher::default(); + value.hash(&mut hasher); + shard_data.shard_id(hasher.finish()) +} + +fn sort_index_chunk(mut pairs: Vec<(Value, RowId)>, shard_data: ShardData) -> SortedIndexChunk { + if pairs.is_empty() { + let mut shard_ranges = IdVec::with_capacity(shard_data.n_shards()); + shard_ranges.resize_with(shard_data.n_shards(), || 0..0); + return SortedIndexChunk { + pairs, + shard_ranges, + }; + } + + // The scan emits rows in RowId order. The stable value radix sort therefore + // produces (Value, RowId) order without a RowId pass. + let original_len = pairs.len(); + let mut scratch = vec![(Value::new_const(0), RowId::new_const(0)); original_len]; + radix_sort_slice_by_value(&mut pairs, &mut scratch); + pairs.dedup(); + + // A final stable counting pass by physical shard turns each shard's values + // into one contiguous run while preserving (Value, RowId) order within it. + let mut counts = vec![0usize; shard_data.n_shards()]; + for &(value, _) in &pairs { + counts[value_shard_id(shard_data, value).index()] += 1; + } + let mut starts = counts; + let mut prefix = 0; + for count in &mut starts { + let size = *count; + *count = prefix; + prefix += size; + } + let mut next = starts.clone(); + for &pair in &pairs { + let shard = value_shard_id(shard_data, pair.0).index(); + scratch[next[shard]] = pair; + next[shard] += 1; + } + scratch.truncate(pairs.len()); + mem::swap(&mut pairs, &mut scratch); + + let mut shard_ranges = IdVec::with_capacity(shard_data.n_shards()); + for (start, end) in starts.into_iter().zip(next) { + shard_ranges.push(start..end); + } + SortedIndexChunk { + pairs, + shard_ranges, + } +} + +const MIN_COARSE_INDEX_SCAN_ROWS: usize = 256; +// Below this worker count, the radix-sort setup costs more than it saves over +// the existing fine-grained scan-and-split path. +const MIN_COARSE_INDEX_SCAN_THREADS: usize = 8; + +fn coarse_index_scan_partition_size(scan_size: usize, workers: usize) -> usize { + scan_size + .div_ceil(workers.max(1)) + .max(MIN_COARSE_INDEX_SCAN_ROWS) +} + +fn subset_partition<'a>(subset: SubsetRef<'a>, start: usize, end: usize) -> SubsetRef<'a> { + debug_assert!(start <= end); + debug_assert!(end <= subset.size()); + match subset { + SubsetRef::Dense(range) => SubsetRef::Dense(OffsetRange::new( + RowId::from_usize(range.start.index() + start), + RowId::from_usize(range.start.index() + end), + )), + SubsetRef::Sparse(rows) => SubsetRef::Sparse(rows.subslice(start, end)), + } +} + impl Clone for ColumnIndexShard { fn clone(&self) -> Self { ColumnIndexShard { @@ -311,6 +419,36 @@ impl Clone for ColumnIndexShard { } } +impl ColumnIndexShard { + fn merge_sorted_pairs(&mut self, pairs: &[(Value, RowId)]) { + let mut start = 0; + while start < pairs.len() { + let key = pairs[start].0; + let mut end = start + 1; + while end < pairs.len() && pairs[end].0 == key { + end += 1; + } + let run = &pairs[start..end]; + match self.table.entry(key) { + Entry::Occupied(mut occupied) => { + // SAFETY: every sparse subset in this table comes from + // `self.subsets`, and source partitions are consumed in + // ascending RowId order. + unsafe { + occupied + .get_mut() + .merge_sorted_pairs(run, &mut self.subsets); + } + } + Entry::Vacant(vacant) => { + vacant.insert(buffered_subset_from_sorted_pairs(run, &mut self.subsets)); + } + } + start = end; + } + } +} + #[derive(Clone)] pub struct ColumnIndex { // A specialized index used when we are indexing on a single column. @@ -374,39 +512,87 @@ impl IndexBase for ColumnIndex { } fn merge_parallel(&mut self, cols: &[ColumnId], table: WrappedTableRef, subset: SubsetRef) { - const BATCH_SIZE: usize = 1024; let shard_data = self.shard_data; - let mut queues = IdVec::>>::with_capacity( - shard_data.n_shards(), - ); - queues.resize_with(shard_data.n_shards(), || { - Mutex::new(Vec::with_capacity((subset.size() / BATCH_SIZE) + 1)) - }); - let split_buf = |buf: TaggedRowBuffer| { - let mut split = IdVec::::default(); - split.resize_with(shard_data.n_shards(), || TaggedRowBuffer::new(1)); - for (row_id, keys) in buf.iter() { - for (i, key) in keys.iter().enumerate() { - // Match `add_row`: a value repeated across this row's covered columns is - // recorded once, so a value's subset never holds a duplicate row id. - if keys[..i].contains(key) { - continue; + let workers = parallel::current_num_threads(); + let use_coarse_scan = workers >= MIN_COARSE_INDEX_SCAN_THREADS; + let partition_size = coarse_index_scan_partition_size(subset.size(), workers); + let n_partitions = subset.size().div_ceil(partition_size); + + run_in_index_thread_pool(|| { + if use_coarse_scan { + // A fixed slot per logical partition transfers ownership of + // each sorted backing vector out of its producer task. After + // the scope joins, shard workers can borrow their contiguous + // ranges directly without Arc ownership or another copy. + let chunks = (0..n_partitions) + .map(|_| OnceLock::new()) + .collect::>>(); + egglog_concurrency::scope(|inner| { + for (partition_id, start) in + (0..subset.size()).step_by(partition_size).enumerate() + { + let end = cmp::min(start + partition_size, subset.size()); + let partition = subset_partition(subset, start, end); + let slot = &chunks[partition_id]; + inner.spawn(move |_| { + let capacity = partition.size().saturating_mul(cols.len()); + let mut sink = ColumnIndexPairSink::with_capacity(capacity); + let next = table.scan_project( + partition, + cols, + Offset::new(0), + usize::MAX, + &[], + &mut sink, + ); + debug_assert!(next.is_none()); + slot.set(sort_index_chunk(sink.pairs, shard_data)) + .unwrap_or_else(|_| unreachable!("partition slot set twice")); + }); } - shard_data - .get_shard_mut(*key, &mut split) - .add_row(row_id, &[*key]); - } - } - for (shard_id, buf) in split.drain() { - if buf.is_empty() { - continue; - } - let first = buf.get_row(RowId::new(0)).0; - queues[shard_id].lock().unwrap().push((first, buf)); + }); + parallel::for_each_id_vec_mut(&mut self.shards, |shard_id, shard| { + for slot in &chunks { + let chunk = slot + .get() + .expect("all index scan partitions completed before population"); + let range = chunk.shard_ranges[shard_id].clone(); + shard.merge_sorted_pairs(&chunk.pairs[range]); + } + }); + return; } - }; - run_in_index_thread_pool(|| { + const BATCH_SIZE: usize = 1024; + let mut queues = IdVec::>>::with_capacity( + shard_data.n_shards(), + ); + queues.resize_with(shard_data.n_shards(), || { + Mutex::new(Vec::with_capacity((subset.size() / BATCH_SIZE) + 1)) + }); + let split_buf = |buf: TaggedRowBuffer| { + let mut split = IdVec::::default(); + split.resize_with(shard_data.n_shards(), || TaggedRowBuffer::new(1)); + for (row_id, keys) in buf.iter() { + for (i, key) in keys.iter().enumerate() { + // Match `add_row`: a value repeated across this row's covered columns is + // recorded once, so a value's subset never holds a duplicate row id. + if keys[..i].contains(key) { + continue; + } + shard_data + .get_shard_mut(*key, &mut split) + .add_row(row_id, &[*key]); + } + } + for (shard_id, buf) in split.drain() { + if buf.is_empty() { + continue; + } + let first = buf.get_row(RowId::new(0)).0; + queues[shard_id].lock().unwrap().push((first, buf)); + } + }; egglog_concurrency::scope(|inner| { let mut cur = Offset::new(0); loop { @@ -422,7 +608,6 @@ impl IndexBase for ColumnIndex { } } }); - parallel::for_each_id_vec_mut(&mut self.shards, |shard_id, shard| { // Sort the vector by start row id to ensure we populate subsets in sorted order. let mut vec = queues[shard_id].lock().unwrap(); @@ -727,25 +912,11 @@ impl ColumnIndex { while i < pairs.len() { let key = pairs[i].0; let start = i; - let mut first = pairs[i].1; - let mut last = pairs[i].1; while i < pairs.len() && pairs[i].0 == key { - last = cmp::max(last, pairs[i].1); - first = cmp::min(first, pairs[i].1); i += 1; } let shard = self.shard_data.get_shard_mut(key, &mut self.shards); - let count = i - start; - let buffered = if last.rep() - first.rep() == (count - 1) as u32 { - // If the row ids are contiguous, we can represent the subset as a dense range - // to avoid allocations - BufferedSubset::Dense(OffsetRange::new(first, last.inc())) - } else { - let bv = shard - .subsets - .new_vec(pairs[start..i].iter().map(|&(_, r)| r)); - BufferedSubset::Sparse(bv) - }; + let buffered = buffered_subset_from_sorted_pairs(&pairs[start..i], &mut shard.subsets); shard.table.insert(key, buffered); } } @@ -1168,6 +1339,48 @@ impl SubsetBuffer { res } + fn extend_vec( + &mut self, + vec: BufferedVec, + rows: impl ExactSizeIterator, + ) -> BufferedVec { + let old_len = vec.len(); + let added = rows.len(); + if added == 0 { + return vec; + } + let new_len = old_len + added; + let old_capacity = old_len.next_power_of_two(); + if new_len <= old_capacity { + let mut end = vec.1; + for row in rows { + self.buf[end.index()] = row; + end = end.inc(); + } + return BufferedVec(vec.0, end); + } + + let start = if let Some(start) = self.free_list.get_size_class(new_len).pop() { + start + } else { + let start = BufferIndex::from_usize(self.buf.len()); + self.buf.resize( + start.index() + new_len.next_power_of_two(), + RowId::new(u32::MAX), + ); + start + }; + self.buf + .copy_within(vec.0.index()..vec.1.index(), start.index()); + let mut end = BufferIndex::from_usize(start.index() + old_len); + for row in rows { + self.buf[end.index()] = row; + end = end.inc(); + } + self.return_vec(vec); + BufferedVec(start, end) + } + fn make_ref<'a>(&'a self, vec: &BufferedVec) -> SubsetRef<'a> { // SAFETY: if `vec` is a valid index into self.buf, it will be sorted. // @@ -1240,6 +1453,50 @@ impl BufferedSubset { } } + /// Merge one nonempty `(Value, RowId)` run whose rows are strictly sorted + /// and newer than every row already in this subset. + /// + /// *Safety:* callers must ensure that `self` is either dense or comes from + /// `buf`, and that the ordering preconditions above hold. + unsafe fn merge_sorted_pairs(&mut self, pairs: &[(Value, RowId)], buf: &mut SubsetBuffer) { + debug_assert!(!pairs.is_empty()); + debug_assert!( + pairs + .windows(2) + .all(|pair| { pair[0].0 == pair[1].0 && pair[0].1 < pair[1].1 }) + ); + let first = pairs[0].1; + let last = pairs[pairs.len() - 1].1; + let incoming_is_dense = last.rep() - first.rep() == u32::try_from(pairs.len() - 1).unwrap(); + + match self { + BufferedSubset::Dense(range) if range.start == range.end => { + *self = buffered_subset_from_sorted_pairs(pairs, buf); + } + BufferedSubset::Dense(range) => { + debug_assert!(range.end <= first); + if range.end == first && incoming_is_dense { + range.end = last.inc(); + } else { + let old_start = range.start.rep(); + let old_len = range.end.index() - range.start.index(); + let rows = (0..old_len + pairs.len()).map(|i| { + if i < old_len { + RowId::new(old_start + u32::try_from(i).unwrap()) + } else { + pairs[i - old_len].1 + } + }); + *self = BufferedSubset::Sparse(buf.new_vec(rows)); + } + } + BufferedSubset::Sparse(vec) => { + debug_assert!(buf.buf[vec.1.index() - 1] < first); + *vec = buf.extend_vec(mem::take(vec), pairs.iter().map(|&(_, row)| row)); + } + } + } + fn empty() -> Self { BufferedSubset::Dense(OffsetRange::new(RowId::new(0), RowId::new(0))) } @@ -1256,6 +1513,25 @@ impl BufferedSubset { } } +fn buffered_subset_from_sorted_pairs( + pairs: &[(Value, RowId)], + buf: &mut SubsetBuffer, +) -> BufferedSubset { + debug_assert!(!pairs.is_empty()); + debug_assert!( + pairs + .windows(2) + .all(|pair| { pair[0].0 == pair[1].0 && pair[0].1 < pair[1].1 }) + ); + let first = pairs[0].1; + let last = pairs[pairs.len() - 1].1; + if last.rep() - first.rep() == u32::try_from(pairs.len() - 1).unwrap() { + BufferedSubset::Dense(OffsetRange::new(first, last.inc())) + } else { + BufferedSubset::Sparse(buf.new_vec(pairs.iter().map(|&(_, row)| row))) + } +} + fn num_shards() -> usize { let n_threads = parallel::current_num_threads(); if n_threads == 1 { 1 } else { n_threads * 2 } diff --git a/core-relations/src/hash_index/tests.rs b/core-relations/src/hash_index/tests.rs index 317fd5fb2..f19cd4e87 100644 --- a/core-relations/src/hash_index/tests.rs +++ b/core-relations/src/hash_index/tests.rs @@ -6,7 +6,7 @@ use rand::{Rng, SeedableRng, rngs::StdRng}; use crate::{ common::{HashSet, Value}, numeric_id::NumericId, - offsets::{Offsets, RowId, SubsetRef}, + offsets::{Offsets, RowId, Subset, SubsetRef}, }; use crate::{ @@ -16,7 +16,9 @@ use crate::{ table_spec::{ColumnId, WrappedTable}, }; -use super::{Index, IndexBase, IndexPosition, PositionedIndexBase}; +use super::{ + Index, IndexBase, IndexPosition, PositionedIndexBase, coarse_index_scan_partition_size, +}; fn subset_rows(subset: SubsetRef<'_>) -> Vec { let mut rows = Vec::new(); @@ -24,6 +26,13 @@ fn subset_rows(subset: SubsetRef<'_>) -> Vec { rows } +#[test] +fn coarse_index_scan_partitions_are_bounded() { + assert_eq!(coarse_index_scan_partition_size(0, 0), 256); + assert_eq!(coarse_index_scan_partition_size(1_000, 4), 256); + assert_eq!(coarse_index_scan_partition_size(4_096, 3), 1_366); +} + #[test] fn basic_updates() { // Get slightly higher coverage with nondeterministic parallelism. @@ -226,7 +235,7 @@ fn column_index_rebuild_matches_oracle() { // Install a multi-thread pool so `ColumnIndex` shards across several shards and // the merge actually forks; correctness must not depend on the shard count. - let parallel = ThreadPool::new(4).install(|| { + let parallel = ThreadPool::new(8).install(|| { let mut ci = ColumnIndex::new(); ci.merge_parallel(&cols, table.as_ref(), table.all().as_ref()); ci @@ -237,6 +246,50 @@ fn column_index_rebuild_matches_oracle() { } } +#[test] +fn parallel_column_index_handles_sparse_source_partitions() { + let rows = (0..1_000) + .map(|i| vec![v(i), v(i % 17), v((i / 3) % 23), v(i % 17)]) + .collect::>(); + let selected = (0..rows.len()).filter(|i| i % 3 == 0).collect::>(); + let mut first_subset = Subset::empty(); + let mut second_subset = Subset::empty(); + for (i, &row) in selected.iter().enumerate() { + let subset = if i < selected.len() / 2 { + &mut first_subset + } else { + &mut second_subset + }; + subset.add_row_sorted(RowId::from_usize(row)); + } + + let mut expected = BTreeMap::>::new(); + for &row_id in &selected { + let row = &rows[row_id]; + let mut seen = Vec::new(); + for value in &row[1..] { + if seen.contains(value) { + continue; + } + seen.push(*value); + expected.entry(value.rep()).or_default().push(row_id); + } + } + + let table = WrappedTable::new(fill_table(rows, 1, None, |old, new| { + assert_eq!(old, new, "unique keys, so no conflicts"); + None + })); + let cols = [ColumnId::new(1), ColumnId::new(2), ColumnId::new(3)]; + let index = ThreadPool::new(8).install(|| { + let mut index = ColumnIndex::new(); + index.merge_parallel(&cols, table.as_ref(), first_subset.as_ref()); + index.merge_parallel(&cols, table.as_ref(), second_subset.as_ref()); + index + }); + assert_matches_oracle(&index, &expected, "sparse source partitions"); +} + #[test] fn physical_shards_partition_index_iteration() { ThreadPool::new(4).install(|| { From 144cb33ee9a49630ea0eff212b946dbb3c318467 Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Sat, 25 Jul 2026 13:09:08 +0200 Subject: [PATCH 05/12] Use coarse partitions for table rebuilds --- core-relations/src/parallel_heuristics.rs | 13 ++ core-relations/src/table/rebuild.rs | 176 ++++++++++++++++------ 2 files changed, 144 insertions(+), 45 deletions(-) diff --git a/core-relations/src/parallel_heuristics.rs b/core-relations/src/parallel_heuristics.rs index e52bf3cf9..1103b4783 100644 --- a/core-relations/src/parallel_heuristics.rs +++ b/core-relations/src/parallel_heuristics.rs @@ -8,6 +8,7 @@ use std::sync::OnceLock; const DEFAULT_DB_LEVEL_OP_CUTOFF: usize = 10_000; const DEFAULT_INDEX_CONSTRUCTION_CUTOFF: usize = 400_000; const DEFAULT_REBUILD_CUTOFF: usize = 400_000; +const DEFAULT_INCREMENTAL_REBUILD_CUTOFF: usize = 16 * 1024; const DEFAULT_INTRA_CONTAINER_CUTOFF: usize = 10_000; const DEFAULT_INTER_CONTAINER_CUTOFF: usize = 8; const DEFAULT_TABLE_OP_CUTOFF: usize = 400_000; @@ -24,6 +25,7 @@ struct Cutoffs { db_level_op: usize, index_construction: usize, rebuild: usize, + incremental_rebuild: usize, intra_container: usize, inter_container: usize, table_op: usize, @@ -48,6 +50,13 @@ pub(crate) fn parallelize_rebuild(table_size: usize) -> bool { should_parallelize(table_size, cutoffs().rebuild) } +/// Whether to process the dirty-id scan of an incremental table rebuild in +/// parallel. Incremental rebuilds reuse coarse worker-local state and need a +/// much lower cutoff than full rebuilds to amortize dispatch. +pub(crate) fn parallelize_incremental_rebuild(dirty_ids: usize) -> bool { + should_parallelize(dirty_ids, cutoffs().incremental_rebuild) +} + /// Whether or not to perform an operation for a given container memo table. pub(crate) fn parallelize_intra_container_op(num_containers: usize) -> bool { should_parallelize(num_containers, cutoffs().intra_container) @@ -89,6 +98,10 @@ fn cutoffs() -> &'static Cutoffs { DEFAULT_INDEX_CONSTRUCTION_CUTOFF, ), rebuild: cutoff("EGGLOG_PARALLEL_REBUILD_CUTOFF", DEFAULT_REBUILD_CUTOFF), + incremental_rebuild: cutoff( + "EGGLOG_PARALLEL_INCREMENTAL_REBUILD_CUTOFF", + DEFAULT_INCREMENTAL_REBUILD_CUTOFF, + ), intra_container: cutoff( "EGGLOG_PARALLEL_INTRA_CONTAINER_CUTOFF", DEFAULT_INTRA_CONTAINER_CUTOFF, diff --git a/core-relations/src/table/rebuild.rs b/core-relations/src/table/rebuild.rs index 64334c3bc..e2702b85d 100644 --- a/core-relations/src/table/rebuild.rs +++ b/core-relations/src/table/rebuild.rs @@ -5,13 +5,13 @@ use std::{cmp, mem}; use crate::numeric_id::NumericId; use crate::{ - ColumnId, ExecutionState, Offset, RowId, Subset, Table, TableId, TaggedRowBuffer, Value, - WrappedTable, + ColumnId, ExecutionState, Offset, OffsetRange, RowId, Subset, SubsetRef, Table, TableId, + TaggedRowBuffer, Value, WrappedTable, common::HashSet, hash_index::{ColumnIndex, Index}, offsets::Offsets, parallel, - parallel_heuristics::parallelize_rebuild, + parallel_heuristics::{parallelize_incremental_rebuild, parallelize_rebuild}, table_spec::{Rebuilder, WrappedTableRef}, }; @@ -127,45 +127,64 @@ impl SortedWritesTable { exec_state: &mut ExecutionState, ) -> bool { self.refresh_rebuild_index(); - let mut buf = TaggedRowBuffer::new(1); - table.scan_project( - to_scan.as_ref(), - &[search_col], - Offset::new(0), - usize::MAX, - &[], - &mut buf, - ); - if parallelize_rebuild(to_scan.size()) { + if parallel::current_num_threads() >= MIN_COARSE_REBUILD_THREADS + && parallelize_incremental_rebuild(to_scan.size()) + { WrappedTableRef::with_wrapper(self, |wrapped| { - let ids = buf.iter().map(|(_, row)| row[0]).collect::>(); - parallel::map(&ids, |_, id| { + let source = table.as_ref(); + let subset = to_scan.as_ref(); + let partition_size = rebuild_partition_size( + subset.size(), + parallel::current_num_threads(), + MIN_REBUILD_PARTITION_ROWS, + ); + let starts = (0..subset.size()) + .step_by(partition_size) + .collect::>(); + parallel::map(&starts, |_, start| { + let partition = subset_partition( + subset, + *start, + cmp::min(*start + partition_size, subset.size()), + ); let mut mutation_buf = self.new_buffer(); let mut exec_state = exec_state.clone(); let mut changed = false; - let Some(subset) = self.rebuild_index.get_subset(id) else { - return changed; - }; let mut scanned = TaggedRowBuffer::new(self.n_columns); - rebuilder.rebuild_subset(wrapped, subset, &mut scanned, &mut exec_state); - for (row_id, row) in scanned.non_stale_mut() { - let to_remove = self.data.get_row(row_id).map(|x| &x[0..self.n_keys]); - if let Some(key) = to_remove { - mutation_buf.stage_remove(key); + source.for_each_col(partition, search_col, &mut |_, id| { + let Some(rows) = self.rebuild_index.get_subset(&id) else { + return; + }; + rebuilder.rebuild_subset(wrapped, rows, &mut scanned, &mut exec_state); + for (row_id, row) in scanned.non_stale_mut() { + let to_remove = self.data.get_row(row_id).map(|x| &x[0..self.n_keys]); + if let Some(key) = to_remove { + mutation_buf.stage_remove(key); + } + changed = true; + insert_row!(self, mutation_buf, row, next_ts); } - changed = true; - insert_row!(self, mutation_buf, row, next_ts); - } + scanned.clear(); + }); changed }) .into_iter() .any(|changed| changed) }) } else { + let mut ids = TaggedRowBuffer::new(1); + table.scan_project( + to_scan.as_ref(), + &[search_col], + Offset::new(0), + usize::MAX, + &[], + &mut ids, + ); let mut scratch = TaggedRowBuffer::new(self.n_columns); let mut changed = false; - for (_, id) in buf.iter() { + for (_, id) in ids.iter() { let Some(subset) = self.rebuild_index.get_subset(&id[0]) else { continue; }; @@ -194,30 +213,67 @@ impl SortedWritesTable { exec_state: &mut ExecutionState, ) -> bool { const STEP_SIZE: usize = 2048; - if parallelize_rebuild(self.data.next_row().index()) { - let max_row = self.data.next_row().index(); - let starts = (0..max_row).step_by(STEP_SIZE).collect::>(); + let max_row = self.data.next_row().index(); + if parallelize_rebuild(max_row) { + let workers = parallel::current_num_threads(); + if workers < MIN_COARSE_REBUILD_THREADS { + // Keep the original fine-grained loop at low thread counts. Reusing + // state across a coarse partition creates larger mutation-buffer + // publications, whose merge cost outweighs the saved setup here. + let starts = (0..max_row).step_by(STEP_SIZE).collect::>(); + return parallel::map(&starts, |_, start| { + let mut mutation_buf = self.new_buffer(); + let mut buf = TaggedRowBuffer::new(self.n_columns); + let mut exec_state = exec_state.clone(); + let mut changed = false; + rebuilder.rebuild_buf( + &self.data.data, + RowId::from_usize(*start), + RowId::from_usize(cmp::min(*start + STEP_SIZE, max_row)), + &mut buf, + &mut exec_state, + ); + for (row_id, row) in buf.non_stale_mut() { + let to_remove = self.data.get_row(row_id).map(|x| &x[0..self.n_keys]); + changed = true; + if let Some(key) = to_remove { + mutation_buf.stage_remove(key); + } + insert_row!(self, mutation_buf, row, next_ts); + } + buf.clear(); + changed + }) + .into_iter() + .any(|changed| changed); + } + + let partition_size = rebuild_partition_size(max_row, workers, STEP_SIZE); + let starts = (0..max_row).step_by(partition_size).collect::>(); parallel::map(&starts, |_, start| { + let partition_end = cmp::min(*start + partition_size, max_row); let mut mutation_buf = self.new_buffer(); let mut buf = TaggedRowBuffer::new(self.n_columns); let mut exec_state = exec_state.clone(); let mut changed = false; - rebuilder.rebuild_buf( - &self.data.data, - RowId::from_usize(*start), - RowId::from_usize(cmp::min(*start + STEP_SIZE, max_row)), - &mut buf, - &mut exec_state, - ); - for (row_id, row) in buf.non_stale_mut() { - let to_remove = self.data.get_row(row_id).map(|x| &x[0..self.n_keys]); - changed = true; - if let Some(key) = to_remove { - mutation_buf.stage_remove(key); + for chunk_start in (*start..partition_end).step_by(STEP_SIZE) { + rebuilder.rebuild_buf( + &self.data.data, + RowId::from_usize(chunk_start), + RowId::from_usize(cmp::min(chunk_start + STEP_SIZE, partition_end)), + &mut buf, + &mut exec_state, + ); + for (row_id, row) in buf.non_stale_mut() { + let to_remove = self.data.get_row(row_id).map(|x| &x[0..self.n_keys]); + changed = true; + if let Some(key) = to_remove { + mutation_buf.stage_remove(key); + } + insert_row!(self, mutation_buf, row, next_ts); } - insert_row!(self, mutation_buf, row, next_ts); + buf.clear(); } - buf.clear(); changed }) .into_iter() @@ -226,7 +282,6 @@ impl SortedWritesTable { let mut buf = TaggedRowBuffer::new(self.n_columns); let mut changed = false; - let max_row = self.data.next_row().index(); for start in (0..max_row).step_by(STEP_SIZE) { rebuilder.rebuild_buf( &self.data.data, @@ -251,6 +306,25 @@ impl SortedWritesTable { } } +const MIN_REBUILD_PARTITION_ROWS: usize = 256; +const MIN_COARSE_REBUILD_THREADS: usize = 4; + +fn rebuild_partition_size(scan_size: usize, workers: usize, minimum: usize) -> usize { + scan_size.div_ceil(workers.max(1)).max(minimum) +} + +fn subset_partition<'a>(subset: SubsetRef<'a>, start: usize, end: usize) -> SubsetRef<'a> { + debug_assert!(start <= end); + debug_assert!(end <= subset.size()); + match subset { + SubsetRef::Dense(range) => SubsetRef::Dense(OffsetRange::new( + RowId::from_usize(range.start.index() + start), + RowId::from_usize(range.start.index() + end), + )), + SubsetRef::Sparse(rows) => SubsetRef::Sparse(rows.subslice(start, end)), + } +} + fn incremental_rebuild(uf_size: usize, table_size: usize, parallel: bool) -> bool { if parallel { table_size > 10_000 && uf_size * 8192 <= table_size @@ -258,3 +332,15 @@ fn incremental_rebuild(uf_size: usize, table_size: usize, parallel: bool) -> boo table_size > 10000 && uf_size * 8 <= table_size } } + +#[cfg(test)] +mod tests { + use super::rebuild_partition_size; + + #[test] + fn rebuild_partitions_are_coarse_and_nonzero() { + assert_eq!(rebuild_partition_size(0, 0, 256), 256); + assert_eq!(rebuild_partition_size(1_000, 4, 256), 256); + assert_eq!(rebuild_partition_size(4_096, 3, 256), 1_366); + } +} From 5c14f69ff6b29b46d3c13658d605bea7c99cd683 Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Sat, 25 Jul 2026 13:09:40 +0200 Subject: [PATCH 06/12] Benchmark production table parallelism --- core-relations/benches/writes.rs | 137 +++++++++++++++++++++++-------- 1 file changed, 104 insertions(+), 33 deletions(-) diff --git a/core-relations/benches/writes.rs b/core-relations/benches/writes.rs index e08689f82..4c50dc513 100644 --- a/core-relations/benches/writes.rs +++ b/core-relations/benches/writes.rs @@ -1,11 +1,10 @@ +use std::sync::Arc; + use divan::{Bencher, counter::ItemsCount}; +use egglog_concurrency::{ThreadPool, scope}; use egglog_core_relations::{Database, SortedWritesTable, Table, Value}; use egglog_numeric_id::NumericId; use rand::{Rng, rng}; -use rayon::{ - ThreadPoolBuilder, - iter::{ParallelBridge, ParallelIterator}, -}; fn main() { divan::main() @@ -65,7 +64,7 @@ fn generate_workload( ops } -#[divan::bench(consts = [1, 2, 4, 8, 16], sample_count=25)] +#[divan::bench(consts = [1, 2, 4, 8, 12, 16], sample_count=25)] fn parallel_insert(bench: Bencher) { const WORKLOAD_SIZE: usize = 4 << 20; bench_workload( @@ -76,7 +75,7 @@ fn parallel_insert(bench: Bencher) { ) } -#[divan::bench(consts = [1, 2, 4, 8, 16], sample_count=25)] +#[divan::bench(consts = [1, 2, 4, 8, 12, 16], sample_count=25)] fn parallel_insert_merge2(bench: Bencher) { const WORKLOAD_SIZE: usize = 4 << 20; bench_workload( @@ -87,7 +86,7 @@ fn parallel_insert_merge2(bench: Bencher) { ) } -#[divan::bench(consts = [1, 2, 4, 8, 16])] +#[divan::bench(consts = [1, 2, 4, 8, 12, 16])] fn parallel_insert_remove_with_collisions(bench: Bencher) { const WORKLOAD_SIZE: usize = 1 << 20; bench_workload( @@ -98,6 +97,87 @@ fn parallel_insert_remove_with_collisions(bench: Bencher) { ) } +#[divan::bench(consts = [1, 2, 4, 8, 12, 16], sample_count = 10)] +fn parallel_delete_only(bench: Bencher) { + const TABLE_SIZE: usize = 1 << 20; + // Stay over the production parallel cutoff (400K) but under the 50% + // compaction threshold, so this benchmark isolates do_delete. + const REMOVALS: usize = 7 << 16; + let rows = Arc::new( + (0..TABLE_SIZE) + .map(|i| { + [ + Value::new(i as u32), + Value::new(i.wrapping_mul(0x9e3779b1) as u32), + Value::new(i.wrapping_mul(0x85ebca6b) as u32), + Value::new(i.wrapping_mul(0xc2b2ae35) as u32), + Value::new(i.wrapping_mul(0x27d4eb2f) as u32), + ] + }) + .collect::>(), + ); + let pool = Arc::new(ThreadPool::new(N)); + bench + .with_inputs({ + let pool = pool.clone(); + let rows = rows.clone(); + move || { + pool.install(|| { + let db = Database::default(); + let mut table = new_table::<3, 5>(); + stage_inserts(&table, &rows); + db.with_execution_state(|es| table.merge(es)); + stage_removals(&table, &rows[..REMOVALS]); + (db, table) + }) + } + }) + .input_counter(|_| ItemsCount::new(REMOVALS)) + .bench_refs(move |input| { + let (db, table) = input; + pool.install(|| db.with_execution_state(|es| table.merge(es))); + }); +} + +fn new_table() -> SortedWritesTable { + SortedWritesTable::new( + K, + C, + None, + vec![], + Box::new(|_, old, new, out: &mut Vec| { + out.extend_from_slice(new); + old != new + }), + ) +} + +fn stage_inserts(table: &SortedWritesTable, rows: &[[Value; C]]) { + scope(|scope| { + for rows in rows.chunks(8 * 1024) { + let mut buf = table.new_buffer(); + scope.spawn(move |_| { + for row in rows { + buf.stage_insert(row); + } + }); + } + }); +} + +fn stage_removals(table: &SortedWritesTable, rows: &[[Value; C]]) { + scope(|scope| { + for rows in rows.chunks(8 * 1024) { + let mut buf = table.new_buffer(); + scope.spawn(move |_| { + for row in rows { + buf.stage_remove(&row[..3]); + } + }); + } + }); +} + fn bench_workload( bench: Bencher, workload: Vec>, @@ -106,38 +186,29 @@ fn bench_workload( ) { const BATCH_SIZE: usize = 1024; let epoch_size = workload.len().next_multiple_of(n_merges) / n_merges; - let pool = ThreadPoolBuilder::new() - .num_threads(threads) - .build() - .unwrap(); + let pool = Arc::new(ThreadPool::new(threads)); let workload_size = workload.len(); bench - .with_inputs(|| { - ( - Database::default(), - SortedWritesTable::new( - K, - C, - None, - vec![], - Box::new(|_, old, new, out: &mut Vec| { - out.extend_from_slice(new); - old != new - }), - ), - ) + .with_inputs({ + let pool = pool.clone(); + move || pool.install(|| (Database::default(), new_table::())) }) .input_counter(move |_| ItemsCount::new(workload_size)) - .bench_values(|(db, mut table)| { + .bench_refs(move |input| { + let (db, table) = input; pool.install(|| { for outer in workload.chunks(epoch_size) { - outer.chunks(BATCH_SIZE).par_bridge().for_each(|batch| { - let mut buf = table.new_buffer(); - for op in batch { - match op { - Operation::Insert(row) => buf.stage_insert(row), - Operation::Remove(key) => buf.stage_remove(key), - } + scope(|scope| { + for batch in outer.chunks(BATCH_SIZE) { + let mut buf = table.new_buffer(); + scope.spawn(move |_| { + for op in batch { + match op { + Operation::Insert(row) => buf.stage_insert(row), + Operation::Remove(key) => buf.stage_remove(key), + } + } + }); } }); db.with_execution_state(|es| table.merge(es)); From 107ef48f77707de094a113b704a26d2e4826cf69 Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Sat, 25 Jul 2026 14:12:43 +0200 Subject: [PATCH 07/12] Optimize parallel table deletion --- core-relations/src/parallel.rs | 54 +++++++++ core-relations/src/table/mod.rs | 182 +++++++++++++++++++++------- core-relations/src/table/rebuild.rs | 28 ++--- core-relations/src/table/tests.rs | 30 +++++ 4 files changed, 234 insertions(+), 60 deletions(-) diff --git a/core-relations/src/parallel.rs b/core-relations/src/parallel.rs index 9d098c8cd..85741fed3 100644 --- a/core-relations/src/parallel.rs +++ b/core-relations/src/parallel.rs @@ -110,6 +110,39 @@ where collect_result_slots(results) } +pub(crate) fn map_chunks_mut(items: &mut [T], f: F) -> Vec +where + T: Send, + R: Send, + F: Fn(usize, &mut [T]) -> R + Sync, +{ + if items.is_empty() { + return Vec::new(); + } + if !enabled_for_len(items.len()) { + return vec![f(0, items)]; + } + + let chunk_len = chunk_len(items.len()); + let chunks = items.len().div_ceil(chunk_len); + let mut results = empty_result_slots(chunks); + egglog_concurrency::scope(|scope| { + let f = &f; + for (chunk_index, (chunk, out)) in items + .chunks_mut(chunk_len) + .zip(results.iter_mut()) + .enumerate() + { + let base = chunk_index * chunk_len; + scope.spawn(move |_| { + *out = Some(f(base, chunk)); + }); + } + }); + + collect_result_slots(results) +} + pub(crate) fn map_dense_id_map_mut(map: &mut DenseIdMap, f: F) -> Vec where K: NumericId, @@ -170,3 +203,24 @@ where collect_result_slots(results) } + +#[cfg(test)] +mod tests { + use super::map_chunks_mut; + + #[test] + fn chunk_map_reports_bases_and_covers_each_item_once() { + let pool = egglog_concurrency::ThreadPool::new(4); + let mut values = vec![0usize; 10]; + let chunks = pool.install(|| { + map_chunks_mut(&mut values, |base, chunk| { + for (offset, value) in chunk.iter_mut().enumerate() { + *value = base + offset + 1; + } + (base, chunk.len()) + }) + }); + assert_eq!(chunks, vec![(0, 3), (3, 3), (6, 3), (9, 1)]); + assert_eq!(values, (1..=10).collect::>()); + } +} diff --git a/core-relations/src/table/mod.rs b/core-relations/src/table/mod.rs index dbfaca9e0..ef4d87c28 100644 --- a/core-relations/src/table/mod.rs +++ b/core-relations/src/table/mod.rs @@ -66,6 +66,21 @@ impl TableEntry { } } +#[derive(Clone, Copy, Debug)] +struct KnownRemoval { + hashcode: HashCode, + row: RowId, +} + +impl KnownRemoval { + fn hashcode(&self) -> u64 { + #[allow(clippy::unnecessary_cast)] + { + self.hashcode as u64 + } + } +} + /// The core data for a table. /// /// This type is a thin wrapper around `RowBuffer`. The big difference is that @@ -227,6 +242,7 @@ impl ArbitraryRowBuffer { struct Buffer { pending_rows: DenseIdMap, pending_removals: DenseIdMap, + pending_known_removals: DenseIdMap>, state: Weak, n_cols: u32, n_keys: u32, @@ -250,6 +266,7 @@ impl MutationBuffer for Buffer { Box::new(Buffer { pending_rows: Default::default(), pending_removals: Default::default(), + pending_known_removals: Default::default(), state: self.state.clone(), n_cols: self.n_cols, n_keys: self.n_keys, @@ -258,6 +275,18 @@ impl MutationBuffer for Buffer { } } +impl Buffer { + fn stage_remove_row(&mut self, row: RowId, key: &[Value]) { + let (shard, hashcode) = hash_code(self.shard_data, key, self.n_keys as _); + self.pending_known_removals + .get_or_insert(shard, Vec::new) + .push(KnownRemoval { + hashcode: hashcode as _, + row, + }); + } +} + impl Drop for Buffer { fn drop(&mut self) { if let Some(state) = self.state.upgrade() { @@ -272,16 +301,25 @@ impl Drop for Buffer { } state.total_rows.fetch_add(rows, Ordering::Relaxed); - let mut rows = 0; + let mut removals = 0; for shard_id in 0..self.pending_removals.n_ids() { let shard = ShardId::from_usize(shard_id); let Some(buf) = self.pending_removals.take(shard) else { continue; }; - rows += buf.len(); + removals += buf.len(); state.pending_removals[shard].push(buf); } - state.total_removals.fetch_add(rows, Ordering::Relaxed); + + for shard_id in 0..self.pending_known_removals.n_ids() { + let shard = ShardId::from_usize(shard_id); + let Some(buf) = self.pending_known_removals.take(shard) else { + continue; + }; + removals += buf.len(); + state.pending_known_removals[shard].push(buf); + } + state.total_removals.fetch_add(removals, Ordering::Relaxed); } } } @@ -536,15 +574,7 @@ impl Table for SortedWritesTable { } fn new_buffer(&self) -> Box { - let n_shards = self.hash.shard_data().n_shards(); - Box::new(Buffer { - pending_rows: DenseIdMap::with_capacity(n_shards), - pending_removals: DenseIdMap::with_capacity(n_shards), - state: Arc::downgrade(&self.pending_state), - n_keys: u32::try_from(self.n_keys).expect("n_keys should fit in u32"), - n_cols: u32::try_from(self.n_columns).expect("n_columns should fit in u32"), - shard_data: self.hash.shard_data(), - }) + Box::new(self.new_table_buffer()) } fn merge(&mut self, exec_state: &mut ExecutionState) -> TableChange { @@ -572,6 +602,19 @@ impl Table for SortedWritesTable { } impl SortedWritesTable { + fn new_table_buffer(&self) -> Buffer { + let n_shards = self.hash.shard_data().n_shards(); + Buffer { + pending_rows: DenseIdMap::with_capacity(n_shards), + pending_removals: DenseIdMap::with_capacity(n_shards), + pending_known_removals: DenseIdMap::with_capacity(n_shards), + state: Arc::downgrade(&self.pending_state), + n_keys: u32::try_from(self.n_keys).expect("n_keys should fit in u32"), + n_cols: u32::try_from(self.n_columns).expect("n_columns should fit in u32"), + shard_data: self.hash.shard_data(), + } + } + /// Create a new [`SortedWritesTable`] with the given number of keys, /// columns, and an optional sort column. /// @@ -609,44 +652,52 @@ impl SortedWritesTable { } } - /// Flush all pending removals, in parallel. + /// Flush pending removals in coarse shard partitions. Each worker first + /// probes and erases all entries in its partition, then marks the removed + /// rows stale as a separate write-only pass. fn parallel_delete(&mut self) -> bool { let shard_data = self.hash.shard_data(); let pending_removals = &self.pending_state.pending_removals; + let pending_known_removals = &self.pending_state.pending_known_removals; let data = &self.data.data; let n_keys = self.n_keys; - let stale_delta: usize = parallel::map_mut(self.hash.mut_shards(), |shard_id, shard| { - let shard_id = ShardId::from_usize(shard_id); - if pending_removals[shard_id].is_empty() { - return 0; - } - let queue = &pending_removals[shard_id]; - let mut marked_stale = 0; - while let Some(buf) = queue.pop() { - buf.for_each(|to_remove| { - let (actual_shard, hc) = hash_code(shard_data, to_remove, n_keys); - assert_eq!(actual_shard, shard_id); - if let Ok(entry) = shard.find_entry(hc, |entry| { - entry.hashcode == (hc as HashCode) - && &data.get_row(entry.row)[0..n_keys] == to_remove - }) { - let (ent, _) = entry.remove(); - // SAFETY: The safety requirements of - // `set_stale_shared` are that there are no - // concurrent accesses to `row`. No other threads - // can access this row within this method because - // different `shards` partition the space - // (guaranteed by the assertion above), and we - // launch at most one thread per shard. - marked_stale += unsafe { !data.set_stale_shared(ent.row) } as usize; + let stale_delta: usize = + parallel::map_chunks_mut(self.hash.mut_shards(), |base_shard, shards| { + let mut removed_rows = with_pool_set(|ps| ps.get::>()); + for (shard_offset, shard) in shards.iter_mut().enumerate() { + let shard_id = ShardId::from_usize(base_shard + shard_offset); + while let Some(buf) = pending_removals[shard_id].pop() { + buf.for_each(|to_remove| { + let (actual_shard, hc) = hash_code(shard_data, to_remove, n_keys); + assert_eq!(actual_shard, shard_id); + if let Ok(entry) = shard.find_entry(hc, |entry| { + entry.hashcode == (hc as HashCode) + && &data.get_row(entry.row)[0..n_keys] == to_remove + }) { + let (ent, _) = entry.remove(); + removed_rows.push(ent.row); + } + }); } - }); - } - marked_stale - }) - .into_iter() - .sum(); - // Update the stale count with the total marked stale. + while let Some(buf) = pending_known_removals[shard_id].pop() { + for removal in buf { + if let Ok(entry) = shard.find_entry(removal.hashcode(), |entry| { + entry.hashcode == removal.hashcode && entry.row == removal.row + }) { + let (ent, _) = entry.remove(); + removed_rows.push(ent.row); + } + } + } + } + // SAFETY: A worker exclusively owns every shard in this + // partition. Live rows have exactly one primary-key entry, + // so successful removals yield disjoint rows across tasks. + unsafe { mark_rows_stale(data, &removed_rows) }; + removed_rows.len() + }) + .into_iter() + .sum(); self.data.stale_rows += stale_delta; stale_delta > 0 } @@ -659,8 +710,7 @@ impl SortedWritesTable { .enumerate() .for_each(|(shard_id, shard)| { let shard_id = ShardId::from_usize(shard_id); - let queue = &self.pending_state.pending_removals[shard_id]; - while let Some(buf) = queue.pop() { + while let Some(buf) = self.pending_state.pending_removals[shard_id].pop() { buf.for_each(|to_remove| { let (actual_shard, hc) = hash_code(shard_data, to_remove, self.n_keys); assert_eq!(actual_shard, shard_id); @@ -675,6 +725,17 @@ impl SortedWritesTable { } }) } + while let Some(buf) = self.pending_state.pending_known_removals[shard_id].pop() { + for removal in buf { + if let Ok(entry) = shard.find_entry(removal.hashcode(), |entry| { + entry.hashcode == removal.hashcode && entry.row == removal.row + }) { + let (ent, _) = entry.remove(); + self.data.set_stale(ent.row); + changed = true; + } + } + } }); changed } @@ -1277,10 +1338,20 @@ fn hash_code(shard_data: ShardData, row: &[Value], n_keys: usize) -> (ShardId, u (shard_data.shard_id(full_code), full_code as HashCode as u64) } +/// # Safety +/// The caller must ensure exclusive access to every row in `rows`. +unsafe fn mark_rows_stale(data: &RowBuffer, rows: &[RowId]) { + for row in rows { + let was_stale = unsafe { data.set_stale_shared(*row) }; + debug_assert!(!was_stale); + } +} + /// A simple struct for packaging up pending mutations to a `SortedWritesTable`. struct PendingState { pending_rows: DenseIdMap>, pending_removals: DenseIdMap>, + pending_known_removals: DenseIdMap>>, total_removals: AtomicUsize, total_rows: AtomicUsize, } @@ -1290,14 +1361,17 @@ impl PendingState { let n_shards = shard_data.n_shards(); let mut pending_rows = DenseIdMap::with_capacity(n_shards); let mut pending_removals = DenseIdMap::with_capacity(n_shards); + let mut pending_known_removals = DenseIdMap::with_capacity(n_shards); for i in 0..n_shards { pending_rows.insert(ShardId::from_usize(i), SegQueue::default()); pending_removals.insert(ShardId::from_usize(i), SegQueue::default()); + pending_known_removals.insert(ShardId::from_usize(i), SegQueue::default()); } PendingState { pending_rows, pending_removals, + pending_known_removals, total_removals: AtomicUsize::new(0), total_rows: AtomicUsize::new(0), } @@ -1310,6 +1384,10 @@ impl PendingState { for (_, queue) in self.pending_removals.iter() { while queue.pop().is_some() {} } + + for (_, queue) in self.pending_known_removals.iter() { + while queue.pop().is_some() {} + } } /// This is only really used in debugging, but it's annoying enough to write @@ -1320,6 +1398,7 @@ impl PendingState { fn deep_copy(&self) -> PendingState { let mut pending_rows = DenseIdMap::new(); let mut pending_removals = DenseIdMap::new(); + let mut pending_known_removals = DenseIdMap::new(); fn drain_queue(queue: &SegQueue) -> Vec { let mut res = Vec::new(); while let Some(x) = queue.pop() { @@ -1347,9 +1426,20 @@ impl PendingState { pending_removals.insert(shard, new_queue); } + for (shard, queue) in self.pending_known_removals.iter() { + let contents = drain_queue(queue); + let new_queue = SegQueue::default(); + for x in contents { + new_queue.push(x.clone()); + queue.push(x); + } + pending_known_removals.insert(shard, new_queue); + } + PendingState { pending_rows, pending_removals, + pending_known_removals, total_removals: AtomicUsize::new(self.total_removals.load(Ordering::Acquire)), total_rows: AtomicUsize::new(self.total_rows.load(Ordering::Acquire)), } diff --git a/core-relations/src/table/rebuild.rs b/core-relations/src/table/rebuild.rs index e2702b85d..d4d82795e 100644 --- a/core-relations/src/table/rebuild.rs +++ b/core-relations/src/table/rebuild.rs @@ -5,14 +5,14 @@ use std::{cmp, mem}; use crate::numeric_id::NumericId; use crate::{ - ColumnId, ExecutionState, Offset, OffsetRange, RowId, Subset, SubsetRef, Table, TableId, + ColumnId, ExecutionState, Offset, OffsetRange, RowId, Subset, SubsetRef, TableId, TaggedRowBuffer, Value, WrappedTable, common::HashSet, hash_index::{ColumnIndex, Index}, offsets::Offsets, parallel, parallel_heuristics::{parallelize_incremental_rebuild, parallelize_rebuild}, - table_spec::{Rebuilder, WrappedTableRef}, + table_spec::{MutationBuffer, Rebuilder, WrappedTableRef}, }; use super::SortedWritesTable; @@ -97,7 +97,7 @@ impl SortedWritesTable { } let mut changed = false; - let mut mutation_buf = self.new_buffer(); + let mut mutation_buf = self.new_table_buffer(); let mut refreshed_row = Vec::::with_capacity(self.n_columns); for row_id in candidate_rows { let Some(current_row) = self.data.get_row(row_id) else { @@ -105,7 +105,7 @@ impl SortedWritesTable { }; // Preserve the logical row and only advance its sort/timestamp // column, so seminaive treats this as a fresh parent-row delta. - mutation_buf.stage_remove(¤t_row[0..self.n_keys]); + mutation_buf.stage_remove_row(row_id, ¤t_row[0..self.n_keys]); refreshed_row.clear(); refreshed_row.extend_from_slice(current_row); if let Some(sort_by) = self.sort_by { @@ -148,7 +148,7 @@ impl SortedWritesTable { *start, cmp::min(*start + partition_size, subset.size()), ); - let mut mutation_buf = self.new_buffer(); + let mut mutation_buf = self.new_table_buffer(); let mut exec_state = exec_state.clone(); let mut changed = false; let mut scanned = TaggedRowBuffer::new(self.n_columns); @@ -160,7 +160,7 @@ impl SortedWritesTable { for (row_id, row) in scanned.non_stale_mut() { let to_remove = self.data.get_row(row_id).map(|x| &x[0..self.n_keys]); if let Some(key) = to_remove { - mutation_buf.stage_remove(key); + mutation_buf.stage_remove_row(row_id, key); } changed = true; insert_row!(self, mutation_buf, row, next_ts); @@ -194,10 +194,10 @@ impl SortedWritesTable { changed |= subset.size() > 0; } if !scratch.is_empty() { - let mut write_buf = self.new_buffer(); + let mut write_buf = self.new_table_buffer(); for (row_id, row) in scratch.non_stale_mut() { if let Some(to_remove) = self.data.get_row(row_id).map(|x| &x[0..self.n_keys]) { - write_buf.stage_remove(to_remove); + write_buf.stage_remove_row(row_id, to_remove); } insert_row!(self, write_buf, row, next_ts); } @@ -222,7 +222,7 @@ impl SortedWritesTable { // publications, whose merge cost outweighs the saved setup here. let starts = (0..max_row).step_by(STEP_SIZE).collect::>(); return parallel::map(&starts, |_, start| { - let mut mutation_buf = self.new_buffer(); + let mut mutation_buf = self.new_table_buffer(); let mut buf = TaggedRowBuffer::new(self.n_columns); let mut exec_state = exec_state.clone(); let mut changed = false; @@ -237,7 +237,7 @@ impl SortedWritesTable { let to_remove = self.data.get_row(row_id).map(|x| &x[0..self.n_keys]); changed = true; if let Some(key) = to_remove { - mutation_buf.stage_remove(key); + mutation_buf.stage_remove_row(row_id, key); } insert_row!(self, mutation_buf, row, next_ts); } @@ -252,7 +252,7 @@ impl SortedWritesTable { let starts = (0..max_row).step_by(partition_size).collect::>(); parallel::map(&starts, |_, start| { let partition_end = cmp::min(*start + partition_size, max_row); - let mut mutation_buf = self.new_buffer(); + let mut mutation_buf = self.new_table_buffer(); let mut buf = TaggedRowBuffer::new(self.n_columns); let mut exec_state = exec_state.clone(); let mut changed = false; @@ -268,7 +268,7 @@ impl SortedWritesTable { let to_remove = self.data.get_row(row_id).map(|x| &x[0..self.n_keys]); changed = true; if let Some(key) = to_remove { - mutation_buf.stage_remove(key); + mutation_buf.stage_remove_row(row_id, key); } insert_row!(self, mutation_buf, row, next_ts); } @@ -292,10 +292,10 @@ impl SortedWritesTable { ); } if !buf.is_empty() { - let mut write_buf = self.new_buffer(); + let mut write_buf = self.new_table_buffer(); for (row_id, row) in buf.non_stale_mut() { if let Some(to_remove) = self.data.get_row(row_id).map(|x| &x[0..self.n_keys]) { - write_buf.stage_remove(to_remove); + write_buf.stage_remove_row(row_id, to_remove); } insert_row!(self, write_buf, row, next_ts); changed = true; diff --git a/core-relations/src/table/tests.rs b/core-relations/src/table/tests.rs index 05feddae6..a96f14e11 100644 --- a/core-relations/src/table/tests.rs +++ b/core-relations/src/table/tests.rs @@ -150,6 +150,36 @@ fn insert_scan_sorted() { ); } +#[test] +fn known_row_removals_match_the_staged_row_id() { + empty_execution_state!(e); + let mut table = fill_table( + vec![ + vec![v(0), v(1), v(2)], + vec![v(1), v(2), v(3)], + vec![v(2), v(3), v(4)], + ], + 2, + None, + |_, new| Some(new.to_vec()), + ); + let key = [v(1), v(2)]; + let row = table.get_row(&key).unwrap().id; + + let mut wrong = table.new_table_buffer(); + wrong.stage_remove_row(RowId::from_usize(row.index() + 1), &key); + drop(wrong); + table.merge(&mut e); + assert!(table.get_row(&key).is_some()); + + let mut exact = table.new_table_buffer(); + exact.stage_remove_row(row, &key); + drop(exact); + table.merge(&mut e); + assert!(table.get_row(&key).is_none()); + assert_eq!(table.len(), 2); +} + #[test] fn shard_math() { let mut table = ShardedHashTable::::with_shards(14); From 641f425c4a629f69a59e2ce7730b09eb1affd66e Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Sat, 25 Jul 2026 14:13:11 +0200 Subject: [PATCH 08/12] Benchmark large parallel table deletion --- core-relations/benches/writes.rs | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/core-relations/benches/writes.rs b/core-relations/benches/writes.rs index 4c50dc513..a31668dae 100644 --- a/core-relations/benches/writes.rs +++ b/core-relations/benches/writes.rs @@ -103,8 +103,19 @@ fn parallel_delete_only(bench: Bencher) { // Stay over the production parallel cutoff (400K) but under the 50% // compaction threshold, so this benchmark isolates do_delete. const REMOVALS: usize = 7 << 16; + bench_delete_only::(bench, TABLE_SIZE, REMOVALS); +} + +#[divan::bench(consts = [1, 2, 4, 8, 12, 16], sample_count = 5)] +fn parallel_delete_only_large(bench: Bencher) { + const TABLE_SIZE: usize = 4 << 20; + const REMOVALS: usize = 7 << 18; + bench_delete_only::(bench, TABLE_SIZE, REMOVALS); +} + +fn bench_delete_only(bench: Bencher, table_size: usize, removals: usize) { let rows = Arc::new( - (0..TABLE_SIZE) + (0..table_size) .map(|i| { [ Value::new(i as u32), @@ -127,12 +138,12 @@ fn parallel_delete_only(bench: Bencher) { let mut table = new_table::<3, 5>(); stage_inserts(&table, &rows); db.with_execution_state(|es| table.merge(es)); - stage_removals(&table, &rows[..REMOVALS]); + stage_removals(&table, &rows[..removals]); (db, table) }) } }) - .input_counter(|_| ItemsCount::new(REMOVALS)) + .input_counter(move |_| ItemsCount::new(removals)) .bench_refs(move |input| { let (db, table) = input; pool.install(|| db.with_execution_state(|es| table.merge(es))); From d83bd18e5babb4889746b84473d0593384f2e1cd Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Sun, 26 Jul 2026 11:23:03 -0700 Subject: [PATCH 09/12] Partition table mutations by cached hashes --- core-relations/benches/writes.rs | 252 +++++- core-relations/src/row_buffer/mod.rs | 31 +- core-relations/src/row_buffer/partitioned.rs | 374 +++++++++ core-relations/src/table/mod.rs | 799 +++++++++++-------- core-relations/src/table/tests.rs | 286 ++++++- 5 files changed, 1392 insertions(+), 350 deletions(-) create mode 100644 core-relations/src/row_buffer/partitioned.rs diff --git a/core-relations/benches/writes.rs b/core-relations/benches/writes.rs index a31668dae..21ded490b 100644 --- a/core-relations/benches/writes.rs +++ b/core-relations/benches/writes.rs @@ -4,7 +4,7 @@ use divan::{Bencher, counter::ItemsCount}; use egglog_concurrency::{ThreadPool, scope}; use egglog_core_relations::{Database, SortedWritesTable, Table, Value}; use egglog_numeric_id::NumericId; -use rand::{Rng, rng}; +use rand::{Rng, SeedableRng, rng, rngs::StdRng}; fn main() { divan::main() @@ -42,6 +42,25 @@ fn generate_workload( collision_pct: f64, ) -> Vec> { let mut rng = rng(); + generate_workload_with_rng(n, insert_pct, collision_pct, &mut rng) +} + +fn generate_workload_seeded( + n: usize, + insert_pct: f64, + collision_pct: f64, + seed: u64, +) -> Vec> { + let mut rng = StdRng::seed_from_u64(seed); + generate_workload_with_rng(n, insert_pct, collision_pct, &mut rng) +} + +fn generate_workload_with_rng( + n: usize, + insert_pct: f64, + collision_pct: f64, + rng: &mut impl Rng, +) -> Vec> { let mut ops = Vec::>::with_capacity(n); for _ in 0..n { if !rng.random_bool(insert_pct) && !ops.is_empty() { @@ -52,18 +71,44 @@ fn generate_workload( ops.push(Operation::Remove(key.try_into().unwrap())); } else if rng.random_bool(collision_pct) && !ops.is_empty() { let key = ops[rng.random_range(0..ops.len())].key(); - let mut row = random_row::(&mut rng); + let mut row = random_row::(rng); for (dst, src) in row.iter_mut().zip(key.iter()) { *dst = *src; } ops.push(Operation::Insert(row)); } else { - ops.push(Operation::Insert(random_row::(&mut rng))); + ops.push(Operation::Insert(random_row::(rng))); } } ops } +fn generate_duplicate_heavy_workload_seeded( + n: usize, + unique_keys: usize, + seed: u64, +) -> Vec> { + assert!(unique_keys > 0); + assert!(unique_keys <= n); + let mut rng = StdRng::seed_from_u64(seed); + let mut unique_rows = Vec::with_capacity(unique_keys); + for _ in 0..unique_keys { + unique_rows.push(random_row::(&mut rng)); + } + + let mut ops = Vec::with_capacity(n); + for row in &unique_rows { + ops.push(Operation::Insert(*row)); + } + while ops.len() < n { + let base = &unique_rows[rng.random_range(0..unique_rows.len())]; + let mut row = random_row::(&mut rng); + row[..K].copy_from_slice(&base[..K]); + ops.push(Operation::Insert(row)); + } + ops +} + #[divan::bench(consts = [1, 2, 4, 8, 12, 16], sample_count=25)] fn parallel_insert(bench: Bencher) { const WORKLOAD_SIZE: usize = 4 << 20; @@ -97,6 +142,129 @@ fn parallel_insert_remove_with_collisions(bench: Bencher) { ) } +const INSERT_SCALABILITY_WORKLOAD_SIZE: usize = 4 << 20; +const INSERT_SCALABILITY_BATCH_SIZE: usize = 8 * 1024; +const DUPLICATE_HEAVY_KEYS: usize = 1 << 16; + +/// Deterministic insertion scalability with mutation staging outside the +/// timer. The input has 5% primary-key collisions. +#[divan::bench(consts = [1, 2, 4, 8, 12], sample_count = 5)] +fn parallel_insert_merge_only_seeded_5pct_collisions(bench: Bencher) { + const SEED: u64 = 0x6d65_7267_652d_6f6e; + bench_insert_merge_only::( + bench, + Arc::new(generate_workload_seeded::<3, 5>( + INSERT_SCALABILITY_WORKLOAD_SIZE, + 1.0, + 0.05, + SEED, + )), + new_table::<3, 5>, + ); +} + +/// Deterministic control containing only vacant-key insertions. +#[divan::bench(consts = [1, 2, 4, 8, 12], sample_count = 5)] +fn parallel_insert_merge_only_all_new(bench: Bencher) { + const SEED: u64 = 0x7265_7365_7276_652d; + bench_insert_merge_only::( + bench, + Arc::new(generate_workload_seeded::<3, 5>( + INSERT_SCALABILITY_WORKLOAD_SIZE, + 1.0, + 0.0, + SEED, + )), + new_table::<3, 5>, + ); +} + +/// Duplicate-heavy insertion accepting incoming values during key conflicts. +#[divan::bench(consts = [1, 2, 4, 8, 12], sample_count = 5)] +fn parallel_insert_merge_only_duplicate_heavy(bench: Bencher) { + const SEED: u64 = 0x6475_706c_6963_6174; + bench_insert_merge_only::( + bench, + Arc::new(generate_duplicate_heavy_workload_seeded::<3, 5>( + INSERT_SCALABILITY_WORKLOAD_SIZE, + DUPLICATE_HEAVY_KEYS, + SEED, + )), + new_table::<3, 5>, + ); +} + +/// Duplicate-heavy insertion whose merge function rejects every update. +/// The timed ordinary `Table::merge` includes stale-row compaction. +#[divan::bench(consts = [1, 2, 4, 8, 12], sample_count = 5)] +fn parallel_insert_merge_only_duplicate_heavy_rejecting(bench: Bencher) { + const SEED: u64 = 0x7265_6a65_6374_6564; + bench_insert_merge_only::( + bench, + Arc::new(generate_duplicate_heavy_workload_seeded::<3, 5>( + INSERT_SCALABILITY_WORKLOAD_SIZE, + DUPLICATE_HEAVY_KEYS, + SEED, + )), + new_rejecting_table, + ); +} + +/// Every mutation updates an existing key. The fixture has enough live rows +/// that the accepted updates stay at (rather than cross) the 50% compaction +/// threshold, isolating occupied-entry insertion from compaction. +#[divan::bench(consts = [1, 2, 4, 8, 12], sample_count = 5)] +fn parallel_insert_merge_only_existing_updates(bench: Bencher) { + const WORKLOAD_SIZE: usize = 1 << 20; + let initial_rows = Arc::new( + (0..WORKLOAD_SIZE) + .map(|i| { + [ + Value::new(i as u32), + Value::new(i.wrapping_mul(0x9e3779b1) as u32), + Value::new(i.wrapping_mul(0x85ebca6b) as u32), + Value::new(i.wrapping_mul(0xc2b2ae35) as u32), + Value::new(i.wrapping_mul(0x27d4eb2f) as u32), + ] + }) + .collect::>(), + ); + let update_rows = Arc::new( + initial_rows + .iter() + .enumerate() + .map(|(i, row)| { + let mut update = *row; + update[3] = Value::new(i.wrapping_mul(0x165667b1).wrapping_add(1) as u32); + update[4] = Value::new(i.wrapping_mul(0xd3a2646c).wrapping_add(1) as u32); + update + }) + .collect::>(), + ); + let pool = Arc::new(ThreadPool::new(N)); + bench + .with_inputs({ + let pool = pool.clone(); + let initial_rows = initial_rows.clone(); + let update_rows = update_rows.clone(); + move || { + pool.install(|| { + let db = Database::default(); + let mut table = new_table::<3, 5>(); + stage_inserts(&table, &initial_rows); + db.with_execution_state(|es| table.merge(es)); + stage_inserts(&table, &update_rows); + (db, table) + }) + } + }) + .input_counter(|_| ItemsCount::new(WORKLOAD_SIZE)) + .bench_refs(move |input| { + let (db, table) = input; + pool.install(|| db.with_execution_state(|es| table.merge(es))); + }); +} + #[divan::bench(consts = [1, 2, 4, 8, 12, 16], sample_count = 10)] fn parallel_delete_only(bench: Bencher) { const TABLE_SIZE: usize = 1 << 20; @@ -113,6 +281,37 @@ fn parallel_delete_only_large(bench: Bencher) { bench_delete_only::(bench, TABLE_SIZE, REMOVALS); } +fn bench_insert_merge_only( + bench: Bencher, + workload: Arc>>, + make_table: fn() -> SortedWritesTable, +) { + let pool = Arc::new(ThreadPool::new(N)); + let workload_size = workload.len(); + bench + .with_inputs({ + let pool = pool.clone(); + let workload = workload.clone(); + move || { + pool.install(|| { + let db = Database::default(); + let table = make_table(); + stage_operations_with_batch_size( + &table, + &workload, + INSERT_SCALABILITY_BATCH_SIZE, + ); + (db, table) + }) + } + }) + .input_counter(move |_| ItemsCount::new(workload_size)) + .bench_refs(move |input| { + let (db, table) = input; + pool.install(|| db.with_execution_state(|es| table.merge(es))); + }); +} + fn bench_delete_only(bench: Bencher, table_size: usize, removals: usize) { let rows = Arc::new( (0..table_size) @@ -163,6 +362,10 @@ fn new_table() -> SortedWritesTable { ) } +fn new_rejecting_table() -> SortedWritesTable { + SortedWritesTable::new(3, 5, None, vec![], Box::new(|_, _, _, _| false)) +} + fn stage_inserts(table: &SortedWritesTable, rows: &[[Value; C]]) { scope(|scope| { for rows in rows.chunks(8 * 1024) { @@ -189,13 +392,40 @@ fn stage_removals(table: &SortedWritesTable, rows: &[[Value; C]] }); } +fn stage_operations( + table: &SortedWritesTable, + operations: &[Operation], +) { + stage_operations_with_batch_size(table, operations, 1024); +} + +fn stage_operations_with_batch_size( + table: &SortedWritesTable, + operations: &[Operation], + batch_size: usize, +) { + assert_ne!(batch_size, 0); + scope(|scope| { + for batch in operations.chunks(batch_size) { + let mut buf = table.new_buffer(); + scope.spawn(move |_| { + for op in batch { + match op { + Operation::Insert(row) => buf.stage_insert(row), + Operation::Remove(key) => buf.stage_remove(key), + } + } + }); + } + }); +} + fn bench_workload( bench: Bencher, workload: Vec>, n_merges: usize, threads: usize, ) { - const BATCH_SIZE: usize = 1024; let epoch_size = workload.len().next_multiple_of(n_merges) / n_merges; let pool = Arc::new(ThreadPool::new(threads)); let workload_size = workload.len(); @@ -209,19 +439,7 @@ fn bench_workload( let (db, table) = input; pool.install(|| { for outer in workload.chunks(epoch_size) { - scope(|scope| { - for batch in outer.chunks(BATCH_SIZE) { - let mut buf = table.new_buffer(); - scope.spawn(move |_| { - for op in batch { - match op { - Operation::Insert(row) => buf.stage_insert(row), - Operation::Remove(key) => buf.stage_remove(key), - } - } - }); - } - }); + stage_operations(table, outer); db.with_execution_state(|es| table.merge(es)); } }) diff --git a/core-relations/src/row_buffer/mod.rs b/core-relations/src/row_buffer/mod.rs index 214f269e3..f6ae7c977 100644 --- a/core-relations/src/row_buffer/mod.rs +++ b/core-relations/src/row_buffer/mod.rs @@ -13,9 +13,12 @@ use crate::{ pool::{Pooled, with_pool_set}, }; +mod partitioned; #[cfg(test)] mod tests; +pub(crate) use partitioned::{HashPartitioning, PartitionedRowBuffer}; + /// A trait for types that can store a vector of `Value`s. /// /// Can either be backed by a pool or SmallVec. @@ -594,7 +597,8 @@ impl]>> ReadHandle<'_, T> { /// The caller must ensure that either `row` is within bounds of the buffer at the creation of /// this handle, or that the row was successfully written to the buffer before it was called. /// - /// Furthermore, no calls to `set_stale_shared` may overlap with this call. + /// Furthermore, no calls to `set_stale_shared` or `replace_row_shared` may + /// overlap with this call. pub(crate) unsafe fn get_row_unchecked(&self, row: RowId) -> &[Value] { // SAFETY: ParallelVecWriter guarantees that data within bounds is not // being modified concurrently. @@ -606,6 +610,31 @@ impl]>> ReadHandle<'_, T> { } } + /// Replace a row in the buffer through shared access. + /// + /// The row may be beyond the vector's original length if it is covered by + /// a completed write to the underlying [`ParallelVecWriter`]. + /// + /// # Safety + /// + /// The caller must ensure that `row` is covered by the original vector or + /// a completed write and that no other reads or writes overlap this row. + pub(crate) unsafe fn replace_row_shared(&self, row: RowId, values: &[Value]) { + assert_eq!( + values.len(), + self.buf.n_columns, + "attempting to replace a row with mismatched arity" + ); + let cells: &[Cell] = &self.data; + let cell_ptr = cells.as_ptr(); + let start = row.index() * self.buf.n_columns; + for (offset, value) in values.iter().copied().enumerate() { + unsafe { + (&*cell_ptr.add(start + offset)).set(value); + } + } + } + /// See the documentation for [`RowBuffer::set_stale_shared`]. /// /// In addition to the requirements there, `row` is allowed to be out of bounds of the initial diff --git a/core-relations/src/row_buffer/partitioned.rs b/core-relations/src/row_buffer/partitioned.rs new file mode 100644 index 000000000..8535099af --- /dev/null +++ b/core-relations/src/row_buffer/partitioned.rs @@ -0,0 +1,374 @@ +//! Hash-partitioned batches of rows. +//! +//! Producers append rows, unchanged full hashes, and their destination +//! partition densely. Before publication the buffer is sealed with a stable +//! counting partition, giving consumers one contiguous sequence per partition. + +use std::iter::FusedIterator; + +use crate::{common::Value, numeric_id::NumericId}; + +/// Selects a power-of-two hash partition within one of several outer groups. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct HashPartitioning { + shift: u8, + bits: u8, + groups: u32, +} + +impl HashPartitioning { + #[cfg(test)] + pub(crate) fn new(shift: u32, bits: u32) -> Self { + Self::grouped(shift, bits, 1) + } + + pub(crate) fn grouped(shift: u32, bits: u32, groups: usize) -> Self { + assert!( + shift.checked_add(bits).is_some_and(|sum| sum <= 64), + "hash partition bits must fit in a u64" + ); + assert!( + bits < usize::BITS, + "hash partition count must fit in a usize" + ); + assert_ne!(groups, 0, "hash partitioning must have at least one group"); + let groups = u32::try_from(groups).expect("hash partition group count must fit in u32"); + (1usize << bits) + .checked_mul(groups as usize) + .expect("total hash partition count overflow"); + Self { + shift: shift as u8, + bits: bits as u8, + groups, + } + } + + pub(crate) fn partitions_per_group(self) -> usize { + 1usize << self.bits + } + + pub(crate) fn partition_count(self) -> usize { + self.partitions_per_group() * self.groups as usize + } + + #[inline] + pub(crate) fn partition(self, group: usize, hash: u64) -> usize { + assert!( + group < self.groups as usize, + "hash partition group out of bounds" + ); + let within_group = if self.bits == 0 { + 0 + } else { + (hash.wrapping_shr(self.shift as u32) & (self.partitions_per_group() as u64 - 1)) + as usize + }; + self.partition_index(group, within_group) + } + + pub(crate) fn partition_index(self, group: usize, within_group: usize) -> usize { + assert!( + group < self.groups as usize, + "hash partition group out of bounds" + ); + assert!( + within_group < self.partitions_per_group(), + "hash subpartition out of bounds" + ); + group * self.partitions_per_group() + within_group + } +} + +/// One row yielded from a [`PartitionedRowBuffer`]. +pub(crate) struct HashedRow<'a> { + pub(crate) hash: u64, + pub(crate) row: &'a [Value], +} + +/// A stable hash-partitioned row batch. +/// +/// Before sealing, all vectors are dense and insertion ordered. +/// `partition_ids[i]` describes `hashes[i]` and row `i`. Sealing performs a +/// stable counting partition and replaces `partition_ids` with prefix offsets. +#[derive(Clone)] +pub(crate) struct PartitionedRowBuffer { + arity: usize, + partitioning: HashPartitioning, + hashes: Vec, + rows: Vec, + partition_ids: Vec, + offsets: Option>, +} + +impl PartitionedRowBuffer { + pub(crate) fn new(arity: usize, partitioning: HashPartitioning) -> Self { + Self { + arity, + partitioning, + hashes: Vec::new(), + rows: Vec::new(), + partition_ids: Vec::new(), + offsets: None, + } + } + + pub(crate) fn arity(&self) -> usize { + self.arity + } + + pub(crate) fn len(&self) -> usize { + self.hashes.len() + } + + #[cfg(test)] + pub(crate) fn is_empty(&self) -> bool { + self.hashes.is_empty() + } + + pub(crate) fn partition_count(&self) -> usize { + self.partitioning.partition_count() + } + + /// Append a row, its outer group, and its unchanged full hash. + pub(crate) fn add_row(&mut self, group: usize, hash: u64, row: &[Value]) { + assert!( + self.offsets.is_none(), + "cannot append to a sealed partitioned row buffer" + ); + assert_eq!( + row.len(), + self.arity, + "attempting to add a row with mismatched arity" + ); + let partition = self.partitioning.partition(group, hash); + let partition = + u32::try_from(partition).expect("partitioned row buffer partition id overflow"); + self.hashes.push(hash); + self.rows.extend_from_slice(row); + self.partition_ids.push(partition); + } + + /// Convert insertion-ordered staging vectors into contiguous stable + /// partition runs. + pub(crate) fn seal(mut self) -> Self { + if self.offsets.is_some() { + return self; + } + debug_assert_eq!(self.partition_ids.len(), self.hashes.len()); + debug_assert_eq!(self.rows.len(), self.hashes.len() * self.arity); + + if self.partition_count() == 1 { + let len = u32::try_from(self.len()).expect("partitioned row buffer row count overflow"); + self.partition_ids.clear(); + self.offsets = Some(Box::new([0, len])); + return self; + } + + let mut offsets = vec![0u32; self.partition_count() + 1]; + for &partition in &self.partition_ids { + offsets[partition as usize + 1] = offsets[partition as usize + 1] + .checked_add(1) + .expect("partitioned row count overflow"); + } + for partition in 0..self.partition_count() { + offsets[partition + 1] = offsets[partition + 1] + .checked_add(offsets[partition]) + .expect("partitioned row offset overflow"); + } + debug_assert_eq!(offsets[self.partition_count()] as usize, self.len()); + + let mut cursors = offsets[..self.partition_count()].to_vec(); + let mut hashes = vec![0u64; self.len()]; + let row_values = self + .len() + .checked_mul(self.arity) + .expect("partitioned row value count overflow"); + let mut rows = vec![Value::new(0); row_values]; + for source in 0..self.len() { + let partition = self.partition_ids[source] as usize; + let destination = cursors[partition] as usize; + cursors[partition] += 1; + hashes[destination] = self.hashes[source]; + if self.arity != 0 { + rows[destination * self.arity..(destination + 1) * self.arity] + .copy_from_slice(&self.rows[source * self.arity..(source + 1) * self.arity]); + } + } + + self.hashes = hashes; + self.rows = rows; + self.partition_ids = Vec::new(); + self.offsets = Some(offsets.into_boxed_slice()); + self + } + + #[cfg(test)] + pub(crate) fn partition_len(&self, partition: usize) -> usize { + let offsets = self.offsets(); + (offsets[partition + 1] - offsets[partition]) as usize + } + + pub(crate) fn group_is_empty(&self, group: usize) -> bool { + self.group_len(group) == 0 + } + + pub(crate) fn group_len(&self, group: usize) -> usize { + let first = self.partitioning.partition_index(group, 0); + let after_last = first + self.partitioning.partitions_per_group(); + let offsets = self.offsets(); + (offsets[after_last] - offsets[first]) as usize + } + + pub(crate) fn partition(&self, partition: usize) -> PartitionIter<'_> { + let offsets = self.offsets(); + PartitionIter { + buffer: self, + next: offsets[partition] as usize, + end: offsets[partition + 1] as usize, + } + } + + fn offsets(&self) -> &[u32] { + self.offsets + .as_deref() + .expect("partitioned row buffer must be sealed before reading") + } +} + +pub(crate) struct PartitionIter<'a> { + buffer: &'a PartitionedRowBuffer, + next: usize, + end: usize, +} + +impl<'a> Iterator for PartitionIter<'a> { + type Item = HashedRow<'a>; + + #[inline] + fn next(&mut self) -> Option { + if self.next == self.end { + return None; + } + let index = self.next; + self.next += 1; + let row = if self.buffer.arity == 0 { + &[] + } else { + &self.buffer.rows[index * self.buffer.arity..(index + 1) * self.buffer.arity] + }; + Some(HashedRow { + hash: self.buffer.hashes[index], + row, + }) + } + + fn size_hint(&self) -> (usize, Option) { + let remaining = self.end - self.next; + (remaining, Some(remaining)) + } +} + +impl ExactSizeIterator for PartitionIter<'_> {} +impl FusedIterator for PartitionIter<'_> {} + +#[cfg(test)] +mod tests { + use super::*; + + fn v(value: usize) -> Value { + Value::from_usize(value) + } + + #[test] + fn partitions_rows_stably_and_keeps_hashes_aligned() { + let partitioning = HashPartitioning::new(4, 2); + let mut buffer = PartitionedRowBuffer::new(2, partitioning); + let mut expected = [Vec::new(), Vec::new(), Vec::new(), Vec::new()]; + for i in 0..80usize { + let partition = (i * 3) & 3; + let hash = ((partition as u64) << 4) | (i as u64 & 0xf); + let row = [v(i), v(i + 1)]; + buffer.add_row(0, hash, &row); + expected[partition].push((hash, row)); + } + let buffer = buffer.seal(); + + assert_eq!(buffer.len(), 80); + assert_eq!(buffer.arity(), 2); + for (partition, expected) in expected.iter().enumerate() { + assert_eq!(buffer.partition_len(partition), expected.len()); + let mut iter = buffer.partition(partition); + assert_eq!(iter.len(), expected.len()); + let actual = iter + .by_ref() + .map(|hashed| (hashed.hash, hashed.row.to_vec())) + .collect::>(); + assert_eq!(iter.len(), 0); + let expected = expected + .iter() + .map(|(hash, row)| (*hash, row.to_vec())) + .collect::>(); + assert_eq!(actual, expected); + } + } + + #[test] + fn supports_groups_zero_arity_and_logical_clone() { + let partitioning = HashPartitioning::grouped(1, 3, 2); + let mut buffer = PartitionedRowBuffer::new(0, partitioning); + for (group, hash) in [(0, 0), (1, 2), (0, 4), (1, 2), (1, 14), (0, 0)] { + buffer.add_row(group, hash, &[]); + } + let buffer = buffer.seal(); + let cloned = buffer.clone(); + + assert_eq!(cloned.len(), 6); + assert!(!cloned.is_empty()); + assert_eq!(cloned.group_len(0), 3); + assert_eq!(cloned.group_len(1), 3); + for partition in 0..partitioning.partition_count() { + let original = buffer + .partition(partition) + .map(|hashed| (hashed.hash, hashed.row.len())) + .collect::>(); + let copied = cloned + .partition(partition) + .map(|hashed| (hashed.hash, hashed.row.len())) + .collect::>(); + assert_eq!(copied, original); + } + } + + #[test] + fn is_send_and_sync() { + fn assert_send_sync() {} + assert_send_sync::(); + } + + #[test] + fn one_partition_seals_without_copying_rows_or_hashes() { + let partitioning = HashPartitioning::new(0, 0); + let mut buffer = PartitionedRowBuffer::new(2, partitioning); + for i in 0..32usize { + buffer.add_row(0, i as u64, &[v(i), v(i + 1)]); + } + let hashes = buffer.hashes.as_ptr(); + let rows = buffer.rows.as_ptr(); + + let buffer = buffer.seal(); + + assert_eq!(buffer.hashes.as_ptr(), hashes); + assert_eq!(buffer.rows.as_ptr(), rows); + assert!(buffer.partition_ids.is_empty()); + assert_eq!(buffer.partition_len(0), 32); + assert_eq!( + buffer + .partition(0) + .map(|hashed| (hashed.hash, hashed.row.to_vec())) + .collect::>(), + (0..32usize) + .map(|i| (i as u64, vec![v(i), v(i + 1)])) + .collect::>() + ); + } +} diff --git a/core-relations/src/table/mod.rs b/core-relations/src/table/mod.rs index ef4d87c28..f710136f7 100644 --- a/core-relations/src/table/mod.rs +++ b/core-relations/src/table/mod.rs @@ -30,7 +30,7 @@ use crate::{ parallel, parallel_heuristics::parallelize_table_op, pool::with_pool_set, - row_buffer::{ParallelRowBufWriter, RowBuffer}, + row_buffer::{HashPartitioning, ParallelRowBufWriter, PartitionedRowBuffer, RowBuffer}, table_spec::{ ColumnId, Constraint, Generation, MutationBuffer, Offset, Row, Table, TableSpec, TableVersion, @@ -49,6 +49,13 @@ mod tests; type HashCode = u64; +// The lowest bits select buckets within a cache-sized destination-table +// window. The next bits group writes to those windows before a shard owner +// probes its HashTable. +const MUTATION_CACHE_WINDOW_BITS: u32 = 11; +const MUTATION_PARTITION_BITS: u32 = 6; +const PARALLEL_INSERT_BATCH_SIZE: usize = 1 << 12; + /// A pointer to a row in the table. #[derive(Clone, Debug)] pub(crate) struct TableEntry { @@ -66,21 +73,6 @@ impl TableEntry { } } -#[derive(Clone, Copy, Debug)] -struct KnownRemoval { - hashcode: HashCode, - row: RowId, -} - -impl KnownRemoval { - fn hashcode(&self) -> u64 { - #[allow(clippy::unnecessary_cast)] - { - self.hashcode as u64 - } - } -} - /// The core data for a table. /// /// This type is a thin wrapper around `RowBuffer`. The big difference is that @@ -185,92 +177,46 @@ impl Clone for SortedWritesTable { } } -/// A variant of [`RowBuffer`] that can handle arity 0. -/// -/// We use this to handle empty keys, where the deletion API needs to handle "row buffers of empty -/// rows". The goal here is to keep most of the API RowBuffer-centric and avoid complicating the -/// code too much: actual code that was optimized to handle arity 0 would look a bit different. -#[derive(Clone)] -enum ArbitraryRowBuffer { - NonEmpty(RowBuffer), - Empty { rows: usize }, -} - -impl ArbitraryRowBuffer { - fn new(arity: usize) -> ArbitraryRowBuffer { - if arity == 0 { - ArbitraryRowBuffer::Empty { rows: 0 } - } else { - ArbitraryRowBuffer::NonEmpty(RowBuffer::new(arity)) - } - } - - fn add_row(&mut self, row: &[Value]) { - match self { - ArbitraryRowBuffer::NonEmpty(buf) => { - buf.add_row(row); - } - ArbitraryRowBuffer::Empty { rows } => { - *rows += 1; - } - } - } - - fn len(&self) -> usize { - match self { - ArbitraryRowBuffer::NonEmpty(buf) => buf.len(), - ArbitraryRowBuffer::Empty { rows } => *rows, - } - } - - fn for_each(&self, mut f: impl FnMut(&[Value])) { - match self { - ArbitraryRowBuffer::NonEmpty(buf) => { - for row in buf.iter() { - f(row); - } - } - ArbitraryRowBuffer::Empty { rows } => { - for _ in 0..*rows { - f(&[]); - } - } - } - } -} - struct Buffer { - pending_rows: DenseIdMap, - pending_removals: DenseIdMap, - pending_known_removals: DenseIdMap>, + pending_rows: Option, + pending_removals: Option, + pending_known_removals: Option, state: Weak, n_cols: u32, n_keys: u32, shard_data: ShardData, + insert_partitioning: HashPartitioning, + removal_partitioning: HashPartitioning, } impl MutationBuffer for Buffer { fn stage_insert(&mut self, row: &[Value]) { - let (shard, _) = hash_code(self.shard_data, row, self.n_keys as _); + let (shard, hash) = hash_code(self.shard_data, row, self.n_keys as _); self.pending_rows - .get_or_insert(shard, || RowBuffer::new(self.n_cols as _)) - .add_row(row); + .get_or_insert_with(|| { + PartitionedRowBuffer::new(self.n_cols as _, self.insert_partitioning) + }) + .add_row(shard.index(), hash, row); } fn stage_remove(&mut self, key: &[Value]) { - let (shard, _) = hash_code(self.shard_data, key, self.n_keys as _); + let (shard, hash) = hash_code(self.shard_data, key, self.n_keys as _); self.pending_removals - .get_or_insert(shard, || ArbitraryRowBuffer::new(self.n_keys as _)) - .add_row(key); + .get_or_insert_with(|| { + PartitionedRowBuffer::new(self.n_keys as _, self.removal_partitioning) + }) + .add_row(shard.index(), hash, key); } fn fresh_handle(&self) -> Box { Box::new(Buffer { - pending_rows: Default::default(), - pending_removals: Default::default(), - pending_known_removals: Default::default(), + pending_rows: None, + pending_removals: None, + pending_known_removals: None, state: self.state.clone(), n_cols: self.n_cols, n_keys: self.n_keys, shard_data: self.shard_data, + insert_partitioning: self.insert_partitioning, + removal_partitioning: self.removal_partitioning, }) } } @@ -279,51 +225,53 @@ impl Buffer { fn stage_remove_row(&mut self, row: RowId, key: &[Value]) { let (shard, hashcode) = hash_code(self.shard_data, key, self.n_keys as _); self.pending_known_removals - .get_or_insert(shard, Vec::new) - .push(KnownRemoval { - hashcode: hashcode as _, - row, - }); + .get_or_insert_with(|| PartitionedRowBuffer::new(1, self.removal_partitioning)) + .add_row(shard.index(), hashcode, &[Value::new(row.rep())]); } } impl Drop for Buffer { fn drop(&mut self) { if let Some(state) = self.state.upgrade() { - let mut rows = 0; - for shard_id in 0..self.pending_rows.n_ids() { - let shard = ShardId::from_usize(shard_id); - let Some(buf) = self.pending_rows.take(shard) else { - continue; - }; - rows += buf.len(); - state.pending_rows[shard].push(buf); - } + let rows = publish_partitioned( + self.pending_rows.take(), + &state.pending_rows, + self.shard_data, + ); state.total_rows.fetch_add(rows, Ordering::Relaxed); - let mut removals = 0; - for shard_id in 0..self.pending_removals.n_ids() { - let shard = ShardId::from_usize(shard_id); - let Some(buf) = self.pending_removals.take(shard) else { - continue; - }; - removals += buf.len(); - state.pending_removals[shard].push(buf); - } - - for shard_id in 0..self.pending_known_removals.n_ids() { - let shard = ShardId::from_usize(shard_id); - let Some(buf) = self.pending_known_removals.take(shard) else { - continue; - }; - removals += buf.len(); - state.pending_known_removals[shard].push(buf); - } + let removals = publish_partitioned( + self.pending_removals.take(), + &state.pending_removals, + self.shard_data, + ) + publish_partitioned( + self.pending_known_removals.take(), + &state.pending_known_removals, + self.shard_data, + ); state.total_removals.fetch_add(removals, Ordering::Relaxed); } } } +fn publish_partitioned( + buffer: Option, + queues: &DenseIdMap>>, + shard_data: ShardData, +) -> usize { + let Some(buffer) = buffer else { + return 0; + }; + let len = buffer.len(); + let buffer = Arc::new(buffer.seal()); + for shard_index in 0..shard_data.n_shards() { + if !buffer.group_is_empty(shard_index) { + queues[ShardId::from_usize(shard_index)].push(buffer.clone()); + } + } + len +} + impl Table for SortedWritesTable { fn dyn_clone(&self) -> Box { Box::new(self.clone()) @@ -603,15 +551,16 @@ impl Table for SortedWritesTable { impl SortedWritesTable { fn new_table_buffer(&self) -> Buffer { - let n_shards = self.hash.shard_data().n_shards(); Buffer { - pending_rows: DenseIdMap::with_capacity(n_shards), - pending_removals: DenseIdMap::with_capacity(n_shards), - pending_known_removals: DenseIdMap::with_capacity(n_shards), + pending_rows: None, + pending_removals: None, + pending_known_removals: None, state: Arc::downgrade(&self.pending_state), n_keys: u32::try_from(self.n_keys).expect("n_keys should fit in u32"), n_cols: u32::try_from(self.n_columns).expect("n_columns should fit in u32"), shard_data: self.hash.shard_data(), + insert_partitioning: self.pending_state.insert_partitioning, + removal_partitioning: self.pending_state.removal_partitioning, } } @@ -644,7 +593,7 @@ impl SortedWritesTable { n_columns, sort_by, offsets: Default::default(), - pending_state: Arc::new(PendingState::new(shard_data)), + pending_state: Arc::new(PendingState::new(shard_data, sort_by.is_some())), merge: merge_fn.into(), to_rebuild, rebuild_index, @@ -659,6 +608,8 @@ impl SortedWritesTable { let shard_data = self.hash.shard_data(); let pending_removals = &self.pending_state.pending_removals; let pending_known_removals = &self.pending_state.pending_known_removals; + let partitioning = self.pending_state.removal_partitioning; + let partition_count = partitioning.partitions_per_group(); let data = &self.data.data; let n_keys = self.n_keys; let stale_delta: usize = @@ -666,26 +617,42 @@ impl SortedWritesTable { let mut removed_rows = with_pool_set(|ps| ps.get::>()); for (shard_offset, shard) in shards.iter_mut().enumerate() { let shard_id = ShardId::from_usize(base_shard + shard_offset); - while let Some(buf) = pending_removals[shard_id].pop() { - buf.for_each(|to_remove| { - let (actual_shard, hc) = hash_code(shard_data, to_remove, n_keys); - assert_eq!(actual_shard, shard_id); - if let Ok(entry) = shard.find_entry(hc, |entry| { - entry.hashcode == (hc as HashCode) - && &data.get_row(entry.row)[0..n_keys] == to_remove - }) { - let (ent, _) = entry.remove(); - removed_rows.push(ent.row); + let removal_buffers = drain_queue(&pending_removals[shard_id]); + for partition in 0..partition_count { + for buf in &removal_buffers { + for hashed in buf.partition( + partitioning.partition_index(shard_id.index(), partition), + ) { + let to_remove = hashed.row; + let hc = hashed.hash; + debug_assert_eq!(shard_data.shard_id(hc), shard_id); + debug_assert_eq!(buf.arity(), n_keys); + if let Ok(entry) = shard.find_entry(hc, |entry| { + entry.hashcode == (hc as HashCode) + && &data.get_row(entry.row)[0..n_keys] == to_remove + }) { + let (ent, _) = entry.remove(); + removed_rows.push(ent.row); + } } - }); + } } - while let Some(buf) = pending_known_removals[shard_id].pop() { - for removal in buf { - if let Ok(entry) = shard.find_entry(removal.hashcode(), |entry| { - entry.hashcode == removal.hashcode && entry.row == removal.row - }) { - let (ent, _) = entry.remove(); - removed_rows.push(ent.row); + let known_buffers = drain_queue(&pending_known_removals[shard_id]); + for partition in 0..partition_count { + for buf in &known_buffers { + for hashed in buf.partition( + partitioning.partition_index(shard_id.index(), partition), + ) { + let hc = hashed.hash; + let row = RowId::new(hashed.row[0].rep()); + debug_assert_eq!(shard_data.shard_id(hc), shard_id); + debug_assert_eq!(buf.arity(), 1); + if let Ok(entry) = shard.find_entry(hc, |entry| { + entry.hashcode == hc as HashCode && entry.row == row + }) { + let (ent, _) = entry.remove(); + removed_rows.push(ent.row); + } } } } @@ -703,6 +670,8 @@ impl SortedWritesTable { } fn serial_delete(&mut self) -> bool { let shard_data = self.hash.shard_data(); + let partitioning = self.pending_state.removal_partitioning; + let partition_count = partitioning.partitions_per_group(); let mut changed = false; self.hash .mut_shards() @@ -710,29 +679,46 @@ impl SortedWritesTable { .enumerate() .for_each(|(shard_id, shard)| { let shard_id = ShardId::from_usize(shard_id); - while let Some(buf) = self.pending_state.pending_removals[shard_id].pop() { - buf.for_each(|to_remove| { - let (actual_shard, hc) = hash_code(shard_data, to_remove, self.n_keys); - assert_eq!(actual_shard, shard_id); - if let Ok(entry) = shard.find_entry(hc, |entry| { - entry.hashcode == (hc as HashCode) - && &self.data.get_row(entry.row).unwrap()[0..self.n_keys] - == to_remove - }) { - let (ent, _) = entry.remove(); - self.data.set_stale(ent.row); - changed = true; + let removal_buffers = drain_queue(&self.pending_state.pending_removals[shard_id]); + for partition in 0..partition_count { + for buf in &removal_buffers { + for hashed in + buf.partition(partitioning.partition_index(shard_id.index(), partition)) + { + let to_remove = hashed.row; + let hc = hashed.hash; + debug_assert_eq!(shard_data.shard_id(hc), shard_id); + debug_assert_eq!(buf.arity(), self.n_keys); + if let Ok(entry) = shard.find_entry(hc, |entry| { + entry.hashcode == (hc as HashCode) + && &self.data.get_row(entry.row).unwrap()[0..self.n_keys] + == to_remove + }) { + let (ent, _) = entry.remove(); + self.data.set_stale(ent.row); + changed = true; + } } - }) + } } - while let Some(buf) = self.pending_state.pending_known_removals[shard_id].pop() { - for removal in buf { - if let Ok(entry) = shard.find_entry(removal.hashcode(), |entry| { - entry.hashcode == removal.hashcode && entry.row == removal.row - }) { - let (ent, _) = entry.remove(); - self.data.set_stale(ent.row); - changed = true; + let known_buffers = + drain_queue(&self.pending_state.pending_known_removals[shard_id]); + for partition in 0..partition_count { + for buf in &known_buffers { + for hashed in + buf.partition(partitioning.partition_index(shard_id.index(), partition)) + { + let hc = hashed.hash; + let row = RowId::new(hashed.row[0].rep()); + debug_assert_eq!(shard_data.shard_id(hc), shard_id); + debug_assert_eq!(buf.arity(), 1); + if let Ok(entry) = shard.find_entry(hc, |entry| { + entry.hashcode == hc as HashCode && entry.row == row + }) { + let (ent, _) = entry.remove(); + self.data.set_stale(ent.row); + changed = true; + } } } } @@ -774,117 +760,73 @@ impl SortedWritesTable { fn serial_insert(&mut self, exec_state: &mut ExecutionState) -> bool { let mut changed = false; let n_keys = self.n_keys; + let partitioning = self.pending_state.insert_partitioning; + let partition_count = partitioning.partitions_per_group(); let mut scratch = with_pool_set(|ps| ps.get::>()); - for (_outer_shard, queue) in self.pending_state.pending_rows.iter() { - if let Some(sort_by) = self.sort_by { - while let Some(buf) = queue.pop() { - for query in buf.non_stale() { + for (outer_shard, queue) in self.pending_state.pending_rows.iter() { + let buffers = drain_queue(queue); + let incoming = buffers + .iter() + .map(|buffer| buffer.group_len(outer_shard.index())) + .sum(); + self.hash.mut_shards()[outer_shard.index()].reserve(incoming, TableEntry::hashcode); + for partition in 0..partition_count { + for buf in &buffers { + for hashed in + buf.partition(partitioning.partition_index(outer_shard.index(), partition)) + { + let query = hashed.row; + if query.first().is_some_and(Value::is_stale) { + continue; + } + let hc = hashed.hash; + debug_assert_eq!(self.hash.shard_data().shard_id(hc), outer_shard); let key = &query[0..n_keys]; - let entry = get_entry_mut(query, n_keys, &mut self.hash, |row| { - let Some(row) = self.data.get_row(row) else { - return false; - }; - &row[0..n_keys] == key - }); - - if let Some(row) = entry { - // First case: overwriting an existing value. Apply merge - // function. Insert new row and update hash table if merge - // changes anything. - let cur = self - .data - .get_row(*row) - .expect("table should not point to stale entry"); - if (self.merge)(exec_state, cur, query, &mut scratch) { - let sort_val = query[sort_by.index()]; - let new = self.data.add_row(&scratch); - if let Some(largest) = self.offsets.last().map(|(v, _)| *v) { - assert!( - sort_val >= largest, - "inserting row that violates sort order ({sort_val:?} vs. {largest:?})" - ); - if sort_val > largest { - self.offsets.push((sort_val, new)); + use hashbrown::hash_table::Entry; + match self.hash.mut_shards()[outer_shard.index()].entry( + hc, + |entry| { + entry.hashcode == hc as HashCode + && self + .data + .get_row(entry.row) + .is_some_and(|row| &row[0..n_keys] == key) + }, + TableEntry::hashcode, + ) { + Entry::Occupied(mut entry) => { + let old_row = entry.get().row; + let cur = self + .data + .get_row(old_row) + .expect("table should not point to stale entry"); + if (self.merge)(exec_state, cur, query, &mut scratch) { + debug_assert_eq!(&scratch[0..n_keys], key); + let new = self.data.add_row(&scratch); + if let Some(sort_by) = self.sort_by { + update_sort_offsets(sort_by, query, new, &mut self.offsets); } - } else { - self.offsets.push((sort_val, new)); + self.data.set_stale(old_row); + entry.get_mut().row = new; + changed = true; } - self.data.set_stale(*row); - *row = new; - changed = true; + scratch.clear(); } - scratch.clear(); - } else { - let sort_val = query[sort_by.index()]; - // New value: update invariants. - let new = self.data.add_row(query); - if let Some(largest) = self.offsets.last().map(|(v, _)| *v) { - assert!( - sort_val >= largest, - "inserting row that violates sort order {sort_val:?} vs. {largest:?}" - ); - if sort_val > largest { - self.offsets.push((sort_val, new)); + Entry::Vacant(entry) => { + let new = self.data.add_row(query); + if let Some(sort_by) = self.sort_by { + update_sort_offsets(sort_by, query, new, &mut self.offsets); } - } else { - self.offsets.push((sort_val, new)); - } - let (shard, hc) = hash_code(self.hash.shard_data(), query, self.n_keys); - debug_assert_eq!(shard, _outer_shard); - self.hash.mut_shards()[shard.index()].insert_unique( - hc as _, - TableEntry { + entry.insert(TableEntry { hashcode: hc as _, row: new, - }, - TableEntry::hashcode, - ); - changed = true; - } - } - } - } else { - // Simplified variant without the sorting constraint. - while let Some(buf) = queue.pop() { - for query in buf.non_stale() { - let key = &query[0..n_keys]; - let entry = get_entry_mut(query, n_keys, &mut self.hash, |row| { - let Some(row) = self.data.get_row(row) else { - return false; - }; - &row[0..n_keys] == key - }); - - if let Some(row) = entry { - let cur = self - .data - .get_row(*row) - .expect("table should not point to stale entry"); - if (self.merge)(exec_state, cur, query, &mut scratch) { - let new = self.data.add_row(&scratch); - self.data.set_stale(*row); - *row = new; + }); changed = true; } - scratch.clear(); - } else { - // New value: update invariants. - let new = self.data.add_row(query); - let (shard, hc) = hash_code(self.hash.shard_data(), query, self.n_keys); - debug_assert_eq!(shard, _outer_shard); - self.hash.mut_shards()[shard.index()].insert_unique( - hc as _, - TableEntry { - hashcode: hc as _, - row: new, - }, - TableEntry::hashcode, - ); - changed = true; } } } - }; + } } changed } @@ -894,7 +836,6 @@ impl SortedWritesTable { exec_state: &ExecutionState, checker: C, ) -> bool { - const BATCH_SIZE: usize = 1 << 18; // Parallel insert uses one giant parallel foreach. We have updates // pre-sharded, and one logical thread can process updates for each // shard independently. Updates happen in three phases, which comments @@ -905,6 +846,8 @@ impl SortedWritesTable { let next_offset = RowId::from_usize(self.data.data.len()); let row_writer = self.data.data.parallel_writer(); let pending_rows = &self.pending_state.pending_rows; + let partitioning = self.pending_state.insert_partitioning; + let partition_count = partitioning.partitions_per_group(); let merge = self.merge.clone(); let pending_adds = parallel::map_mut(self.hash.mut_shards(), |shard_id, shard| { let shard_id = ShardId::from_usize(shard_id); @@ -912,13 +855,20 @@ impl SortedWritesTable { let mut exec_state = exec_state.clone(); let mut scratch = with_pool_set(|ps| ps.get::>()); let queue = &pending_rows[shard_id]; + let buffers = drain_queue(queue); + let incoming = buffers + .iter() + .map(|buffer| buffer.group_len(shard_id.index())) + .sum(); + shard.reserve(incoming, TableEntry::hashcode); let mut marked_stale = 0usize; - let mut staged = StagedOutputs::new(n_keys, n_cols, BATCH_SIZE); + let mut staged = CoalescedInsertBatch::new(n_keys, n_cols, PARALLEL_INSERT_BATCH_SIZE); let mut changed = false; - // The core flush loop: We call once `staged` reaches `BATCH_SIZE` or - // when we're done. - macro_rules! flush_staged_outputs { + // Flush after a bounded number of staged physical rows or at a + // destination-table partition boundary. + macro_rules! flush_insert_batch { () => {{ + if !staged.is_empty() { // Phase 2: Write the staged rows to the row writer. This only // works due to the `ParallelRowBufWriter` machinery. let (start_row, stale) = staged.write_output(&row_writer); @@ -933,16 +883,13 @@ impl SortedWritesTable { // contention. let mut cur_row = start_row; let read_handle = row_writer.read_handle(); - for row in staged.rows() { - if row.first().map(Value::is_stale).unwrap_or(false) { + for (hc, row) in staged.rows_hashed() { + if row.first().is_some_and(Value::is_stale) { cur_row = cur_row.inc(); continue; } use hashbrown::hash_table::Entry; - checker.check_local(row); - changed = true; let key = &row[0..n_keys]; - let (_actual_shard, hc) = hash_code(shard_data, row, n_keys); #[cfg(any(debug_assertions, test))] { unsafe { @@ -951,7 +898,7 @@ impl SortedWritesTable { assert_eq!(read_handle.get_row_unchecked(cur_row), row); } } - debug_assert_eq!(_actual_shard, shard_id); + debug_assert_eq!(shard_data.shard_id(hc), shard_id); match shard.entry( hc, // SAFETY: `ent` must point to a valid row @@ -964,19 +911,37 @@ impl SortedWritesTable { Entry::Occupied(mut occ) => { // SAFETY: `occ` must point to a valid row: we only insert valid rows // into the map. - let cur = unsafe { read_handle.get_row_unchecked(occ.get().row) }; + let old_row = occ.get().row; + let merged = { + let cur = + unsafe { read_handle.get_row_unchecked(old_row) }; + (merge)(&mut exec_state, cur, row, &mut scratch) + }; // SAFETY: The safety requirements of - // `set_stale_shared` are that there are no - // concurrent accesses to `row`. We have - // exclusive access to any row whose hash matches this - // shard. - if (merge)(&mut exec_state, cur, row, &mut scratch) { + // `replace_row_shared` and `set_stale_shared` + // are that there are no concurrent accesses + // to either row. `cur_row` is a completed, + // uniquely reserved write that has not yet + // been published, and we have exclusive + // access to every existing row whose hash + // maps to this shard. + if merged { + debug_assert_eq!( + &scratch[0..n_keys], + key, + "merge functions must preserve table keys" + ); + checker.check_local(&scratch); unsafe { - let _was_stale = read_handle.set_stale_shared(occ.get().row); - debug_assert!(!_was_stale); + read_handle.replace_row_shared(cur_row, &scratch); } occ.get_mut().row = cur_row; + unsafe { + let _was_stale = + read_handle.set_stale_shared(old_row); + debug_assert!(!_was_stale); + } changed = true; } else { // Mark the new row as stale: we didn't end up needing it. @@ -989,6 +954,7 @@ impl SortedWritesTable { scratch.clear(); } Entry::Vacant(v) => { + checker.check_local(row); changed = true; v.insert(TableEntry { hashcode: hc as HashCode, @@ -1000,24 +966,31 @@ impl SortedWritesTable { cur_row = cur_row.inc(); } staged.clear(); + } }}; } // Phase 1: process all incoming updates: - // * Add new values to `staged` - // * Removing entries in `shard` and mark them as stale in - // `data` if they will be overwritten. - while let Some(buf) = queue.pop() { - // We create a read_handle once per batch to avoid blocking - // too many threads if someone needs to resize the row - // writer. - for row in buf.non_stale() { - staged.insert(row, |cur, new, out| (merge)(&mut exec_state, cur, new, out)); - if staged.len() >= BATCH_SIZE { - flush_staged_outputs!(); + // * Coalesce rows with the same key in a bounded temporary table. + // * Copy the staged rows to the shared row buffer in one + // reservation. + // * Apply each live row to the exclusively owned destination + // shard in `flush_insert_batch`. + for partition in 0..partition_count { + for buf in &buffers { + for hashed in + buf.partition(partitioning.partition_index(shard_id.index(), partition)) + { + staged.insert_hashed(hashed.hash, hashed.row, |cur, new, out| { + (merge)(&mut exec_state, cur, new, out) + }); + if staged.physical_len() >= PARALLEL_INSERT_BATCH_SIZE { + flush_insert_batch!(); + } } } + // Do not carry rows into the next destination-table window. + flush_insert_batch!(); } - flush_staged_outputs!(); (checker, marked_stale, changed) }); self.data.data = row_writer.finish(); @@ -1113,17 +1086,11 @@ impl SortedWritesTable { // Parallel rehashes go "hash-first" rather than "rows-first". // // We iterate over each shard and then write out new contents to a fresh row, in parallel. + self.generation = self.generation.inc(); let Some(sort_by) = self.sort_by else { - // Just do a serial rehash for now. We currently do not have a use-case for parallel - // compaction of unsorted tables. - // - // Implementing parallel compaction for an unsorted table is much easier: each shard - // can write to a contiguous chunk of the `scratch` buffer, with the offsets being - // pre-chunked based on the size of each shard. - self.rehash(); + self.parallel_rehash_unsorted(); return; }; - self.generation = self.generation.inc(); assert!(!self.offsets.is_empty()); struct TimestampStats { value: Value, @@ -1255,6 +1222,86 @@ impl SortedWritesTable { mem::swap(&mut self.data.data, &mut self.data.scratch); self.data.stale_rows = 0; } + + /// Compact an unsorted table by walking its live hash entries rather than + /// scanning the row store, which may be mostly stale. + /// + /// Every live row has exactly one entry in exactly one hash shard. Shard + /// lengths therefore give both the live-row count and a cheap histogram + /// for assigning each worker a disjoint output range. + fn parallel_rehash_unsorted(&mut self) { + debug_assert!(self.sort_by.is_none()); + debug_assert!(self.offsets.is_empty()); + + let shards = self.hash.mut_shards(); + let mut shard_offsets = Vec::with_capacity(shards.len() + 1); + shard_offsets.push(0usize); + for shard in shards.iter() { + let next = shard_offsets + .last() + .copied() + .unwrap() + .checked_add(shard.len()) + .expect("live row count should fit in usize"); + shard_offsets.push(next); + } + let live_rows = *shard_offsets.last().unwrap(); + debug_assert_eq!( + live_rows, + self.data.data.len() - self.data.stale_rows, + "each live row should have exactly one hash entry" + ); + + self.data.scratch.clear(); + self.data.scratch.reserve(live_rows); + let scratch_ptr = self.data.scratch.raw_rows() as usize; + let data = &self.data.data; + let n_columns = self.n_columns; + + // Workers mutate disjoint hash shards and copy into the corresponding + // disjoint half-open ranges in `scratch`. The raw pointer is passed as + // an integer because raw pointers are not Sync. + parallel::for_each_mut(shards, |shard_index, shard| { + let scratch_ptr = scratch_ptr as *mut Value; + let mut next = shard_offsets[shard_index]; + let end = shard_offsets[shard_index + 1]; + for TableEntry { row: row_id, .. } in shard.iter_mut() { + let row = data.get_row(*row_id); + debug_assert!( + !row[0].is_stale(), + "a hash entry should never point to a stale row" + ); + let new_row = RowId::from_usize(next); + + // SAFETY: + // * `scratch` was reserved once before any worker started and + // is not resized until all workers have joined. + // * The prefix offsets give each shard a disjoint destination + // range large enough for all of its entries. + // * `data` and `scratch` are distinct row buffers, so source + // and destination cannot overlap. + // * No references into the uninitialized part of `scratch` + // are created before `set_len` below. + unsafe { + std::ptr::copy_nonoverlapping( + row.as_ptr(), + scratch_ptr.add(next * n_columns), + n_columns, + ); + } + *row_id = new_row; + next += 1; + } + debug_assert_eq!(next, end); + }); + + // SAFETY: every row up to `live_rows` was initialized exactly once by + // the shard workers above, and the parallel scope has joined. + unsafe { self.data.scratch.set_len(live_rows) }; + mem::swap(&mut self.data.data, &mut self.data.scratch); + self.data.stale_rows = 0; + } + fn rehash_impl( sort_by: Option, n_keys: usize, @@ -1338,6 +1385,46 @@ fn hash_code(shard_data: ShardData, row: &[Value], n_keys: usize) -> (ShardId, u (shard_data.shard_id(full_code), full_code as HashCode as u64) } +fn update_sort_offsets( + sort_by: ColumnId, + row: &[Value], + row_id: RowId, + offsets: &mut Vec<(Value, RowId)>, +) { + let sort_val = row[sort_by.index()]; + if let Some(largest) = offsets.last().map(|(value, _)| *value) { + assert!( + sort_val >= largest, + "inserting row that violates sort order ({sort_val:?} vs. {largest:?})" + ); + if sort_val > largest { + offsets.push((sort_val, row_id)); + } + } else { + offsets.push((sort_val, row_id)); + } +} + +fn mutation_partitioning(shard_data: ShardData, ordered: bool) -> HashPartitioning { + if ordered || shard_data.n_shards() == 1 { + HashPartitioning::grouped(0, 0, shard_data.n_shards()) + } else { + HashPartitioning::grouped( + MUTATION_CACHE_WINDOW_BITS, + MUTATION_PARTITION_BITS, + shard_data.n_shards(), + ) + } +} + +fn drain_queue(queue: &SegQueue) -> Vec { + let mut drained = Vec::new(); + while let Some(item) = queue.pop() { + drained.push(item); + } + drained +} + /// # Safety /// The caller must ensure exclusive access to every row in `rows`. unsafe fn mark_rows_stale(data: &RowBuffer, rows: &[RowId]) { @@ -1349,16 +1436,20 @@ unsafe fn mark_rows_stale(data: &RowBuffer, rows: &[RowId]) { /// A simple struct for packaging up pending mutations to a `SortedWritesTable`. struct PendingState { - pending_rows: DenseIdMap>, - pending_removals: DenseIdMap>, - pending_known_removals: DenseIdMap>>, + pending_rows: DenseIdMap>>, + pending_removals: DenseIdMap>>, + pending_known_removals: DenseIdMap>>, total_removals: AtomicUsize, total_rows: AtomicUsize, + insert_partitioning: HashPartitioning, + removal_partitioning: HashPartitioning, } impl PendingState { - fn new(shard_data: ShardData) -> PendingState { + fn new(shard_data: ShardData, ordered: bool) -> PendingState { let n_shards = shard_data.n_shards(); + let insert_partitioning = mutation_partitioning(shard_data, ordered); + let removal_partitioning = mutation_partitioning(shard_data, false); let mut pending_rows = DenseIdMap::with_capacity(n_shards); let mut pending_removals = DenseIdMap::with_capacity(n_shards); let mut pending_known_removals = DenseIdMap::with_capacity(n_shards); @@ -1374,6 +1465,8 @@ impl PendingState { pending_known_removals, total_removals: AtomicUsize::new(0), total_rows: AtomicUsize::new(0), + insert_partitioning, + removal_partitioning, } } fn clear(&self) { @@ -1388,6 +1481,8 @@ impl PendingState { for (_, queue) in self.pending_known_removals.iter() { while queue.pop().is_some() {} } + self.total_rows.store(0, Ordering::Relaxed); + self.total_removals.store(0, Ordering::Relaxed); } /// This is only really used in debugging, but it's annoying enough to write @@ -1442,6 +1537,8 @@ impl PendingState { pending_known_removals, total_removals: AtomicUsize::new(self.total_removals.load(Ordering::Acquire)), total_rows: AtomicUsize::new(self.total_rows.load(Ordering::Acquire)), + insert_partitioning: self.insert_partitioning, + removal_partitioning: self.removal_partitioning, } } } @@ -1541,87 +1638,127 @@ impl OrderingChecker for SortChecker { } } -/// A type similar to a SortedWritesTable used to buffer outputs. The main thing -/// that StagedOutputs handles is running the merge function for a table on -/// multiple updates to the same key that show up in the same round of -/// insertions. -struct StagedOutputs { - shard_data: ShardData, +/// A bounded, cached-hash coalescer for pending insertion rows. +/// +/// Multiple rows with the same key are merged before the surviving value is +/// combined with the destination table's current value. Physical rows and +/// their unchanged full hashes stay aligned, including stale intermediate +/// merge outputs. +struct CoalescedInsertBatch { n_keys: usize, hash: Pooled>, rows: RowBuffer, + hashes: Vec, n_stale: usize, scratch: Pooled>, } -impl StagedOutputs { - fn rows(&self) -> impl Iterator { - self.rows.iter() +impl CoalescedInsertBatch { + fn rows_hashed(&self) -> impl Iterator { + debug_assert_eq!(self.hashes.len(), self.rows.len()); + self.hashes + .iter() + .copied() + .zip(self.rows.iter()) + .map(|(hash, row)| { + // Keep this conversion at the boundary so HashCode can still + // be changed independently. + #[allow(clippy::unnecessary_cast)] + let hash = hash as u64; + (hash, row) + }) } + fn new(n_keys: usize, n_cols: usize, capacity: usize) -> Self { - let mut res = with_pool_set(|ps| StagedOutputs { - shard_data: ShardData::new(1), + let mut res = with_pool_set(|ps| CoalescedInsertBatch { n_keys, - n_stale: 0, hash: ps.get(), rows: RowBuffer::new(n_cols), + hashes: Vec::with_capacity(capacity), + n_stale: 0, scratch: ps.get(), }); - res.hash.reserve(capacity, TableEntry::hashcode); + res.hash + .reserve(capacity, |entry| coalesced_hash(entry.hashcode())); res.rows.reserve(capacity); res } + fn clear(&mut self) { self.hash.clear(); self.rows.clear(); + self.hashes.clear(); self.n_stale = 0; } - fn len(&self) -> usize { - self.rows.len() - self.n_stale + + fn is_empty(&self) -> bool { + self.rows.len() == 0 + } + + fn physical_len(&self) -> usize { + self.rows.len() } - fn insert( + fn insert_hashed( &mut self, + hc: u64, row: &[Value], mut merge_fn: impl FnMut(&[Value], &[Value], &mut Vec) -> bool, ) { - if row[0].is_stale() { + if row.first().is_some_and(Value::is_stale) { return; } + + debug_assert_eq!(self.hashes.len(), self.rows.len()); use hashbrown::hash_table::Entry; - let (_, hc) = hash_code(self.shard_data, row, self.n_keys); let entry = self.hash.entry( - hc, - |te| { - te.hashcode() == hc - && self.rows.get_row(te.row)[0..self.n_keys] == row[0..self.n_keys] + coalesced_hash(hc), + |entry| { + entry.hashcode() == hc + && self.rows.get_row(entry.row)[0..self.n_keys] == row[0..self.n_keys] }, - TableEntry::hashcode, + |entry| coalesced_hash(entry.hashcode()), ); match entry { - Entry::Occupied(mut occupied_entry) => { - let cur = self.rows.get_row(occupied_entry.get().row); + Entry::Occupied(mut occupied) => { + let cur = self.rows.get_row(occupied.get().row); if merge_fn(cur, row, &mut self.scratch) { + debug_assert_eq!( + &self.scratch[0..self.n_keys], + &row[0..self.n_keys], + "merge functions must preserve table keys" + ); let new = self.rows.add_row(&self.scratch); - self.rows.set_stale(occupied_entry.get().row); + self.hashes.push(hc as HashCode); + self.rows.set_stale(occupied.get().row); self.n_stale += 1; - occupied_entry.get_mut().row = new; + occupied.get_mut().row = new; } self.scratch.clear(); } - Entry::Vacant(vacant_entry) => { + Entry::Vacant(vacant) => { let next = self.rows.add_row(row); - vacant_entry.insert(TableEntry { - hashcode: hc as _, + self.hashes.push(hc as HashCode); + vacant.insert(TableEntry { + hashcode: hc as HashCode, row: next, }); } } + debug_assert_eq!(self.hashes.len(), self.rows.len()); } - /// Write the contents of the staged outputs to the given writer, returning the initial RowId - /// of the new output. Returns the number of stale values in the buffer that was appended. + /// Append all physical rows, returning the first RowId and stale count. fn write_output(&self, output: &ParallelRowBufWriter) -> (RowId, usize) { + debug_assert_eq!(self.hashes.len(), self.rows.len()); (output.append_contents(&self.rows), self.n_stale) } } + +#[inline] +fn coalesced_hash(hash: u64) -> u64 { + // Cache partitioning fixes a run of middle bits in the full hash. Mix + // higher bits down before probing the temporary table so those fixed bits + // do not artificially restrict its usable buckets. + hash ^ hash.wrapping_shr(17) +} diff --git a/core-relations/src/table/tests.rs b/core-relations/src/table/tests.rs index a96f14e11..ae275f33b 100644 --- a/core-relations/src/table/tests.rs +++ b/core-relations/src/table/tests.rs @@ -5,10 +5,11 @@ use crate::{ common::{HashMap, ShardId, Value}, offsets::{RowId, SubsetRef}, row_buffer::TaggedRowBuffer, - table::{TableEntry, hash_code}, + table::{SortedWritesTable, TableEntry, hash_code}, table_shortcuts::{fill_table, v}, table_spec::{ColumnId, Constraint, Offset, Table, WrappedTable}, }; +use egglog_concurrency::ThreadPool; use super::sharded_hash_table::ShardedHashTable; @@ -180,6 +181,289 @@ fn known_row_removals_match_the_staged_row_id() { assert_eq!(table.len(), 2); } +#[test] +fn pending_partitioned_mutations_survive_table_clone() { + empty_execution_state!(e); + let mut table = fill_table( + vec![vec![v(1), v(10)], vec![v(2), v(20)]], + 1, + None, + |_, new| Some(new.to_vec()), + ); + let mut pending = table.new_buffer(); + pending.stage_remove(&[v(1)]); + pending.stage_insert(&[v(3), v(30)]); + drop(pending); + + let mut cloned = table.clone(); + table.merge(&mut e); + cloned.merge(&mut e); + + for table in [&table, &cloned] { + assert!(table.get_row(&[v(1)]).is_none()); + assert_eq!(&*table.get_row(&[v(2)]).unwrap().vals, &[v(2), v(20)]); + assert_eq!(&*table.get_row(&[v(3)]).unwrap().vals, &[v(3), v(30)]); + } +} + +#[test] +fn cached_hashes_stay_aligned_through_parallel_duplicate_staging() { + const KEYS: usize = 10_000; + let pool = ThreadPool::new(4); + pool.install(|| { + empty_execution_state!(e); + let mut table = SortedWritesTable::new( + 1, + 2, + None, + vec![], + Box::new(|_, old, new, out: &mut Vec| { + out.extend_from_slice(new); + old != new + }), + ); + + for base in (0..KEYS).step_by(257) { + let mut buffer = table.new_buffer(); + for key in base..(base + 257).min(KEYS) { + buffer.stage_insert(&[v(key), v(key)]); + buffer.stage_insert(&[v(key), v(key + 1)]); + } + } + + assert!(table.parallel_insert(&e, ())); + assert_eq!(table.len(), KEYS); + for key in 0..KEYS { + assert_eq!( + &*table.get_row(&[v(key)]).unwrap().vals, + &[v(key), v(key + 1)] + ); + } + }); +} + +#[test] +fn parallel_insert_uses_merge_output_for_existing_row() { + let pool = ThreadPool::new(4); + pool.install(|| { + empty_execution_state!(e); + let mut table = SortedWritesTable::new( + 1, + 2, + None, + vec![], + Box::new(|_, old, new, out: &mut Vec| { + out.extend_from_slice(&[old[0], Value::new(old[1].rep() + new[1].rep())]); + true + }), + ); + + table.new_buffer().stage_insert(&[v(1), v(10)]); + table.merge(&mut e); + table.new_buffer().stage_insert(&[v(1), v(7)]); + + assert!(table.parallel_insert(&e, ())); + assert_eq!(&*table.get_row(&[v(1)]).unwrap().vals, &[v(1), v(17)]); + assert_eq!(table.len(), 1); + }); +} + +#[test] +fn parallel_insert_merges_duplicate_split_across_flushes() { + let pool = ThreadPool::new(1); + pool.install(|| { + empty_execution_state!(e); + let mut table = SortedWritesTable::new( + 1, + 2, + None, + vec![], + Box::new(|_, old, new, out: &mut Vec| { + out.extend_from_slice(&[old[0], Value::new(old[1].rep() + new[1].rep())]); + true + }), + ); + let target = super::PARALLEL_INSERT_BATCH_SIZE + 1; + let mut buffer = table.new_buffer(); + buffer.stage_insert(&[v(target), v(10)]); + for key in 0..super::PARALLEL_INSERT_BATCH_SIZE { + buffer.stage_insert(&[v(key), v(key)]); + } + buffer.stage_insert(&[v(target), v(7)]); + drop(buffer); + + assert!(table.parallel_insert(&e, ())); + assert_eq!( + &*table.get_row(&[v(target)]).unwrap().vals, + &[v(target), v(17)] + ); + assert_eq!(table.len(), super::PARALLEL_INSERT_BATCH_SIZE + 1); + }); +} + +#[test] +fn parallel_insert_reports_no_change_when_merge_rejects_update() { + let pool = ThreadPool::new(4); + pool.install(|| { + empty_execution_state!(e); + let mut table = SortedWritesTable::new(1, 2, None, vec![], Box::new(|_, _, _, _| false)); + + table.new_buffer().stage_insert(&[v(1), v(10)]); + table.merge(&mut e); + table.new_buffer().stage_insert(&[v(1), v(7)]); + + assert!(!table.parallel_insert(&e, ())); + assert_eq!(&*table.get_row(&[v(1)]).unwrap().vals, &[v(1), v(10)]); + assert_eq!(table.len(), 1); + }); +} + +#[test] +fn parallel_insert_preserves_coalesced_merge_parenthesization() { + let pool = ThreadPool::new(4); + pool.install(|| { + empty_execution_state!(e); + let mut table = SortedWritesTable::new( + 1, + 2, + None, + vec![], + Box::new(|_, old, new, out: &mut Vec| { + out.extend_from_slice(&[ + old[0], + Value::new(old[1].rep().checked_sub(new[1].rep()).unwrap()), + ]); + true + }), + ); + + table.new_buffer().stage_insert(&[v(1), v(20)]); + table.merge(&mut e); + + let mut buffer = table.new_buffer(); + buffer.stage_insert(&[v(1), v(7)]); + buffer.stage_insert(&[v(1), v(3)]); + drop(buffer); + + assert!(table.parallel_insert(&e, ())); + // Pending rows are coalesced before consulting the destination: + // 20 - (7 - 3) = 16. Direct left-folding would produce 10. + assert_eq!(&*table.get_row(&[v(1)]).unwrap().vals, &[v(1), v(16)]); + assert_eq!(table.len(), 1); + }); +} + +#[test] +fn parallel_unsorted_rehash_compacts_heavy_stale_rows_and_repoints_entries() { + const KEYS: usize = 70_000; + const UPDATES: usize = 5; + + let pool = ThreadPool::new(4); + pool.install(|| { + empty_execution_state!(e); + let mut table = SortedWritesTable::new( + 1, + 2, + None, + vec![], + Box::new(|_, old, new, out: &mut Vec| { + out.extend_from_slice(new); + old != new + }), + ); + + // Deliberately bypass `merge`'s compaction step so the final call + // crosses the normal parallel-compaction threshold with a row store + // that is overwhelmingly stale. + for update in 0..=UPDATES { + let mut buffer = table.new_buffer(); + for key in 0..KEYS { + buffer.stage_insert(&[v(key), v(update)]); + } + drop(buffer); + assert!(table.do_insert(&mut e)); + } + + assert_eq!(table.data.data.len(), KEYS * (UPDATES + 1)); + assert_eq!(table.data.stale_rows, KEYS * UPDATES); + let old_generation = table.generation; + + table.maybe_rehash(); + + assert_eq!(table.generation.index(), old_generation.index() + 1); + assert_eq!(table.data.data.len(), KEYS); + assert_eq!(table.data.stale_rows, 0); + assert_eq!(table.len(), KEYS); + + let mut seen_rows = vec![false; KEYS]; + let mut seen_keys = vec![false; KEYS]; + let shard_data = table.hash.shard_data(); + for (shard_index, shard) in table.hash.mut_shards().iter().enumerate() { + for entry in shard.iter() { + let row = table.data.data.get_row(entry.row); + assert_eq!(row[1], v(UPDATES)); + assert_eq!(entry.hashcode, hash_code(shard_data, row, 1).1); + assert_eq!(shard_data.shard_id(entry.hashcode).index(), shard_index); + + assert!(entry.row.index() < KEYS); + assert!(!seen_rows[entry.row.index()]); + seen_rows[entry.row.index()] = true; + + let key = row[0].rep() as usize; + assert!(key < KEYS); + assert!(!seen_keys[key]); + seen_keys[key] = true; + } + } + assert!(seen_rows.into_iter().all(|seen| seen)); + assert!(seen_keys.into_iter().all(|seen| seen)); + + // Exercise the repointed hash entries after compaction. + table.new_buffer().stage_insert(&[v(17), v(UPDATES + 1)]); + assert!(table.do_insert(&mut e)); + assert_eq!( + &*table.get_row(&[v(17)]).unwrap().vals, + &[v(17), v(UPDATES + 1)] + ); + }); +} + +#[test] +fn parallel_unsorted_rehash_handles_no_live_rows() { + const KEYS: usize = 1_024; + + let pool = ThreadPool::new(4); + pool.install(|| { + empty_execution_state!(e); + let mut table = SortedWritesTable::new(1, 2, None, vec![], Box::new(|_, _, _, _| false)); + + let mut inserts = table.new_buffer(); + for key in 0..KEYS { + inserts.stage_insert(&[v(key), v(key + 1)]); + } + drop(inserts); + assert!(table.do_insert(&mut e)); + + let mut removals = table.new_buffer(); + for key in 0..KEYS { + removals.stage_remove(&[v(key)]); + } + drop(removals); + assert!(table.do_delete()); + assert_eq!(table.data.stale_rows, KEYS); + assert!(table.hash.mut_shards().iter().all(|shard| shard.is_empty())); + + let old_generation = table.generation; + table.parallel_rehash(); + + assert_eq!(table.generation.index(), old_generation.index() + 1); + assert_eq!(table.data.data.len(), 0); + assert_eq!(table.data.stale_rows, 0); + assert_eq!(table.len(), 0); + assert!(table.hash.mut_shards().iter().all(|shard| shard.is_empty())); + }); +} + #[test] fn shard_math() { let mut table = ShardedHashTable::::with_shards(14); From 808db869ad6f11fdc0e95d17b0506cd9bb22878a Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Sun, 26 Jul 2026 11:23:13 -0700 Subject: [PATCH 10/12] Add math scalability benchmark harness --- scripts/math_scalability.py | 531 ++++++++++++++++++++++++++++++++++++ 1 file changed, 531 insertions(+) create mode 100755 scripts/math_scalability.py diff --git a/scripts/math_scalability.py b/scripts/math_scalability.py new file mode 100755 index 000000000..157230053 --- /dev/null +++ b/scripts/math_scalability.py @@ -0,0 +1,531 @@ +#!/usr/bin/env python3 +"""Benchmark math-microbenchmark across egglog thread counts. + +The harness intentionally always passes ``--no-decomp``. It builds the +release egglog binary by default, performs per-configuration warmups, then +interleaves measured runs to reduce ordering bias. Aggregate results are +written as CSV and raw plus aggregate results are written as JSON. + +Examples: + + python3 scripts/math_scalability.py + python3 scripts/math_scalability.py --threads 1,2,4,8,12 --repetitions 7 + python3 scripts/math_scalability.py --no-build --output-dir /private/tmp +""" + +from __future__ import annotations + +import argparse +import csv +import json +import os +import platform +import statistics +import subprocess +import sys +import time +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Sequence + +try: + import resource +except ImportError: # pragma: no cover - resource is available on Unix/macOS. + resource = None # type: ignore[assignment] + + +SCRIPT_DIR = Path(__file__).resolve().parent +REPO_ROOT = SCRIPT_DIR.parent +DEFAULT_BINARY = REPO_ROOT / "target" / "release" / "egglog" +DEFAULT_BENCHMARK = REPO_ROOT / "tests" / "math-microbenchmark.egg" +DEFAULT_OUTPUT_DIR = REPO_ROOT / "target" / "math-scalability" +DEFAULT_THREADS = (1, 2, 4, 8, 12) + + +def positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("must be greater than zero") + return parsed + + +def nonnegative_int(value: str) -> int: + parsed = int(value) + if parsed < 0: + raise argparse.ArgumentTypeError("must be nonnegative") + return parsed + + +def positive_float(value: str) -> float: + parsed = float(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("must be greater than zero") + return parsed + + +def thread_counts(value: str) -> tuple[int, ...]: + try: + parsed = tuple(int(part.strip()) for part in value.split(",")) + except ValueError as error: + raise argparse.ArgumentTypeError( + "must be a comma-separated list of integers" + ) from error + if not parsed or any(thread <= 0 for thread in parsed): + raise argparse.ArgumentTypeError("thread counts must be greater than zero") + if len(set(parsed)) != len(parsed): + raise argparse.ArgumentTypeError("thread counts must be unique") + return parsed + + +def parse_args(argv: Sequence[str]) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Measure math-microbenchmark scalability. Every benchmark command " + "uses --no-decomp." + ) + ) + parser.add_argument( + "--threads", + type=thread_counts, + default=DEFAULT_THREADS, + metavar="N,N,...", + help="thread counts to measure (default: 1,2,4,8,12)", + ) + parser.add_argument( + "--warmups", + type=nonnegative_int, + default=1, + help="warmup runs per thread count (default: 1)", + ) + parser.add_argument( + "--repetitions", + type=positive_int, + default=5, + help="measured runs per thread count (default: 5)", + ) + parser.add_argument( + "--timeout", + type=positive_float, + default=300.0, + metavar="SECONDS", + help="timeout for each egglog invocation (default: 300)", + ) + parser.add_argument( + "--binary", + type=Path, + default=DEFAULT_BINARY, + help="egglog binary to run (default: target/release/egglog)", + ) + parser.add_argument( + "--benchmark", + type=Path, + default=DEFAULT_BENCHMARK, + help="benchmark input (default: tests/math-microbenchmark.egg)", + ) + parser.add_argument( + "--output-dir", + type=Path, + default=DEFAULT_OUTPUT_DIR, + help="directory for timestamped CSV and JSON results", + ) + parser.add_argument( + "--label", + default="", + help="optional label stored in the output metadata", + ) + parser.add_argument( + "--no-build", + action="store_true", + help="use the existing binary without running cargo build --release", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="print the build and benchmark commands without running them", + ) + return parser.parse_args(argv) + + +def clean_runtime_environment() -> tuple[dict[str, str], list[str]]: + """Remove ambient egglog logging/tuning settings for reproducible runs.""" + environment = os.environ.copy() + removed = sorted( + key + for key in environment + if key == "RUST_LOG" or key.startswith("EGGLOG_") + ) + for key in removed: + del environment[key] + return environment, removed + + +def git_output(*arguments: str) -> str: + try: + return subprocess.check_output( + ["git", "-C", str(REPO_ROOT), *arguments], + stderr=subprocess.DEVNULL, + text=True, + ).strip() + except (FileNotFoundError, subprocess.CalledProcessError): + return "" + + +def benchmark_command( + binary: Path, benchmark: Path, threads: int +) -> list[str]: + # Keep --no-decomp unconditional: results from this harness should remain + # comparable while decomposition is disabled for this workload. + return [ + str(binary), + "--no-decomp", + "-j", + str(threads), + str(benchmark), + ] + + +def display_command(command: Sequence[str]) -> str: + import shlex + + return shlex.join(command) + + +def child_cpu_times() -> tuple[float, float] | None: + """Return cumulative (user, system) CPU seconds for completed children.""" + if resource is None: + return None + try: + usage = resource.getrusage(resource.RUSAGE_CHILDREN) + except (AttributeError, OSError): + return None + return usage.ru_utime, usage.ru_stime + + +def run_once( + command: Sequence[str], + *, + environment: dict[str, str], + timeout: float, +) -> dict[str, float | None]: + cpu_before = child_cpu_times() + started = time.perf_counter_ns() + try: + completed = subprocess.run( + command, + cwd=REPO_ROOT, + env=environment, + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + text=True, + timeout=timeout, + ) + except subprocess.TimeoutExpired as error: + raise RuntimeError( + f"benchmark exceeded the {timeout:g}s timeout: " + f"{display_command(command)}" + ) from error + elapsed_seconds = (time.perf_counter_ns() - started) / 1_000_000_000 + if completed.returncode != 0: + stderr = completed.stderr.strip() + detail = f"\n{stderr}" if stderr else "" + raise RuntimeError( + f"benchmark exited with status {completed.returncode}: " + f"{display_command(command)}{detail}" + ) + cpu_after = child_cpu_times() + user_cpu_seconds: float | None = None + system_cpu_seconds: float | None = None + if cpu_before is not None and cpu_after is not None: + user_cpu_seconds = max(0.0, cpu_after[0] - cpu_before[0]) + system_cpu_seconds = max(0.0, cpu_after[1] - cpu_before[1]) + total_cpu_seconds = ( + user_cpu_seconds + system_cpu_seconds + if user_cpu_seconds is not None and system_cpu_seconds is not None + else None + ) + return { + "elapsed_seconds": elapsed_seconds, + "user_cpu_seconds": user_cpu_seconds, + "system_cpu_seconds": system_cpu_seconds, + "cpu_seconds": total_cpu_seconds, + "effective_utilized_cores": ( + total_cpu_seconds / elapsed_seconds + if total_cpu_seconds is not None + else None + ), + } + + +def summarize( + samples: dict[int, list[dict[str, float | None]]], + ordered_threads: Sequence[int], +) -> list[dict[str, Any]]: + baseline_threads = min(ordered_threads) + baseline_mean = statistics.fmean( + sample["elapsed_seconds"] for sample in samples[baseline_threads] + ) + results: list[dict[str, Any]] = [] + for threads in ordered_threads: + thread_samples = samples[threads] + values = [sample["elapsed_seconds"] for sample in thread_samples] + assert all(value is not None for value in values) + elapsed_values = [float(value) for value in values] + mean = statistics.fmean(elapsed_values) + speedup = baseline_mean / mean + efficiency = speedup * baseline_threads / threads + cpu_available = all( + sample["cpu_seconds"] is not None for sample in thread_samples + ) + mean_user_cpu = ( + statistics.fmean( + float(sample["user_cpu_seconds"]) for sample in thread_samples + ) + if cpu_available + else None + ) + mean_system_cpu = ( + statistics.fmean( + float(sample["system_cpu_seconds"]) for sample in thread_samples + ) + if cpu_available + else None + ) + mean_cpu = ( + statistics.fmean( + float(sample["cpu_seconds"]) for sample in thread_samples + ) + if cpu_available + else None + ) + results.append( + { + "threads": threads, + "repetitions": len(elapsed_values), + "mean_seconds": mean, + "median_seconds": statistics.median(elapsed_values), + "min_seconds": min(elapsed_values), + "max_seconds": max(elapsed_values), + "stddev_seconds": statistics.stdev(elapsed_values) + if len(elapsed_values) > 1 + else 0.0, + "mean_user_cpu_seconds": mean_user_cpu, + "mean_system_cpu_seconds": mean_system_cpu, + "mean_cpu_seconds": mean_cpu, + "effective_utilized_cores": mean_cpu / mean + if mean_cpu is not None + else None, + "speedup_vs_baseline": speedup, + "parallel_efficiency_pct": efficiency * 100.0, + } + ) + return results + + +CSV_COLUMNS = ( + "threads", + "repetitions", + "mean_seconds", + "median_seconds", + "min_seconds", + "max_seconds", + "stddev_seconds", + "mean_user_cpu_seconds", + "mean_system_cpu_seconds", + "mean_cpu_seconds", + "effective_utilized_cores", + "speedup_vs_baseline", + "parallel_efficiency_pct", +) + + +def write_outputs( + output_dir: Path, document: dict[str, Any] +) -> tuple[Path, Path]: + output_dir.mkdir(parents=True, exist_ok=True) + timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ") + stem = f"math-scalability-{timestamp}" + csv_path = output_dir / f"{stem}.csv" + json_path = output_dir / f"{stem}.json" + + with csv_path.open("x", newline="", encoding="utf-8") as output: + writer = csv.DictWriter(output, fieldnames=CSV_COLUMNS) + writer.writeheader() + writer.writerows(document["results"]) + with json_path.open("x", encoding="utf-8") as output: + json.dump(document, output, indent=2, sort_keys=True) + output.write("\n") + return csv_path, json_path + + +def print_summary(results: Sequence[dict[str, Any]], baseline_threads: int) -> None: + print() + print("math-microbenchmark scalability (--no-decomp)") + print( + f"{'threads':>7} {'mean (s)':>10} {'stddev':>10} " + f"{'speedup':>8} {'efficiency':>10} {'CPU cores':>9}" + ) + for result in results: + cpu_cores = result["effective_utilized_cores"] + cpu_cores_display = f"{cpu_cores:.2f}" if cpu_cores is not None else "n/a" + print( + f"{result['threads']:>7d} " + f"{result['mean_seconds']:>10.4f} " + f"{result['stddev_seconds']:>10.4f} " + f"{result['speedup_vs_baseline']:>7.2f}x " + f"{result['parallel_efficiency_pct']:>9.1f}% " + f"{cpu_cores_display:>9}" + ) + print( + f"Speedup and efficiency use the {baseline_threads}-thread mean as " + "the baseline. CPU cores is (child user + system CPU time) / wall time." + ) + + +def main(argv: Sequence[str] | None = None) -> int: + args = parse_args(argv if argv is not None else sys.argv[1:]) + binary = args.binary.expanduser().resolve() + benchmark = args.benchmark.expanduser().resolve() + output_dir = args.output_dir.expanduser().resolve() + environment, removed_environment = clean_runtime_environment() + build_command = [ + "cargo", + "build", + "--release", + "--manifest-path", + str(REPO_ROOT / "Cargo.toml"), + "--bin", + "egglog", + ] + + if args.dry_run: + if not args.no_build: + print(display_command(build_command)) + for threads in args.threads: + print( + f"# {args.warmups} warmup(s), {args.repetitions} measured run(s)" + ) + print(display_command(benchmark_command(binary, benchmark, threads))) + print(f"# output directory: {output_dir}") + return 0 + + if not benchmark.is_file(): + raise RuntimeError(f"benchmark input does not exist: {benchmark}") + if not args.no_build: + print("Building release egglog...", flush=True) + subprocess.run(build_command, cwd=REPO_ROOT, check=True) + if not binary.is_file(): + raise RuntimeError( + f"egglog binary does not exist: {binary} " + "(remove --no-build or pass --binary)" + ) + + print( + "Benchmarking " + + ", ".join(f"{threads}t" for threads in args.threads) + + f" with {args.warmups} warmup(s) and " + + f"{args.repetitions} measured run(s) each.", + flush=True, + ) + if removed_environment: + print( + "Ignoring ambient runtime settings: " + + ", ".join(removed_environment), + flush=True, + ) + + for threads in args.threads: + command = benchmark_command(binary, benchmark, threads) + for warmup in range(args.warmups): + print( + f" warmup {warmup + 1}/{args.warmups}: {threads} threads", + flush=True, + ) + run_once(command, environment=environment, timeout=args.timeout) + + samples = {threads: [] for threads in args.threads} + raw_samples: list[dict[str, Any]] = [] + # Alternate traversal direction each round to reduce monotonic thermal and + # background-load bias without adding random, irreproducible ordering. + for repetition in range(args.repetitions): + order = args.threads if repetition % 2 == 0 else tuple(reversed(args.threads)) + for threads in order: + sample = run_once( + benchmark_command(binary, benchmark, threads), + environment=environment, + timeout=args.timeout, + ) + samples[threads].append(sample) + raw_samples.append( + { + "threads": threads, + "repetition": repetition + 1, + **sample, + } + ) + elapsed = sample["elapsed_seconds"] + cpu_cores = sample["effective_utilized_cores"] + cpu_display = ( + f", {cpu_cores:.2f} effective CPU cores" + if cpu_cores is not None + else "" + ) + print( + f" run {repetition + 1}/{args.repetitions}: " + f"{threads:>2} threads {elapsed:.4f}s{cpu_display}", + flush=True, + ) + + results = summarize(samples, args.threads) + metadata = { + "schema_version": 1, + "created_at_utc": datetime.now(timezone.utc).isoformat(), + "label": args.label, + "benchmark": str(benchmark), + "binary": str(binary), + "no_decomp": True, + "threads": list(args.threads), + "warmups": args.warmups, + "repetitions": args.repetitions, + "timeout_seconds": args.timeout, + "command_template": [ + str(binary), + "--no-decomp", + "-j", + "{threads}", + str(benchmark), + ], + "git": { + "commit": git_output("rev-parse", "HEAD"), + "branch": git_output("rev-parse", "--abbrev-ref", "HEAD"), + "dirty": bool(git_output("status", "--porcelain")), + }, + "machine": { + "platform": platform.platform(), + "machine": platform.machine(), + "logical_cpu_count": os.cpu_count(), + }, + "cpu_time_available": child_cpu_times() is not None, + "effective_utilized_cores_definition": ( + "(child user CPU seconds + child system CPU seconds) / wall seconds" + ), + "ignored_environment_variables": removed_environment, + } + document = { + "metadata": metadata, + "samples": raw_samples, + "results": results, + } + csv_path, json_path = write_outputs(output_dir, document) + + print_summary(results, min(args.threads)) + print(f"CSV: {csv_path}") + print(f"JSON: {json_path}") + return 0 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except (OSError, RuntimeError, subprocess.CalledProcessError) as error: + print(f"error: {error}", file=sys.stderr) + raise SystemExit(1) From ddc35f0f3a027ee2e42baddd51f9a4fb96674a59 Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Sun, 26 Jul 2026 16:00:08 -0700 Subject: [PATCH 11/12] Improve parallel union-find and join scalability --- core-relations/src/free_join/execute.rs | 482 +++++++++++++++++------- core-relations/src/tests.rs | 282 +++++++++++++- core-relations/src/uf/mod.rs | 198 +++++++++- core-relations/src/uf/tests.rs | 118 ++++++ numeric-id/src/lib.rs | 21 ++ union-find/src/concurrent/atomic_int.rs | 31 +- union-find/src/concurrent/uf.rs | 163 +++++--- union-find/src/lib.rs | 96 ++++- union-find/src/tests.rs | 23 ++ 9 files changed, 1223 insertions(+), 191 deletions(-) diff --git a/core-relations/src/free_join/execute.rs b/core-relations/src/free_join/execute.rs index 92c4fe99e..7df097155 100644 --- a/core-relations/src/free_join/execute.rs +++ b/core-relations/src/free_join/execute.rs @@ -359,15 +359,103 @@ fn top_index_shape_is_eligible( && nonempty_shards >= workers } -/// One physical cached-index shard used to drive the top join level. +const TOP_COVER_PARTITIONS_PER_WORKER: usize = 4; +const MIN_TOP_COVER_ROWS_PER_PARTITION: usize = 256; +const TOP_INDEX_RANGES_PER_WORKER: usize = 32; + +/// Split a cover subset into bounded ordinal ranges without materializing its +/// row IDs. Keeping several ranges per worker gives work stealing room to +/// absorb join-selectivity skew while the minimum size keeps handoff coarse. +fn top_cover_partitions( + cover_rows: usize, + workers: usize, + min_rows_per_partition: usize, +) -> Option> { + if workers <= 1 || cover_rows < min_rows_per_partition.saturating_mul(workers) { + return None; + } + + let target_partitions = workers.saturating_mul(TOP_COVER_PARTITIONS_PER_WORKER); + let rows_per_partition = cover_rows + .div_ceil(target_partitions) + .max(min_rows_per_partition); + let mut partitions = Vec::with_capacity(cover_rows.div_ceil(rows_per_partition)); + for start in (0..cover_rows).step_by(rows_per_partition) { + partitions.push(TopLevelPartition::CoverRange { + start, + scan_size: cmp::min(rows_per_partition, cover_rows - start), + }); + } + Some(partitions) +} + +/// Split an unsharded scalar index into disjoint key-ordinal ranges. +/// +/// Unlike scan-and-project partitioning, this borrows the scalar index that +/// normal join execution has already built for the filtered root. It therefore +/// adds no row-ID copy or second sort for seminaive and constrained subsets. +/// The aggressive overpartitioning leaves enough independent work to absorb +/// skew in the number of join results produced by each leading key. +fn top_index_range_partitions( + leader_keys: usize, + workers: usize, + min_keys_per_worker: usize, +) -> Option> { + if workers <= 1 || leader_keys < min_keys_per_worker.saturating_mul(workers) { + return None; + } + + let target_partitions = workers.saturating_mul(TOP_INDEX_RANGES_PER_WORKER); + let keys_per_partition = leader_keys + .div_ceil(target_partitions) + .max(min_keys_per_worker); + let mut partitions = Vec::with_capacity(leader_keys.div_ceil(keys_per_partition)); + for start in (0..leader_keys).step_by(keys_per_partition) { + partitions.push(TopLevelPartition::IndexRange { + start, + scan_size: cmp::min(keys_per_partition, leader_keys - start), + }); + } + Some(partitions) +} + +/// One coarse, disjoint partition used to drive the top join level. /// -/// `scan_size` is the exact number of leader keys in the shard. Besides -/// avoiding another lookup during execution, it lets a demand-driven executor -/// coalesce many top bindings into roughly one recursive job per worker. +/// Cached scalar intersections naturally partition by physical index shard. +/// Filtered scalar indexes use borrowed key-ordinal ranges, and fused +/// intersections use ordinal ranges of their cover subset. The range forms +/// work for dense, sparse, whole-table, and seminaive subsets without copying +/// row IDs. `scan_size` is the number of leader keys or cover rows in the +/// partition. #[derive(Clone, Copy, Debug, Eq, PartialEq)] -struct TopIndexPartition { - shard: usize, - scan_size: usize, +enum TopLevelPartition { + IndexShard { shard: usize, scan_size: usize }, + IndexRange { start: usize, scan_size: usize }, + CoverRange { start: usize, scan_size: usize }, +} + +impl TopLevelPartition { + fn scan_size(self) -> usize { + match self { + Self::IndexShard { scan_size, .. } + | Self::IndexRange { scan_size, .. } + | Self::CoverRange { scan_size, .. } => scan_size, + } + } +} + +fn cover_scan_bounds(partition: Option, cover_rows: usize) -> (Offset, usize) { + match partition { + Some(TopLevelPartition::CoverRange { start, scan_size }) => { + debug_assert!(start <= cover_rows); + debug_assert!(scan_size <= cover_rows - start); + (Offset::from_usize(start), scan_size) + } + None => (Offset::new(0), cover_rows), + Some(TopLevelPartition::IndexShard { .. } | TopLevelPartition::IndexRange { .. }) => { + unreachable!("an index partition can only drive a scalar intersection") + } + } } /// A key ordinal in a root continuation grid. @@ -1007,10 +1095,6 @@ where self.size() == 0 } - fn is_root(&self) -> bool { - matches!(self, Self::Root(_)) - } - #[cfg(test)] fn root_arc(&self) -> &Arc { let Self::Root(root) = self else { @@ -1632,6 +1716,64 @@ where } } + /// Visit a disjoint ordinal range of an unsharded scalar index. + /// + /// `Intersect` stages always probe one column, so partitioning the first + /// (and only) key level preserves complete key groups. The backing + /// projected/packed index remains borrowed by every coarse task. + fn for_each_range( + &self, + start: usize, + scan_size: usize, + mut f: impl FnMut(&[Value], ProbeMatch<'rows, 'exec>), + ) { + let end = start + .checked_add(scan_size) + .expect("top index range overflow"); + match &self.ix { + ProbeIndex::ProjectedRoot(projected) => { + debug_assert_eq!(projected.columns.len(), 1); + assert!(end <= projected.first.len()); + for key_index in start..end { + let key = [projected.first.value_at(key_index)]; + f( + &key, + Self::keep_or_discard(projected.scalar_rows(key_index), self.keep_rows), + ); + } + } + ProbeIndex::SmallColumn(index) => { + assert!(end <= index.n_keys); + for key_index in start..end { + let rows = if self.keep_rows { + ProbeMatch::Rows(AtomRows::Inline(index.rows_at(key_index))) + } else { + ProbeMatch::Present + }; + f(&index.keys[key_index..key_index + 1], rows); + } + } + ProbeIndex::Packed(packed) => { + debug_assert_eq!(packed.columns.len(), 1); + let values = packed.first.values(); + assert!(end <= values.len()); + for key_index in start..end { + let rows = AtomRows::Packed(PackedCursor::new(packed.first, key_index)); + f( + &values[key_index..key_index + 1], + Self::keep_or_discard(rows, self.keep_rows), + ); + } + } + ProbeIndex::CachedTuple { .. } | ProbeIndex::CachedColumn { .. } => { + unreachable!("persistent indexes use physical shard partitions") + } + ProbeIndex::SmallExact(..) => { + unreachable!("a scalar intersection cannot use an exact tuple probe") + } + } + } + fn shard_count(&self) -> Option { match &self.ix { ProbeIndex::CachedTuple { table, .. } => Some(table.shard_count()), @@ -1645,25 +1787,9 @@ where fn shard_len(&self, shard: usize) -> Option { match &self.ix { - ProbeIndex::CachedTuple { - intersect_outer: None, - table, - .. - } => Some(table.shard_len(shard)), - ProbeIndex::CachedColumn { - intersect_outer: None, - table, - .. - } => Some(table.shard_len(shard)), - ProbeIndex::CachedTuple { - intersect_outer: Some(_), - .. - } - | ProbeIndex::CachedColumn { - intersect_outer: Some(_), - .. - } - | ProbeIndex::ProjectedRoot(..) + ProbeIndex::CachedTuple { table, .. } => Some(table.shard_len(shard)), + ProbeIndex::CachedColumn { table, .. } => Some(table.shard_len(shard)), + ProbeIndex::ProjectedRoot(..) | ProbeIndex::SmallColumn(..) | ProbeIndex::SmallExact(..) | ProbeIndex::Packed(..) => None, @@ -3050,71 +3176,66 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { } } - /// Describe an eligible cached index that could drive a top-level + /// Describe an eligible scalar index that could drive a top-level /// generic-join intersection. /// - /// This mirrors the leader selection in [`Self::run_plan`]. Probers move - /// trie nodes out of `binding_info`, so restore every node before returning. - /// Every scan must cover its whole table through an unfiltered cached - /// hash/column index. Requiring this before constructing any prober avoids - /// building a dynamic nonleader index during discovery and then rebuilding - /// it in every shard job. Tiny or badly skewed leaders stay on the existing - /// path as well. - fn top_index_shards<'rows>( + /// This mirrors leader selection and child-index shape in + /// [`Self::run_plan`]. Probers move trie nodes out of `binding_info`, so + /// restore every node before returning. Persistent catalog indexes use + /// their physical hash shards, including when a dense timestamp range is + /// intersected at probe time. Projected and packed filtered roots use + /// disjoint ordinal ranges of the already-built scalar index, avoiding a + /// scan-and-project handoff copy. Tiny leaders stay on the existing path. + #[allow(clippy::too_many_arguments)] + fn top_index_partitions<'rows>( &self, - stage: &JoinStage, - prepared_join: &'rows PreparedJoinIndexes, - prepared: &[PreparedIndexSlot], + stages: &JoinStages, + stage_index: usize, + prepared: &'rows PreparedJoinIndexes, atoms: &Arc>, + order: &InstrOrder, binding_info: &mut BindingInfo<'rows, 'exec>, workers: usize, - ) -> Option> + ) -> Option> where 'exec: 'rows, { - let JoinStage::Intersect { scans, .. } = stage else { + let JoinStage::Intersect { scans, .. } = &stages.instrs[stage_index] else { return None; }; - if scans.is_empty() || workers <= 1 { + if scans.is_empty() + || workers <= 1 + || estimate_size(&stages.instrs[stage_index], binding_info) + < MIN_TOP_INDEX_KEYS_PER_WORKER.saturating_mul(workers) + { return None; } - debug_assert_eq!(scans.len(), prepared.len()); - - // Mirror the cached/unfiltered branch in `get_index` without actually - // constructing any dynamic fallback indexes. - for scan in scans { - let rows = &binding_info.subsets[scan.atom]; - let info = &self.db.tables[atoms[scan.atom].table]; - if !rows.is_root() - || !scan.cs.is_empty() - || info.table.has_stale_rows() - || !columns_are_cacheable(info, &[scan.column]) - { - return None; - } - let SubsetRef::Dense(subset) = rows.subset() else { - return None; - }; - let Subset::Dense(whole_table) = info.table.all() else { - return None; - }; - if subset != whole_table { - return None; - } - } + let stage_prepared = prepared.stage(stage_index); + debug_assert_eq!(scans.len(), stage_prepared.len()); + let remaining_after_current = prepared + .all_stage_mask() + .map(|remaining| remaining & !(1u64 << stage_index)); let mut leader = 0; let mut leader_size = usize::MAX; let mut probers = Vec::with_capacity(scans.len()); - for (i, (scan, prepared)) in scans.iter().zip(prepared).enumerate() { + for (i, (scan, prepared_slot)) in scans.iter().zip(stage_prepared).enumerate() { + let tail = atom_tail_use( + scan.atom, + &stages.instrs, + prepared, + remaining_after_current, + order, + 1, + ); let prober = self.get_index( atoms, binding_info, ProbeRequest::column( scan, false, - ChildShape::Leaf, - prepared_join.resolve(prepared), + tail.child_shape, + prepared.resolve(prepared_slot), ), ); let size = prober.len(); @@ -3124,11 +3245,11 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { } probers.push(prober); } - let shards = probers[leader].shard_count().and_then(|shard_count| { + let partitions = if let Some(shard_count) = probers[leader].shard_count() { let shards = (0..shard_count) .filter_map(|shard| { let scan_size = probers[leader].shard_len(shard).unwrap_or(0); - (scan_size != 0).then_some(TopIndexPartition { shard, scan_size }) + (scan_size != 0).then_some(TopLevelPartition::IndexShard { shard, scan_size }) }) .collect::>(); top_index_shape_is_eligible( @@ -3138,24 +3259,27 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { MIN_TOP_INDEX_KEYS_PER_WORKER, ) .then_some(shards) - }); + } else { + top_index_range_partitions(leader_size, workers, MIN_TOP_INDEX_KEYS_PER_WORKER) + }; for (scan, prober) in scans.iter().zip(probers) { binding_info.move_back(scan.atom, prober); } - shards + partitions } - /// Return a coarse partition for the variable already sorted to the top of - /// the join. Later variables are deliberately not promoted: benchmarking - /// found that probing and reordering them was not broadly beneficial. - fn select_top_index_shards<'rows>( + /// Return coarse partitions for the stage already sorted to the top of the + /// join. Cached scalar intersections use physical index shards, while + /// fused intersections use ordinal ranges of their cover subset. Later + /// stages are deliberately not promoted. + fn select_top_partitions<'rows>( &self, stages: &JoinStages, prepared: &'rows PreparedJoinIndexes, atoms: &Arc>, order: &InstrOrder, binding_info: &mut BindingInfo<'rows, 'exec>, - ) -> Option> + ) -> Option> where 'exec: 'rows, { @@ -3164,14 +3288,25 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { } let stage_index = order.get(0); - self.top_index_shards( - &stages.instrs[stage_index], - prepared, - prepared.stage(stage_index), - atoms, - binding_info, - crate::parallel::current_num_threads(), - ) + let workers = crate::parallel::current_num_threads(); + let stage = &stages.instrs[stage_index]; + match stage { + JoinStage::Intersect { .. } => self.top_index_partitions( + stages, + stage_index, + prepared, + atoms, + order, + binding_info, + workers, + ), + JoinStage::FusedIntersect { cover, .. } => top_cover_partitions( + binding_info.subsets[cover.to_index.atom].size(), + workers, + MIN_TOP_COVER_ROWS_PER_PARTITION, + ), + JoinStage::FusedIntersectMat { .. } => None, + } } /// Runs the free join plan, starting with the header. @@ -3218,21 +3353,20 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { sort_plan_by_size(&mut order, &mut leaf_scans, 0, &stages.instrs, binding_info); let all_stages = prepared.all_stage_mask(); - // A cached top index is already partitioned by hash. Schedule one - // coarse global job per nonempty shard and run each shard's subtree - // serially. Besides avoiding an intermediate key copy, this keeps - // related nested probes on one worker. Selecting this path commits the - // whole subtree to coarse-only parallelism: `recur_global_serial` - // supplies a serial buffer whose recursive work stays inline. Buffers - // that cannot construct an independent partition (the in-place - // executors) decline this path. + // Schedule each coarse top partition as a global job and run its + // complete subtree inline under the production `Serial` policy. The + // experimental top-shard policy may instead fan out later recursive + // work. Cached scalar indexes already provide physical hash shards; + // fused covers use borrowed ordinal ranges, so seminaive and other + // filtered roots need no row-ID copy. Buffers that cannot construct an + // independent partition decline this path. if !stages.instrs.is_empty() && action_buf.supports_global_partition() - && let Some(shards) = - self.select_top_index_shards(stages, prepared, atoms, &order, binding_info) + && let Some(partitions) = + self.select_top_partitions(stages, prepared, atoms, &order, binding_info) { let mut updates = FrameUpdates::with_capacity(0); - for partition in shards { + for partition in partitions { let db = self.db; let exec_state_for_factory = self.exec_state; let exec_state_for_work = self.exec_state; @@ -3296,8 +3430,8 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { /// /// This method takes the plan, mutable data structures for variable binding /// and staging actions, and `cur`, the current stage of the plan. A top-level - /// coarse partition also passes `index_partition`; recursive calls clear it - /// so only the first intersection is restricted to that physical shard. + /// coarse partition also passes `top_partition`; recursive calls clear it + /// so only the first intersection or fused cover scan is restricted. /// The associated top-shard buffer policy, rather than the presence of the /// partition itself, determines whether later recursive work may fan out. #[allow(clippy::too_many_arguments)] @@ -3310,7 +3444,7 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { instr_order: &mut InstrOrder, leaf_scans: &mut LeafScans, cur: usize, - index_partition: Option, + top_partition: Option, remaining_stages: Option, binding_info: &mut BindingInfo<'rows, 'exec>, action_buf: &mut BUF, @@ -3345,9 +3479,9 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { ); return; } - let mut cur_size = index_partition.map_or_else( + let mut cur_size = top_partition.map_or_else( || estimate_size(&stages.instrs[instr_order.get(cur)], binding_info), - |partition| partition.scan_size, + TopLevelPartition::scan_size, ); if cur_size > 32 && cur % 3 == 1 && cur < instr_order.len() - 1 { // Re-evaluate the remaining suffix after observing the residuals @@ -3449,7 +3583,7 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { // `drain_updates_parallel!` is a bit slower because of the additional // `ExecutionState` clone. if cur < free_join_fork_depth() - && action_buf.supports_parallel_drain(cur, index_partition.is_some()) + && action_buf.supports_parallel_drain(cur, top_partition.is_some()) { drain_updates_parallel!($updates) } else { @@ -3510,14 +3644,21 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { } // A sharded top-level job enumerates only its assigned physical index - // shard. Every recursive call clears `index_partition`, so this macro is + // shard. Every recursive call clears `top_partition`, so this macro is // used only by the leading prober of the initial `Intersect` stage. macro_rules! for_each_leader { ($prober:expr, $callback:expr) => { - if let Some(partition) = index_partition { - $prober.for_each_shard(partition.shard, $callback) - } else { - $prober.for_each($callback) + match top_partition { + Some(TopLevelPartition::IndexShard { shard, .. }) => { + $prober.for_each_shard(shard, $callback) + } + Some(TopLevelPartition::IndexRange { + start, scan_size, .. + }) => $prober.for_each_range(start, scan_size, $callback), + None => $prober.for_each($callback), + Some(TopLevelPartition::CoverRange { .. }) => { + unreachable!("a cover range can only drive a fused intersection") + } } }; } @@ -3713,11 +3854,13 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { SmallVec::<[ColumnId; 4]>::from_iter(bind.iter().map(|(col, _)| *col)); let vars = bind.iter().map(|(_, var)| *var).collect(); let mut buf = TaggedRowBuffer::new_inline(bind.len()); + let (scan_start, scan_size) = + cover_scan_bounds(top_partition, cover_subset.size()); table.scan_project( cover_subset, &proj, - Offset::new(0), - usize::MAX, + scan_start, + scan_size, &cover.constraints, &mut buf, ); @@ -3747,20 +3890,23 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { SmallVec::<[ColumnId; 4]>::from_iter(bind.iter().map(|(col, _)| *col)); let cover_node = binding_info.unwrap_val(cover_atom); let cover_subset = cover_node.subset(); - let mut offset = Offset::new(0); + let (mut offset, mut scan_remaining) = + cover_scan_bounds(top_partition, cover_subset.size()); let mut buffer = TaggedRowBuffer::new(bind.len()); let mut updates = FrameUpdates::with_capacity(cmp::min(chunk_size, cur_size)); - loop { + while scan_remaining != 0 { buffer.clear(); let table = &self.db.tables[atoms[cover_atom].table].table; + let scan_size = cmp::min(chunk_size, scan_remaining); let next = table.scan_project( cover_subset, &proj, offset, - chunk_size, + scan_size, &cover.constraints, &mut buffer, ); + scan_remaining -= scan_size; for (row, key) in buffer.iter() { if keep_cover { updates.refine_atom_dense( @@ -3777,11 +3923,12 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { drain_updates!(updates); } } - if let Some(next) = next { + if scan_remaining != 0 { + let Some(next) = next else { + break; + }; offset = next; - continue; } - break; } drain_updates!(updates, false); // Restore the subsets we swapped out. @@ -3837,20 +3984,23 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { let proj = SmallVec::<[ColumnId; 4]>::from_iter(bind.iter().map(|(col, _)| *col)); let cover_node = binding_info.unwrap_val(cover_atom); let cover_subset = cover_node.subset(); - let mut cur = Offset::new(0); + let (mut scan_offset, mut scan_remaining) = + cover_scan_bounds(top_partition, cover_subset.size()); let mut buffer = TaggedRowBuffer::new(bind.len()); let mut updates = FrameUpdates::with_capacity(cmp::min(chunk_size, cur_size)); - loop { + while scan_remaining != 0 { buffer.clear(); let table = &self.db.tables[atoms[cover_atom].table].table; + let scan_size = cmp::min(chunk_size, scan_remaining); let next = table.scan_project( cover_subset, &proj, - cur, - chunk_size, + scan_offset, + scan_size, &cover.constraints, &mut buffer, ); + scan_remaining -= scan_size; 'mid: for (row, key) in buffer.iter() { if keep_cover { updates.refine_atom_dense(cover_atom, OffsetRange::new(row, row.inc())); @@ -3884,11 +4034,12 @@ impl<'a, 'state, 'exec> JoinState<'a, 'state, 'exec> { drain_updates!(updates); } } - if let Some(next) = next { - cur = next; - continue; + if scan_remaining != 0 { + let Some(next) = next else { + break; + }; + scan_offset = next; } - break; } // TODO: special-case the scenario when the cover doesn't need // deduping (and hence we can do a straight scan: e.g. when the @@ -5527,9 +5678,11 @@ mod top_index_tests { ExperimentalTopShardPolicy, InstrOrder, LazyArenaHandle, PreparedIndexKind, PreparedIndexSlot, PreparedIndexState, PreparedIndexStateId, PreparedJoinIndexes, PreparedTailMasks, RootContinuationCache, RootProjection, SmallColumnIndex, - SmallColumnSink, TrieNode, cached_coarse_scan_morsel_size, coarse_scan_morsel_size, + SmallColumnSink, TOP_INDEX_RANGES_PER_WORKER, TopLevelPartition, TrieNode, + cached_coarse_scan_morsel_size, coarse_scan_morsel_size, cover_scan_bounds, for_each_stage_atom, materialization_is_live_in_tail, packed_child_shape_in_tail, - scan_atom_tail_use, sort_plan_by_size_inner, top_index_shape_is_eligible, + scan_atom_tail_use, sort_plan_by_size_inner, top_cover_partitions, + top_index_range_partitions, top_index_shape_is_eligible, }; use crate::free_join::packed_trie::ChildShape; @@ -6122,6 +6275,77 @@ mod top_index_tests { assert!(top_index_shape_is_eligible(4, 40, 4, 10)); } + #[test] + fn fused_cover_partitions_are_complete_disjoint_and_coarse() { + assert!(top_cover_partitions(10_000, 1, 16).is_none()); + assert!(top_cover_partitions(63, 4, 16).is_none()); + + for (cover_rows, workers) in [(64, 4), (1_003, 4), (585_276, 12)] { + let partitions = top_cover_partitions(cover_rows, workers, 16).unwrap(); + assert!(partitions.len() <= workers * 4); + + let mut expected_start = 0; + for partition in &partitions { + let TopLevelPartition::CoverRange { start, scan_size } = *partition else { + panic!("fused covers must produce ordinal ranges") + }; + assert_eq!(start, expected_start, "ranges must be contiguous"); + assert!(scan_size > 0); + if start + scan_size != cover_rows { + assert!( + scan_size >= 16, + "all non-tail partitions must remain coarse" + ); + } + expected_start += scan_size; + } + assert_eq!(expected_start, cover_rows, "ranges must cover every row"); + } + } + + #[test] + fn fused_cover_scan_bounds_preserve_nonzero_subset_origins() { + let partition = TopLevelPartition::CoverRange { + start: 17, + scan_size: 23, + }; + let (start, len) = cover_scan_bounds(Some(partition), 100); + assert_eq!(start.index(), 17); + assert_eq!(len, 23); + + let (start, len) = cover_scan_bounds(None, 100); + assert_eq!(start.index(), 0); + assert_eq!(len, 100); + } + + #[test] + fn filtered_index_ranges_are_complete_disjoint_and_coarse() { + assert!(top_index_range_partitions(10_000, 1, 16).is_none()); + assert!(top_index_range_partitions(63, 4, 16).is_none()); + + for (leader_keys, workers) in [(64, 4), (1_003, 4), (585_276, 12)] { + let partitions = top_index_range_partitions(leader_keys, workers, 16).unwrap(); + assert!(partitions.len() <= workers * TOP_INDEX_RANGES_PER_WORKER); + + let mut expected_start = 0; + for partition in &partitions { + let TopLevelPartition::IndexRange { start, scan_size } = *partition else { + panic!("an unsharded index must produce key-ordinal ranges") + }; + assert_eq!(start, expected_start, "ranges must be contiguous"); + assert!(scan_size > 0); + if start + scan_size != leader_keys { + assert!( + scan_size >= 16, + "all non-tail partitions must remain coarse" + ); + } + expected_start += scan_size; + } + assert_eq!(expected_start, leader_keys, "ranges must cover every key"); + } + } + #[test] fn task_clone_keeps_only_atoms_in_the_dynamic_join_tail() { let stages = vec![ diff --git a/core-relations/src/tests.rs b/core-relations/src/tests.rs index 21577edc1..530ce0a95 100644 --- a/core-relations/src/tests.rs +++ b/core-relations/src/tests.rs @@ -9,7 +9,7 @@ use egglog_reports::ReportLevel; use crate::numeric_id::NumericId; use crate::{ - PlanStrategy, + PlanStrategy, Subset, action::WriteVal, common::Value, free_join::{ @@ -1801,6 +1801,286 @@ fn gj_top_index_shards_preserve_count_and_materialization_inner() { }); } +#[test] +fn gj_filtered_top_index_ranges_preserve_count_and_materialization() { + let serial = gj_filtered_top_index_range_fixture(1); + let parallel = gj_filtered_top_index_range_fixture(4); + assert_eq!(parallel.0, serial.0, "parallel match count changed"); + assert_eq!(parallel.1, serial.1, "parallel actions changed"); + assert!( + parallel.2 > 4, + "filtered top-index ranges did not enqueue global work" + ); + assert_eq!( + parallel.3, 0, + "filtered top-index range subtrees must stay serial" + ); +} + +/// Return `(matches, materialized rows, global pushes, local pushes)`. +/// +/// The filtered left input is deliberately smaller than half of its table. +/// It therefore builds the normal execution-scoped packed scalar index rather +/// than using the persistent whole-table catalog. Coarse search jobs must +/// borrow disjoint key ranges from that index without rescanning the table. +fn gj_filtered_top_index_range_fixture(num_threads: usize) -> (usize, Vec>, u64, u64) { + let pool = egglog_concurrency::ThreadPool::new(num_threads); + pool.install(|| { + let mut db = Database::default(); + let left = db.add_table( + SortedWritesTable::new( + 3, + 3, + Some(ColumnId::new(2)), + vec![], + Box::new(|_, old, new, _| { + assert_eq!(old, new, "left rows are unique"); + false + }), + ), + iter::empty(), + iter::empty(), + ); + let right = add_set_table(&mut db, 2); + let output = add_set_table(&mut db, 2); + + const ROWS: usize = 12_001; + const NEW_START: usize = 9_001; + { + let mut left_buf = db.new_buffer(left); + let mut right_buf = db.new_buffer(right); + for key in 0..NEW_START { + left_buf.stage_insert(&[v(key), v(100_000 + key), v(0)]); + right_buf.stage_insert(&[v(key), v(200_000 + key)]); + } + } + db.merge_all(); + { + let mut left_buf = db.new_buffer(left); + let mut right_buf = db.new_buffer(right); + for key in NEW_START..ROWS { + left_buf.stage_insert(&[v(key), v(100_000 + key), v(1)]); + right_buf.stage_insert(&[v(key), v(200_000 + key)]); + } + } + db.merge_all(); + + let recent = Constraint::GeConst { + col: ColumnId::new(2), + val: v(1), + }; + let mut rsb = db.new_rule_set(); + let mut query = rsb.new_rule(); + query.set_plan_strategy(PlanStrategy::Gj); + query.set_no_decomp(true); + let key = query.new_var_named("key"); + let left_value = query.new_var_named("left-value"); + let timestamp = query.new_var_named("timestamp"); + let right_value = query.new_var_named("right-value"); + query + .add_atom( + left, + &[key.into(), left_value.into(), timestamp.into()], + std::slice::from_ref(&recent), + ) + .unwrap(); + query + .add_atom(right, &[key.into(), right_value.into()], &[]) + .unwrap(); + let mut rule = query.build(); + rule.insert(output, &[left_value.into(), right_value.into()]) + .unwrap(); + rule.build_with_description("filtered-index-range-gj"); + let rules = rsb.build(); + + let (plan, _, _) = rules.plans.values().next().unwrap(); + let Plan::SinglePlan(plan) = plan else { + panic!("set_no_decomp must produce a single plan") + }; + assert!(matches!( + &plan.stages.instrs[0], + JoinStage::Intersect { scans, .. } if scans.len() == 2 + )); + + pool.reset_scheduler_metrics(); + let report = db.run_rule_set(&rules, ReportLevel::TimeOnly); + let scheduler = pool.scheduler_metrics(); + let expected = ROWS - NEW_START; + assert_eq!(report.num_matches("filtered-index-range-gj"), expected); + let rows = table_rows(&db, output); + assert_eq!(rows.len(), expected); + ( + report.num_matches("filtered-index-range-gj"), + rows, + scheduler.global_pushes, + scheduler.local_pushes, + ) + }) +} + +#[test] +fn fused_cover_partitions_preserve_matches_and_actions() { + let serial = fused_cover_partition_fixture(1, false); + let parallel = fused_cover_partition_fixture(4, false); + assert_eq!(parallel.0, serial.0, "parallel match count changed"); + assert_eq!(parallel.1, serial.1, "parallel actions changed"); + assert!( + parallel.2 > 4, + "coarse fused-cover partitioning did not enqueue global work" + ); + assert_eq!( + parallel.3, 0, + "coarse fused-cover subtrees must stay serial" + ); +} + +#[test] +fn fused_cover_partitions_preserve_seminaive_subset() { + let serial = fused_cover_partition_fixture(1, true); + let parallel = fused_cover_partition_fixture(4, true); + assert_eq!(parallel.0, serial.0, "parallel match count changed"); + assert_eq!(parallel.1, serial.1, "parallel actions changed"); + assert!( + parallel.2 > 4, + "filtered fused cover did not enqueue global work" + ); + assert_eq!( + parallel.3, 0, + "filtered fused-cover subtrees must stay serial" + ); +} + +/// Return `(matches, materialized rows, global pushes, local pushes)`. +/// +/// The timestamp-filtered case models a seminaive root: its cover is a dense +/// suffix with a nonzero origin rather than the table's whole-row subset. +fn fused_cover_partition_fixture( + num_threads: usize, + timestamp_filtered: bool, +) -> (usize, Vec>, u64, u64) { + let pool = egglog_concurrency::ThreadPool::new(num_threads); + pool.install(|| { + let mut db = Database::default(); + let facts = db.add_table( + SortedWritesTable::new( + 3, + 3, + Some(ColumnId::new(2)), + vec![], + Box::new(|_, old, new, _| { + assert_eq!(old, new, "facts are unique"); + false + }), + ), + iter::empty(), + iter::empty(), + ); + let gate = add_set_table(&mut db, 1); + let output = add_set_table(&mut db, 3); + + const FACT_ROWS: usize = 12_001; + const ZERO_TIMESTAMP_ROWS: usize = 4_001; + const GATE_ROWS: usize = 16_001; + { + let mut facts_buf = db.new_buffer(facts); + for i in 0..ZERO_TIMESTAMP_ROWS { + facts_buf.stage_insert(&[v(i), v(100_000 + i), v(0)]); + } + } + { + let mut gate_buf = db.new_buffer(gate); + for i in 0..GATE_ROWS { + gate_buf.stage_insert(&[v(i)]); + } + } + db.merge_all(); + { + let mut facts_buf = db.new_buffer(facts); + for i in ZERO_TIMESTAMP_ROWS..FACT_ROWS { + facts_buf.stage_insert(&[v(i), v(100_000 + i), v(1)]); + } + } + db.merge_all(); + + let timestamp = Constraint::GeConst { + col: ColumnId::new(2), + val: v(1), + }; + let mut rsb = db.new_rule_set(); + let mut query = rsb.new_rule(); + query.set_plan_strategy(PlanStrategy::PureSize); + query.set_no_decomp(true); + let key = query.new_var_named("key"); + let payload = query.new_var_named("payload"); + let ts = query.new_var_named("timestamp"); + let facts_atom = query + .add_atom( + facts, + &[key.into(), payload.into(), ts.into()], + timestamp_filtered.then_some(×tamp), + ) + .unwrap(); + query.add_atom(gate, &[key.into()], &[]).unwrap(); + let mut rule = query.build(); + rule.insert(output, &[key.into(), payload.into(), ts.into()]) + .unwrap(); + rule.build_with_description("coarse-fused-cover"); + let rules = rsb.build(); + + let (plan, _, _) = rules.plans.values().next().unwrap(); + let Plan::SinglePlan(plan) = plan else { + panic!("set_no_decomp must produce a single plan") + }; + assert_eq!(plan.stages.instrs.len(), 1); + assert!(matches!( + &plan.stages.instrs[0], + JoinStage::FusedIntersect { + cover, + to_intersect, + .. + } if cover.to_index.atom == facts_atom && !to_intersect.is_empty() + )); + let expected = if timestamp_filtered { + FACT_ROWS - ZERO_TIMESTAMP_ROWS + } else { + FACT_ROWS + }; + if timestamp_filtered { + let cover_header = plan + .header + .iter() + .find(|header| header.atom == facts_atom) + .unwrap(); + assert_eq!(cover_header.subset.size(), expected); + let Subset::Dense(range) = &cover_header.subset else { + panic!("sorted timestamp filter must produce a dense suffix") + }; + assert!( + range.start.index() > 0, + "the seminaive fixture must exercise a nonzero subset origin" + ); + } else { + assert!( + plan.header.iter().all(|header| header.atom != facts_atom), + "an unconstrained cover should use the whole table" + ); + } + + pool.reset_scheduler_metrics(); + let report = db.run_rule_set(&rules, ReportLevel::TimeOnly); + let scheduler = pool.scheduler_metrics(); + assert_eq!(report.num_matches("coarse-fused-cover"), expected); + let rows = table_rows(&db, output); + assert_eq!(rows.len(), expected); + ( + report.num_matches("coarse-fused-cover"), + rows, + scheduler.global_pushes, + scheduler.local_pushes, + ) + }) +} + #[test] fn gj_small_top_fallback_uses_local_queue() { let pool = egglog_concurrency::ThreadPool::new(4); diff --git a/core-relations/src/uf/mod.rs b/core-relations/src/uf/mod.rs index 9667dce0d..60ab8e00e 100644 --- a/core-relations/src/uf/mod.rs +++ b/core-relations/src/uf/mod.rs @@ -15,6 +15,7 @@ use crate::{ action::ExecutionState, common::{HashMap, Value}, offsets::{OffsetRange, RowId, Subset, SubsetRef}, + parallel, pool::with_pool_set, row_buffer::RowBuffer, table_spec::{ @@ -28,6 +29,22 @@ mod tests; type UnionFind = crate::union_find::UnionFind; +const PARALLEL_UNION_CHUNK_ROWS: usize = 16 * 1024; + +#[derive(Clone, Copy)] +struct UnionChunk { + buffer: usize, + start: usize, + end: usize, +} + +struct TimestampGroup { + timestamp: Value, + chunks: std::ops::Range, + rows: usize, + max_value: Value, +} + /// A special table backed by a union-find used to efficiently implement /// egglog-style canonicaliztion. /// @@ -439,22 +456,101 @@ impl Table for DisplacedTable { } fn merge(&mut self, _: &mut ExecutionState) -> TableChange { - while let Some(rowbuf) = self.buffered_writes.pop() { - for row in rowbuf.iter() { - self.changed |= self.insert_impl(row).is_some(); + if parallel::current_num_threads() < 2 { + while let Some(rowbuf) = self.buffered_writes.pop() { + for row in rowbuf.iter() { + self.changed |= self.insert_impl(row).is_some(); + } } + return self.take_change(); } + + let mut buffers = Vec::new(); + let mut total_rows = 0; + while let Some(buffer) = self.buffered_writes.pop() { + total_rows += buffer.len(); + buffers.push(buffer); + } + if total_rows <= PARALLEL_UNION_CHUNK_ROWS { + for buffer in &buffers { + for row in buffer.iter() { + self.changed |= self.insert_impl(row).is_some(); + } + } + return self.take_change(); + } + + let (chunks, groups) = Self::union_work(&buffers); + // Timestamps are semantic barriers: a value must be recorded at the + // first timestamp at which it ceases to be canonical. Union order + // within one timestamp is intentionally unconstrained, but a later + // timestamp cannot race an earlier one. + for group in groups { + if group.rows <= PARALLEL_UNION_CHUNK_ROWS { + for chunk in &chunks[group.chunks] { + for row_index in chunk.start..chunk.end { + let row = buffers[chunk.buffer].get_row(RowId::from_usize(row_index)); + self.changed |= self.insert_impl(row).is_some(); + } + } + continue; + } + + let successes = self.uf.with_atomic_access(group.max_value, |uf| { + if !parallel::enabled_for_len(group.chunks.len()) { + let mut children = Vec::with_capacity(group.rows); + for chunk in &chunks[group.chunks.clone()] { + for row_index in chunk.start..chunk.end { + let row = buffers[chunk.buffer].get_row(RowId::from_usize(row_index)); + let (parent, child) = uf.union(row[0], row[1]); + if parent != child { + children.push(child); + } + } + } + vec![children] + } else { + parallel::map(&chunks[group.chunks.clone()], |_, chunk| { + let mut children = Vec::with_capacity(chunk.end - chunk.start); + for row_index in chunk.start..chunk.end { + let row = buffers[chunk.buffer].get_row(RowId::from_usize(row_index)); + let (parent, child) = uf.union(row[0], row[1]); + if parent != child { + children.push(child); + } + } + children + }) + } + }); + + let added = successes.iter().map(Vec::len).sum(); + self.displaced.reserve(added); + self.lookup_table.reserve(added); + // A successful CAS links a root to a smaller root. It can never + // become a root again, so each returned child is globally unique + // and may be published once without deduplication. Publication + // stays serial because the backing Vec and HashMap are serial. + for child in successes.into_iter().flatten() { + self.publish_displaced(child, group.timestamp); + self.changed = true; + } + } + self.take_change() + } +} + +impl DisplacedTable { + fn take_change(&mut self) -> TableChange { let changed = mem::take(&mut self.changed); - // UF table rows can be updated "in place", we count both added and removed as changed in - // this case. + // UF table rows can be updated "in place", so count a change as both + // an addition and a removal. TableChange { added: changed, removed: changed, } } -} -impl DisplacedTable { pub fn underlying_uf(&self) -> &UnionFind { &self.uf } @@ -491,7 +587,11 @@ impl DisplacedTable { // Compress paths somewhat, given that we perform naive finds everywhere else. let _ = self.uf.find(parent); let _ = self.uf.find(child); - let ts = row[2]; + self.publish_displaced(child, row[2]); + Some((parent, child)) + } + + fn publish_displaced(&mut self, child: Value, ts: Value) { if let Some((_, highest)) = self.displaced.last() { assert!( *highest <= ts, @@ -501,7 +601,87 @@ impl DisplacedTable { let next = RowId::from_usize(self.displaced.len()); self.displaced.push((child, ts)); self.lookup_table.insert(child, next); - Some((parent, child)) + } + + fn union_work(buffers: &[RowBuffer]) -> (Vec, Vec) { + let mut chunks = Vec::new(); + let mut groups = Vec::::new(); + + for (buffer_index, buffer) in buffers.iter().enumerate() { + assert_eq!( + buffer.arity(), + 3, + "attempt to insert a row with the wrong arity" + ); + let mut run_start = 0; + let mut run_timestamp = None; + let mut run_max = None; + for (row_index, row) in buffer.iter().enumerate() { + let timestamp = row[2]; + if run_timestamp.is_some_and(|current| current != timestamp) { + Self::push_union_run( + &mut chunks, + &mut groups, + buffer_index, + run_start, + row_index, + run_timestamp.unwrap(), + run_max.unwrap(), + ); + run_start = row_index; + run_max = None; + } + run_timestamp = Some(timestamp); + let row_max = row[0].max(row[1]); + run_max = Some(run_max.map_or(row_max, |current: Value| current.max(row_max))); + } + if let Some(timestamp) = run_timestamp { + Self::push_union_run( + &mut chunks, + &mut groups, + buffer_index, + run_start, + buffer.len(), + timestamp, + run_max.unwrap(), + ); + } + } + (chunks, groups) + } + + fn push_union_run( + chunks: &mut Vec, + groups: &mut Vec, + buffer: usize, + start: usize, + end: usize, + timestamp: Value, + max_value: Value, + ) { + let group = match groups.last_mut() { + Some(group) if group.timestamp == timestamp => group, + _ => { + let start = chunks.len(); + groups.push(TimestampGroup { + timestamp, + chunks: start..start, + rows: 0, + max_value, + }); + groups.last_mut().unwrap() + } + }; + for start in (start..end).step_by(PARALLEL_UNION_CHUNK_ROWS) { + chunks.push(UnionChunk { + buffer, + start, + end: (start + PARALLEL_UNION_CHUNK_ROWS).min(end), + }); + } + group.chunks.end = chunks.len(); + group.rows += end - start; + group.max_value = group.max_value.max(max_value); } } diff --git a/core-relations/src/uf/tests.rs b/core-relations/src/uf/tests.rs index 106fd61db..5d163ed71 100644 --- a/core-relations/src/uf/tests.rs +++ b/core-relations/src/uf/tests.rs @@ -51,3 +51,121 @@ fn displaced() { d.scan_generic(all.as_ref(), |_, row| updates_2.push((row[0], row[1]))); assert!(updates_2.windows(2).all(|x| x[0].1 == x[1].1)); } + +#[test] +fn parallel_merge_preserves_timestamp_order_and_lookup() { + empty_execution_state!(e); + let mut d = DisplacedTable::default(); + let n = 40_000; + { + let mut buf = d.new_buffer(); + for i in (0..n).step_by(2) { + buf.stage_insert(&[v(i), v(i + 1), v(0)]); + } + for i in (2..n).step_by(2) { + buf.stage_insert(&[v(0), v(i), v(1)]); + } + } + + let pool = egglog_concurrency::ThreadPool::new(4); + let change = pool.install(|| d.merge(&mut e)); + assert!(change.added); + assert_eq!(d.len(), n - 1); + + let mut timestamps = Vec::new(); + d.scan_generic(d.all().as_ref(), |_, row| timestamps.push(row[2])); + assert_eq!( + timestamps + .iter() + .filter(|&×tamp| timestamp == v(0)) + .count(), + n / 2 + ); + assert_eq!( + timestamps + .iter() + .filter(|&×tamp| timestamp == v(1)) + .count(), + n / 2 - 1 + ); + assert!(timestamps.windows(2).all(|pair| pair[0] <= pair[1])); + + for i in 1..n { + let row = d.get_row(&[v(i)]).expect("every displaced id is indexed"); + assert_eq!(row.vals[0], v(i)); + assert_eq!(row.vals[1], v(0)); + } +} + +#[test] +fn parallel_merge_records_each_displaced_root_once_under_contention() { + empty_execution_state!(e); + let mut d = DisplacedTable::default(); + let n = 20_000; + { + let mut buf = d.new_buffer(); + for repeat in 0..4 { + for i in 1..n { + let (left, right) = if (i + repeat) % 2 == 0 { + (i - 1, i) + } else { + (i, i - 1) + }; + buf.stage_insert(&[v(left), v(right), v(0)]); + } + } + } + + let pool = egglog_concurrency::ThreadPool::new(4); + pool.install(|| d.merge(&mut e)); + assert_eq!(d.len(), n - 1); + for i in 1..n { + assert_eq!(d.get_row_column(&[v(i)], ColumnId::new(1)), Some(v(0))); + } + + let mut clone = d.clone(); + clone.new_buffer().stage_insert(&[v(0), v(n), v(1)]); + clone.merge(&mut e); + assert_eq!(clone.underlying_uf().find_naive(v(n)), v(0)); + assert_eq!(d.underlying_uf().find_naive(v(n)), v(n)); +} + +#[test] +fn parallel_merge_matches_serial_timestamp_oracle() { + empty_execution_state!(e); + let mut d = DisplacedTable::default(); + let mut oracle = crate::union_find::UnionFind::default(); + let mut expected = crate::common::HashMap::default(); + let mut random = 0x1234_5678_u64; + { + let mut buf = d.new_buffer(); + for timestamp in 0..4 { + for _ in 0..20_000 { + random = random + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1); + let left = v((random as usize) % 50_000); + random = random + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1); + let right = v((random as usize) % 50_000); + buf.stage_insert(&[left, right, v(timestamp)]); + + let (parent, child) = oracle.union(left, right); + if parent != child { + assert_eq!(expected.insert(child, v(timestamp)), None); + } + } + } + } + + let pool = egglog_concurrency::ThreadPool::new(8); + pool.install(|| d.merge(&mut e)); + assert_eq!(d.len(), expected.len()); + d.scan_generic(d.all().as_ref(), |row_id, row| { + assert_eq!(expected.remove(&row[0]), Some(row[2])); + assert_eq!(row[1], oracle.find_naive(row[0])); + assert_eq!(d.get_row(&[row[0]]).map(|found| found.id), Some(row_id)); + }); + assert!(expected.is_empty()); +} diff --git a/numeric-id/src/lib.rs b/numeric-id/src/lib.rs index d8fccddda..917b342cd 100644 --- a/numeric-id/src/lib.rs +++ b/numeric-id/src/lib.rs @@ -22,6 +22,22 @@ pub trait NumericId: Copy + Clone + PartialEq + Eq + PartialOrd + Ord + Hash + S } } +/// A [`NumericId`] whose in-memory representation may be accessed atomically. +/// +/// # Safety +/// +/// The ID must have the same size, byte representation, and set of valid bit +/// patterns as both `Self::Rep` and `Self::Atomic`. (The atomic type may require +/// stricter alignment; code creating an atomic view must check that separately.) +/// It must therefore be sound to reinterpret a sufficiently aligned initialized +/// slice of IDs as a slice of atomic integers, provided that atomic and +/// non-atomic accesses do not overlap. +pub unsafe trait AtomicNumericId: NumericId {} + +// SAFETY: `usize` and `AtomicUsize` have the same size, byte representation, +// and valid bit patterns. +unsafe impl AtomicNumericId for usize {} + impl NumericId for usize { type Rep = usize; type Atomic = std::sync::atomic::AtomicUsize; @@ -448,6 +464,7 @@ macro_rules! define_id { ($v:vis $name:ident, $repr:tt, pretty $pretty_name:expr) => { define_id!($v $name, $repr, "", pretty $pretty_name); }; ($v:vis $name:ident, $repr:tt, $doc:tt, pretty $pretty_name:tt) => { #[derive(Copy, Clone)] + #[repr(transparent)] #[doc = $doc] $v struct $name { rep: $repr, @@ -516,6 +533,10 @@ macro_rules! define_id { } } + // SAFETY: this ID is `repr(transparent)` over its integer + // representation, and `Atomic` is the corresponding atomic integer. + unsafe impl $crate::AtomicNumericId for $name {} + impl std::fmt::Debug for $name { fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result { let name = if $pretty_name.is_empty() { diff --git a/union-find/src/concurrent/atomic_int.rs b/union-find/src/concurrent/atomic_int.rs index 72fd6f60b..d213ec79c 100644 --- a/union-find/src/concurrent/atomic_int.rs +++ b/union-find/src/concurrent/atomic_int.rs @@ -11,12 +11,18 @@ pub trait AtomicInt: Default + Send + Sync + 'static { fn from_usize(value: usize) -> Self; fn as_usize(value: Self::Underlying) -> usize; fn load(&self) -> Self::Underlying; + fn load_relaxed(&self) -> Self::Underlying; fn store(&self, value: Self::Underlying); fn cas( &self, current: Self::Underlying, new: Self::Underlying, ) -> Result; + fn cas_relaxed( + &self, + current: Self::Underlying, + new: Self::Underlying, + ) -> Result; } impl AtomicInt for AtomicU32 { @@ -30,12 +36,18 @@ impl AtomicInt for AtomicU32 { fn load(&self) -> u32 { self.load(LOAD_ORDERING) } + fn load_relaxed(&self) -> u32 { + self.load(Ordering::Relaxed) + } fn store(&self, value: u32) { self.store(value, STORE_ORDERING); } fn cas(&self, current: u32, new: u32) -> Result { self.compare_exchange(current, new, CAS_SUCCESS_ORDERING, CAS_FAILURE_ORDERING) } + fn cas_relaxed(&self, current: u32, new: u32) -> Result { + self.compare_exchange(current, new, Ordering::Relaxed, Ordering::Relaxed) + } } impl AtomicInt for AtomicU64 { @@ -49,12 +61,18 @@ impl AtomicInt for AtomicU64 { fn load(&self) -> u64 { self.load(LOAD_ORDERING) } + fn load_relaxed(&self) -> u64 { + self.load(Ordering::Relaxed) + } fn store(&self, value: u64) { self.store(value, STORE_ORDERING); } fn cas(&self, current: u64, new: u64) -> Result { self.compare_exchange(current, new, CAS_SUCCESS_ORDERING, CAS_FAILURE_ORDERING) } + fn cas_relaxed(&self, current: u64, new: u64) -> Result { + self.compare_exchange(current, new, Ordering::Relaxed, Ordering::Relaxed) + } } impl AtomicInt for AtomicUsize { type Underlying = usize; @@ -69,19 +87,26 @@ impl AtomicInt for AtomicUsize { fn load(&self) -> usize { self.load(LOAD_ORDERING) } + fn load_relaxed(&self) -> usize { + self.load(Ordering::Relaxed) + } fn store(&self, value: usize) { self.store(value, STORE_ORDERING); } fn cas(&self, current: usize, new: usize) -> Result { self.compare_exchange(current, new, CAS_SUCCESS_ORDERING, CAS_FAILURE_ORDERING) } + fn cas_relaxed(&self, current: usize, new: usize) -> Result { + self.compare_exchange(current, new, Ordering::Relaxed, Ordering::Relaxed) + } } /* -We could do the following, but we don't have good reason to believe it is -correct, or preserves the linearizability guarantees of the baseline -data-structure. +We could do the following for the growable concurrent data structure, but we +don't have good reason to believe it is correct, or preserves the +linearizability guarantees of the baseline data-structure. The fixed-capacity +scoped view has a narrower contract and a separate relaxed-ordering argument. Still, code like this is helpful in microbenchmarks to measure the overhead of memory barriers for the code. We've observed 5-10% overhead on an M1 Ultra CPU diff --git a/union-find/src/concurrent/uf.rs b/union-find/src/concurrent/uf.rs index 11298da14..3f1efac37 100644 --- a/union-find/src/concurrent/uf.rs +++ b/union-find/src/concurrent/uf.rs @@ -66,7 +66,7 @@ impl ConcurrentUnionFind { pub fn find(&self, elt: T::Underlying) -> T::Underlying { self.data.with_access( T::as_usize(elt) + 1, - |buf| Self::find_impl(buf, elt), + |buf| find_impl(buf, elt), T::from_usize, ) } @@ -77,15 +77,15 @@ impl ConcurrentUnionFind { self.data.with_access( T::as_usize(max_elt) + 1, |buf| { - let mut l = Self::find_impl(buf, l); - let mut r = Self::find_impl(buf, r); + let mut l = find_impl(buf, l); + let mut r = find_impl(buf, r); while l != r { let next = buf[T::as_usize(l)].load(); if next == l { return false; } - l = Self::find_impl(buf, l); - r = Self::find_impl(buf, r); + l = find_impl(buf, l); + r = find_impl(buf, r); } true }, @@ -106,52 +106,121 @@ impl ConcurrentUnionFind { ) { self.data.with_access( T::as_usize(cmp::max(l, r)) + 1, - |buf| { - let mut l = l; - let mut r = r; - loop { - l = Self::find_impl(buf, l); - r = Self::find_impl(buf, r); - if l != r { - // We do "union by min": common in egraphs due to the - // hypothesis that smaller ids will be - // better-represented in the egraph, and hence - // perturbing the maximum id is likely to result in less - // work for rebuilding. - let parent = cmp::min(l, r); - let child = cmp::max(l, r); - match buf[T::as_usize(child)].cas(child, parent) { - Ok(_) => return (parent, child), - Err(_) => continue, - } - } else { - return (l, l); - } - } - }, + |buf| merge_impl(buf, l, r), T::from_usize, ) } +} - fn find_impl(buf: &[T], elt: T::Underlying) -> T::Underlying { - macro_rules! load { - ($x:expr) => { - buf[T::as_usize($x)].load() - }; - } - let mut cur = elt; - let mut next = load!(cur); - let mut grand = load!(next); - while next != grand { - let _ = buf[T::as_usize(cur)].cas(next, grand); - // This is what the paper calls "two-try" splitting. - // next = load!(cur); - // grand = load!(next); - // let _ = buf[T::as_usize(cur)].cas(next, grand); - cur = next; - next = load!(cur); - grand = load!(next); +// The ordinary concurrent union-find uses acquire/release operations because +// its growable buffer has independent publication requirements. The scoped +// variants operate on an already initialized, fixed-capacity parent slice and +// therefore use relaxed operations. Their safety argument is documented on +// `AtomicUnionFindAccess`: links only move to smaller same-component ancestors, +// CAS validates every mutation, and the enclosing scope join publishes the +// completed phase before serial access resumes. +pub(crate) fn find_impl(buf: &[T], elt: T::Underlying) -> T::Underlying { + find_impl_with_order::(buf, elt) +} + +fn find_impl_with_order( + buf: &[T], + elt: T::Underlying, +) -> T::Underlying { + macro_rules! load { + ($x:expr) => { + load::(&buf[T::as_usize($x)]) + }; + } + let mut cur = elt; + let mut next = load!(cur); + let mut grand = load!(next); + while next != grand { + let _ = cas::(&buf[T::as_usize(cur)], next, grand); + // This is what the paper calls "two-try" splitting. + // next = load!(cur); + // grand = load!(next); + // let _ = buf[T::as_usize(cur)].cas(next, grand); + cur = next; + next = load!(cur); + grand = load!(next); + } + next +} + +pub(crate) fn merge_impl( + buf: &[T], + l: T::Underlying, + r: T::Underlying, +) -> ( + T::Underlying, /* parent */ + T::Underlying, /* child */ +) { + merge_impl_with_order::(buf, l, r) +} + +pub(crate) fn merge_scoped_impl( + buf: &[T], + l: T::Underlying, + r: T::Underlying, +) -> ( + T::Underlying, /* parent */ + T::Underlying, /* child */ +) { + merge_impl_with_order::(buf, l, r) +} + +fn merge_impl_with_order( + buf: &[T], + mut l: T::Underlying, + mut r: T::Underlying, +) -> ( + T::Underlying, /* parent */ + T::Underlying, /* child */ +) { + assert!( + T::as_usize(cmp::max(l, r)) < buf.len(), + "fixed union-find access is too small for the requested union" + ); + loop { + l = find_impl_with_order::(buf, l); + r = find_impl_with_order::(buf, r); + if l != r { + // We do "union by min": common in egraphs due to the + // hypothesis that smaller ids will be + // better-represented in the egraph, and hence + // perturbing the maximum id is likely to result in less + // work for rebuilding. + let parent = cmp::min(l, r); + let child = cmp::max(l, r); + match cas::(&buf[T::as_usize(child)], child, parent) { + Ok(_) => return (parent, child), + Err(_) => continue, + } + } else { + return (l, l); } - next + } +} + +#[inline] +fn load(value: &T) -> T::Underlying { + if SCOPED_RELAXED { + value.load_relaxed() + } else { + value.load() + } +} + +#[inline] +fn cas( + value: &T, + current: T::Underlying, + new: T::Underlying, +) -> Result { + if SCOPED_RELAXED { + value.cas_relaxed(current, new) + } else { + value.cas(current, new) } } diff --git a/union-find/src/lib.rs b/union-find/src/lib.rs index e92c3cc47..40e5910ed 100644 --- a/union-find/src/lib.rs +++ b/union-find/src/lib.rs @@ -11,8 +11,8 @@ //! closure. There's likely more to do in this area but for now it seems to work //! well enough. It doesn't hurt that it's also simpler to implement. use egglog_numeric_id as numeric_id; -use numeric_id::NumericId; -use std::cmp; +use numeric_id::{AtomicNumericId, NumericId}; +use std::{cmp, marker::PhantomData, mem, slice}; pub mod concurrent; @@ -25,6 +25,41 @@ pub struct UnionFind { parents: Vec, } +/// A concurrent, fixed-capacity view borrowed from a serial [`UnionFind`]. +/// +/// Creating this view requires exclusive access to the serial union-find. The +/// view may be shared with scoped worker threads, and the serial representation +/// becomes accessible again only after all of those borrows have ended. +/// +/// Operations on this fixed view use relaxed atomics. Parent entries are the +/// only shared state: a successful union changes a root from itself to a +/// strictly smaller ID, and path splitting changes a parent to a strictly +/// smaller ancestor. Consequently links never increase, cannot exhibit ABA, +/// and cannot form a cycle. A stale load is still a same-component ancestor, +/// while every mutation uses compare-exchange to verify the observed link. +/// The scope join before this view is dropped provides the happens-before edge +/// needed before ordinary reads resume; parent links do not publish any other +/// payload. +pub struct AtomicUnionFindAccess<'a, Value: AtomicNumericId> { + parents: &'a [Value::Atomic], + marker: PhantomData, +} + +impl AtomicUnionFindAccess<'_, Value> +where + Value::Atomic: concurrent::atomic_int::AtomicInt, +{ + /// Merge two equivalence classes. + /// + /// Panics if either ID exceeds the fixed capacity selected when this view + /// was created. + #[inline] + pub fn union(&self, l: Value, r: Value) -> (Value, Value) { + let (parent, child) = concurrent::uf::merge_scoped_impl(self.parents, l.rep(), r.rep()); + (Value::new(parent), Value::new(child)) + } +} + impl Default for UnionFind { fn default() -> Self { Self { @@ -102,3 +137,60 @@ impl UnionFind { cur } } + +impl UnionFind +where + Value::Atomic: concurrent::atomic_int::AtomicInt, +{ + /// Temporarily borrow this union-find as a fixed-capacity atomic view. + /// + /// The backing vector is grown before the atomic view is created and + /// cannot resize while `f` runs. Atomic and ordinary accesses never + /// overlap: the mutable receiver excludes ordinary readers until `f` and + /// every scoped worker borrowing its view have completed. No allocation, + /// resize, or per-operation access guard occurs while the view is active. + /// `max_id` is inclusive, and `f` must not request a union containing a + /// larger ID. + /// + /// # Panics + /// + /// Panics if the ID and its atomic integer do not have identical size and + /// alignment on the target platform. + pub fn with_atomic_access( + &mut self, + max_id: Value, + f: impl FnOnce(&AtomicUnionFindAccess<'_, Value>) -> R, + ) -> R { + self.reserve(max_id); + assert_eq!( + mem::size_of::(), + mem::size_of::(), + "numeric id and atomic representation must have equal size" + ); + assert_eq!( + mem::align_of::(), + mem::align_of::(), + "numeric id and atomic representation must have equal alignment" + ); + + // SAFETY: + // - `AtomicNumericId` guarantees a transparent integer + // representation compatible with the corresponding atomic; the + // assertions above defensively verify its size and alignment. + // - Integer bit patterns are valid atomic integer bit patterns. + // - `&mut self` excludes every ordinary access to `parents` for the + // lifetime of this slice. + // - `f` cannot let the borrowed view escape, so every atomic access + // ends before ordinary access resumes. + let parents = unsafe { + slice::from_raw_parts( + self.parents.as_mut_ptr().cast::(), + self.parents.len(), + ) + }; + f(&AtomicUnionFindAccess { + parents, + marker: PhantomData, + }) + } +} diff --git a/union-find/src/tests.rs b/union-find/src/tests.rs index d895dbf92..05cff51c9 100644 --- a/union-find/src/tests.rs +++ b/union-find/src/tests.rs @@ -1,4 +1,5 @@ use crate::numeric_id::{NumericId, define_id}; +use std::thread; use crate::UnionFind; @@ -38,3 +39,25 @@ fn basic_uf() { .all(|&id| uf.find(id) == target) ); } + +#[test] +fn serial_union_find_supports_scoped_atomic_updates() { + let mut uf = UnionFind::::default(); + let n = 10_000; + uf.with_atomic_access(v(n), |access| { + thread::scope(|scope| { + for thread_id in 0..8 { + scope.spawn(move || { + for i in (thread_id + 1..n).step_by(8) { + access.union(v(i - 1), v(i)); + } + }); + } + }); + }); + + for i in 0..n { + assert_eq!(uf.find_naive(v(i)), v(0)); + } + assert_eq!(uf.find_naive(v(n + 1)), v(n + 1)); +} From 5ec34b7bf76d46a305b799a34fb3526848aaec5a Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Sun, 26 Jul 2026 23:07:52 -0700 Subject: [PATCH 12/12] Use packed cached hashes for table mutations --- core-relations/src/row_buffer/mod.rs | 3 - core-relations/src/row_buffer/partitioned.rs | 374 ------------ core-relations/src/table/mod.rs | 609 ++++++++++--------- core-relations/src/table/tests.rs | 194 +++++- 4 files changed, 493 insertions(+), 687 deletions(-) delete mode 100644 core-relations/src/row_buffer/partitioned.rs diff --git a/core-relations/src/row_buffer/mod.rs b/core-relations/src/row_buffer/mod.rs index f6ae7c977..71bb756f5 100644 --- a/core-relations/src/row_buffer/mod.rs +++ b/core-relations/src/row_buffer/mod.rs @@ -13,12 +13,9 @@ use crate::{ pool::{Pooled, with_pool_set}, }; -mod partitioned; #[cfg(test)] mod tests; -pub(crate) use partitioned::{HashPartitioning, PartitionedRowBuffer}; - /// A trait for types that can store a vector of `Value`s. /// /// Can either be backed by a pool or SmallVec. diff --git a/core-relations/src/row_buffer/partitioned.rs b/core-relations/src/row_buffer/partitioned.rs deleted file mode 100644 index 8535099af..000000000 --- a/core-relations/src/row_buffer/partitioned.rs +++ /dev/null @@ -1,374 +0,0 @@ -//! Hash-partitioned batches of rows. -//! -//! Producers append rows, unchanged full hashes, and their destination -//! partition densely. Before publication the buffer is sealed with a stable -//! counting partition, giving consumers one contiguous sequence per partition. - -use std::iter::FusedIterator; - -use crate::{common::Value, numeric_id::NumericId}; - -/// Selects a power-of-two hash partition within one of several outer groups. -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub(crate) struct HashPartitioning { - shift: u8, - bits: u8, - groups: u32, -} - -impl HashPartitioning { - #[cfg(test)] - pub(crate) fn new(shift: u32, bits: u32) -> Self { - Self::grouped(shift, bits, 1) - } - - pub(crate) fn grouped(shift: u32, bits: u32, groups: usize) -> Self { - assert!( - shift.checked_add(bits).is_some_and(|sum| sum <= 64), - "hash partition bits must fit in a u64" - ); - assert!( - bits < usize::BITS, - "hash partition count must fit in a usize" - ); - assert_ne!(groups, 0, "hash partitioning must have at least one group"); - let groups = u32::try_from(groups).expect("hash partition group count must fit in u32"); - (1usize << bits) - .checked_mul(groups as usize) - .expect("total hash partition count overflow"); - Self { - shift: shift as u8, - bits: bits as u8, - groups, - } - } - - pub(crate) fn partitions_per_group(self) -> usize { - 1usize << self.bits - } - - pub(crate) fn partition_count(self) -> usize { - self.partitions_per_group() * self.groups as usize - } - - #[inline] - pub(crate) fn partition(self, group: usize, hash: u64) -> usize { - assert!( - group < self.groups as usize, - "hash partition group out of bounds" - ); - let within_group = if self.bits == 0 { - 0 - } else { - (hash.wrapping_shr(self.shift as u32) & (self.partitions_per_group() as u64 - 1)) - as usize - }; - self.partition_index(group, within_group) - } - - pub(crate) fn partition_index(self, group: usize, within_group: usize) -> usize { - assert!( - group < self.groups as usize, - "hash partition group out of bounds" - ); - assert!( - within_group < self.partitions_per_group(), - "hash subpartition out of bounds" - ); - group * self.partitions_per_group() + within_group - } -} - -/// One row yielded from a [`PartitionedRowBuffer`]. -pub(crate) struct HashedRow<'a> { - pub(crate) hash: u64, - pub(crate) row: &'a [Value], -} - -/// A stable hash-partitioned row batch. -/// -/// Before sealing, all vectors are dense and insertion ordered. -/// `partition_ids[i]` describes `hashes[i]` and row `i`. Sealing performs a -/// stable counting partition and replaces `partition_ids` with prefix offsets. -#[derive(Clone)] -pub(crate) struct PartitionedRowBuffer { - arity: usize, - partitioning: HashPartitioning, - hashes: Vec, - rows: Vec, - partition_ids: Vec, - offsets: Option>, -} - -impl PartitionedRowBuffer { - pub(crate) fn new(arity: usize, partitioning: HashPartitioning) -> Self { - Self { - arity, - partitioning, - hashes: Vec::new(), - rows: Vec::new(), - partition_ids: Vec::new(), - offsets: None, - } - } - - pub(crate) fn arity(&self) -> usize { - self.arity - } - - pub(crate) fn len(&self) -> usize { - self.hashes.len() - } - - #[cfg(test)] - pub(crate) fn is_empty(&self) -> bool { - self.hashes.is_empty() - } - - pub(crate) fn partition_count(&self) -> usize { - self.partitioning.partition_count() - } - - /// Append a row, its outer group, and its unchanged full hash. - pub(crate) fn add_row(&mut self, group: usize, hash: u64, row: &[Value]) { - assert!( - self.offsets.is_none(), - "cannot append to a sealed partitioned row buffer" - ); - assert_eq!( - row.len(), - self.arity, - "attempting to add a row with mismatched arity" - ); - let partition = self.partitioning.partition(group, hash); - let partition = - u32::try_from(partition).expect("partitioned row buffer partition id overflow"); - self.hashes.push(hash); - self.rows.extend_from_slice(row); - self.partition_ids.push(partition); - } - - /// Convert insertion-ordered staging vectors into contiguous stable - /// partition runs. - pub(crate) fn seal(mut self) -> Self { - if self.offsets.is_some() { - return self; - } - debug_assert_eq!(self.partition_ids.len(), self.hashes.len()); - debug_assert_eq!(self.rows.len(), self.hashes.len() * self.arity); - - if self.partition_count() == 1 { - let len = u32::try_from(self.len()).expect("partitioned row buffer row count overflow"); - self.partition_ids.clear(); - self.offsets = Some(Box::new([0, len])); - return self; - } - - let mut offsets = vec![0u32; self.partition_count() + 1]; - for &partition in &self.partition_ids { - offsets[partition as usize + 1] = offsets[partition as usize + 1] - .checked_add(1) - .expect("partitioned row count overflow"); - } - for partition in 0..self.partition_count() { - offsets[partition + 1] = offsets[partition + 1] - .checked_add(offsets[partition]) - .expect("partitioned row offset overflow"); - } - debug_assert_eq!(offsets[self.partition_count()] as usize, self.len()); - - let mut cursors = offsets[..self.partition_count()].to_vec(); - let mut hashes = vec![0u64; self.len()]; - let row_values = self - .len() - .checked_mul(self.arity) - .expect("partitioned row value count overflow"); - let mut rows = vec![Value::new(0); row_values]; - for source in 0..self.len() { - let partition = self.partition_ids[source] as usize; - let destination = cursors[partition] as usize; - cursors[partition] += 1; - hashes[destination] = self.hashes[source]; - if self.arity != 0 { - rows[destination * self.arity..(destination + 1) * self.arity] - .copy_from_slice(&self.rows[source * self.arity..(source + 1) * self.arity]); - } - } - - self.hashes = hashes; - self.rows = rows; - self.partition_ids = Vec::new(); - self.offsets = Some(offsets.into_boxed_slice()); - self - } - - #[cfg(test)] - pub(crate) fn partition_len(&self, partition: usize) -> usize { - let offsets = self.offsets(); - (offsets[partition + 1] - offsets[partition]) as usize - } - - pub(crate) fn group_is_empty(&self, group: usize) -> bool { - self.group_len(group) == 0 - } - - pub(crate) fn group_len(&self, group: usize) -> usize { - let first = self.partitioning.partition_index(group, 0); - let after_last = first + self.partitioning.partitions_per_group(); - let offsets = self.offsets(); - (offsets[after_last] - offsets[first]) as usize - } - - pub(crate) fn partition(&self, partition: usize) -> PartitionIter<'_> { - let offsets = self.offsets(); - PartitionIter { - buffer: self, - next: offsets[partition] as usize, - end: offsets[partition + 1] as usize, - } - } - - fn offsets(&self) -> &[u32] { - self.offsets - .as_deref() - .expect("partitioned row buffer must be sealed before reading") - } -} - -pub(crate) struct PartitionIter<'a> { - buffer: &'a PartitionedRowBuffer, - next: usize, - end: usize, -} - -impl<'a> Iterator for PartitionIter<'a> { - type Item = HashedRow<'a>; - - #[inline] - fn next(&mut self) -> Option { - if self.next == self.end { - return None; - } - let index = self.next; - self.next += 1; - let row = if self.buffer.arity == 0 { - &[] - } else { - &self.buffer.rows[index * self.buffer.arity..(index + 1) * self.buffer.arity] - }; - Some(HashedRow { - hash: self.buffer.hashes[index], - row, - }) - } - - fn size_hint(&self) -> (usize, Option) { - let remaining = self.end - self.next; - (remaining, Some(remaining)) - } -} - -impl ExactSizeIterator for PartitionIter<'_> {} -impl FusedIterator for PartitionIter<'_> {} - -#[cfg(test)] -mod tests { - use super::*; - - fn v(value: usize) -> Value { - Value::from_usize(value) - } - - #[test] - fn partitions_rows_stably_and_keeps_hashes_aligned() { - let partitioning = HashPartitioning::new(4, 2); - let mut buffer = PartitionedRowBuffer::new(2, partitioning); - let mut expected = [Vec::new(), Vec::new(), Vec::new(), Vec::new()]; - for i in 0..80usize { - let partition = (i * 3) & 3; - let hash = ((partition as u64) << 4) | (i as u64 & 0xf); - let row = [v(i), v(i + 1)]; - buffer.add_row(0, hash, &row); - expected[partition].push((hash, row)); - } - let buffer = buffer.seal(); - - assert_eq!(buffer.len(), 80); - assert_eq!(buffer.arity(), 2); - for (partition, expected) in expected.iter().enumerate() { - assert_eq!(buffer.partition_len(partition), expected.len()); - let mut iter = buffer.partition(partition); - assert_eq!(iter.len(), expected.len()); - let actual = iter - .by_ref() - .map(|hashed| (hashed.hash, hashed.row.to_vec())) - .collect::>(); - assert_eq!(iter.len(), 0); - let expected = expected - .iter() - .map(|(hash, row)| (*hash, row.to_vec())) - .collect::>(); - assert_eq!(actual, expected); - } - } - - #[test] - fn supports_groups_zero_arity_and_logical_clone() { - let partitioning = HashPartitioning::grouped(1, 3, 2); - let mut buffer = PartitionedRowBuffer::new(0, partitioning); - for (group, hash) in [(0, 0), (1, 2), (0, 4), (1, 2), (1, 14), (0, 0)] { - buffer.add_row(group, hash, &[]); - } - let buffer = buffer.seal(); - let cloned = buffer.clone(); - - assert_eq!(cloned.len(), 6); - assert!(!cloned.is_empty()); - assert_eq!(cloned.group_len(0), 3); - assert_eq!(cloned.group_len(1), 3); - for partition in 0..partitioning.partition_count() { - let original = buffer - .partition(partition) - .map(|hashed| (hashed.hash, hashed.row.len())) - .collect::>(); - let copied = cloned - .partition(partition) - .map(|hashed| (hashed.hash, hashed.row.len())) - .collect::>(); - assert_eq!(copied, original); - } - } - - #[test] - fn is_send_and_sync() { - fn assert_send_sync() {} - assert_send_sync::(); - } - - #[test] - fn one_partition_seals_without_copying_rows_or_hashes() { - let partitioning = HashPartitioning::new(0, 0); - let mut buffer = PartitionedRowBuffer::new(2, partitioning); - for i in 0..32usize { - buffer.add_row(0, i as u64, &[v(i), v(i + 1)]); - } - let hashes = buffer.hashes.as_ptr(); - let rows = buffer.rows.as_ptr(); - - let buffer = buffer.seal(); - - assert_eq!(buffer.hashes.as_ptr(), hashes); - assert_eq!(buffer.rows.as_ptr(), rows); - assert!(buffer.partition_ids.is_empty()); - assert_eq!(buffer.partition_len(0), 32); - assert_eq!( - buffer - .partition(0) - .map(|hashed| (hashed.hash, hashed.row.to_vec())) - .collect::>(), - (0..32usize) - .map(|i| (i as u64, vec![v(i), v(i + 1)])) - .collect::>() - ); - } -} diff --git a/core-relations/src/table/mod.rs b/core-relations/src/table/mod.rs index f710136f7..5d078751d 100644 --- a/core-relations/src/table/mod.rs +++ b/core-relations/src/table/mod.rs @@ -30,7 +30,7 @@ use crate::{ parallel, parallel_heuristics::parallelize_table_op, pool::with_pool_set, - row_buffer::{HashPartitioning, ParallelRowBufWriter, PartitionedRowBuffer, RowBuffer}, + row_buffer::{ParallelRowBufWriter, RowBuffer}, table_spec::{ ColumnId, Constraint, Generation, MutationBuffer, Offset, Row, Table, TableSpec, TableVersion, @@ -42,35 +42,154 @@ mod sharded_hash_table; #[cfg(test)] mod tests; -// NB: Having this type def lets us switch between 64 and 32 bits of hashcode. -// -// We should consider just using u64 everywhere though. Hashbrown doesn't play nicely with 32-bit -// hashcodes because it uses both the high and low bits of a 64-bit code. +const PARALLEL_INSERT_BATCH_SIZE: usize = 1 << 12; -type HashCode = u64; +/// The complete hash used to choose a physical shard. +/// +/// Only this representation may be passed to [`ShardData::shard_id`]. The +/// compact hash deliberately omits the shard-selection bits. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +#[repr(transparent)] +struct FullHash(u64); -// The lowest bits select buckets within a cache-sized destination-table -// window. The next bits group writes to those windows before a shard owner -// probes its HashTable. -const MUTATION_CACHE_WINDOW_BITS: u32 = 11; -const MUTATION_PARTITION_BITS: u32 = 6; -const PARALLEL_INSERT_BATCH_SIZE: usize = 1 << 12; +/// The 32-bit hash fingerprint cached beside rows and table entries. +/// +/// Bits 0..25 preserve the full hash's low bucket-index bits, while bits +/// 25..32 preserve hashbrown's exact seven-bit H2 tag from the top of the +/// target's effective hash word. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +#[repr(transparent)] +struct CompactHash(u32); + +impl CompactHash { + const LOW_BITS: u32 = 25; + const LOW_MASK: u64 = (1 << Self::LOW_BITS) - 1; + const H2_BITS: u32 = 7; + const H2_MASK: u64 = (1 << Self::H2_BITS) - 1; + const H2_SHIFT: u32 = if usize::BITS < u64::BITS { + usize::BITS - Self::H2_BITS + } else { + u64::BITS - Self::H2_BITS + }; + + #[inline] + fn from_full(full: FullHash) -> Self { + let low = full.0 & Self::LOW_MASK; + let h2 = (full.0 >> Self::H2_SHIFT) & Self::H2_MASK; + Self((low | (h2 << Self::LOW_BITS)) as u32) + } + + #[inline] + fn from_value(value: Value) -> Self { + Self(value.rep()) + } + + #[inline] + fn as_value(self) -> Value { + Value::new(self.0) + } + + #[inline] + fn probe(self) -> ProbeHash { + let low = self.0 as u64; + let h2 = low >> Self::LOW_BITS; + ProbeHash(low | (h2 << Self::H2_SHIFT)) + } +} + +/// A reconstructed 64-bit hash used only at hashbrown's raw API boundary. +/// +/// This has the cached low bucket bits and exact H2 tag, but it does not +/// contain the bits needed to recover a physical shard. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +#[repr(transparent)] +struct ProbeHash(u64); + +impl ProbeHash { + #[inline] + fn raw(self) -> u64 { + self.0 + } +} + +/// The two independent products of hashing a row. +/// +/// `shard` is computed from the full hash before compaction. `compact` is the +/// fingerprint cached after the row has already been routed to that shard. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +struct ShardHash { + shard: ShardId, + compact: CompactHash, +} + +impl ShardHash { + #[inline] + fn from_full(shard_data: ShardData, full: FullHash) -> Self { + Self { + shard: shard_data.shard_id(full.0), + compact: CompactHash::from_full(full), + } + } +} /// A pointer to a row in the table. #[derive(Clone, Debug)] +#[repr(C)] pub(crate) struct TableEntry { - hashcode: HashCode, + hash: CompactHash, row: RowId, } impl TableEntry { - fn hashcode(&self) -> u64 { - // We keep the cast here to make it easy to switch to HashCode=u32. - #[allow(clippy::unnecessary_cast)] - { - self.hashcode as u64 + /// Adapter for hashbrown callbacks, which require the raw `u64`. + fn raw_probe_hash(&self) -> u64 { + self.hash.probe().raw() + } +} + +/// Producer-local rows for one physical hash-table shard. +/// +/// This preserves producer order and performs no seal/scatter step. A compact +/// hash is stored as a trailing `Value` in each physical row, matching +/// `TaggedRowBuffer`'s one-allocation layout. +#[derive(Clone)] +struct HashedRowBuffer { + arity: usize, + values: Vec, +} + +impl HashedRowBuffer { + fn new(arity: usize) -> Self { + Self { + arity, + values: Vec::new(), } } + + fn add_row(&mut self, hash: CompactHash, row: &[Value]) { + assert_eq!(row.len(), self.arity); + self.values.extend_from_slice(row); + self.values.push(hash.as_value()); + } + + fn len(&self) -> usize { + self.values.len() / (self.arity + 1) + } + + fn rows_hashed(&self) -> impl Iterator { + debug_assert_eq!(self.values.len() % (self.arity + 1), 0); + self.values.chunks_exact(self.arity + 1).map(|tagged| { + let hash = CompactHash::from_value(tagged[self.arity]); + (hash, &tagged[..self.arity]) + }) + } +} + +#[derive(Clone, Copy, Debug)] +#[repr(C)] +struct KnownRemoval { + hash: CompactHash, + row: RowId, } /// The core data for a table. @@ -178,100 +297,86 @@ impl Clone for SortedWritesTable { } struct Buffer { - pending_rows: Option, - pending_removals: Option, - pending_known_removals: Option, + pending_rows: DenseIdMap, + pending_removals: DenseIdMap, + pending_known_removals: DenseIdMap>, state: Weak, n_cols: u32, n_keys: u32, shard_data: ShardData, - insert_partitioning: HashPartitioning, - removal_partitioning: HashPartitioning, } impl MutationBuffer for Buffer { fn stage_insert(&mut self, row: &[Value]) { - let (shard, hash) = hash_code(self.shard_data, row, self.n_keys as _); + let hash = shard_hash(self.shard_data, row, self.n_keys as _); self.pending_rows - .get_or_insert_with(|| { - PartitionedRowBuffer::new(self.n_cols as _, self.insert_partitioning) - }) - .add_row(shard.index(), hash, row); + .get_or_insert(hash.shard, || HashedRowBuffer::new(self.n_cols as _)) + .add_row(hash.compact, row); } fn stage_remove(&mut self, key: &[Value]) { - let (shard, hash) = hash_code(self.shard_data, key, self.n_keys as _); + let hash = shard_hash(self.shard_data, key, self.n_keys as _); self.pending_removals - .get_or_insert_with(|| { - PartitionedRowBuffer::new(self.n_keys as _, self.removal_partitioning) - }) - .add_row(shard.index(), hash, key); + .get_or_insert(hash.shard, || HashedRowBuffer::new(self.n_keys as _)) + .add_row(hash.compact, key); } fn fresh_handle(&self) -> Box { Box::new(Buffer { - pending_rows: None, - pending_removals: None, - pending_known_removals: None, + pending_rows: DenseIdMap::with_capacity(self.shard_data.n_shards()), + pending_removals: DenseIdMap::with_capacity(self.shard_data.n_shards()), + pending_known_removals: DenseIdMap::with_capacity(self.shard_data.n_shards()), state: self.state.clone(), n_cols: self.n_cols, n_keys: self.n_keys, shard_data: self.shard_data, - insert_partitioning: self.insert_partitioning, - removal_partitioning: self.removal_partitioning, }) } } impl Buffer { fn stage_remove_row(&mut self, row: RowId, key: &[Value]) { - let (shard, hashcode) = hash_code(self.shard_data, key, self.n_keys as _); + let hash = shard_hash(self.shard_data, key, self.n_keys as _); self.pending_known_removals - .get_or_insert_with(|| PartitionedRowBuffer::new(1, self.removal_partitioning)) - .add_row(shard.index(), hashcode, &[Value::new(row.rep())]); + .get_or_insert(hash.shard, Vec::new) + .push(KnownRemoval { + hash: hash.compact, + row, + }); } } impl Drop for Buffer { fn drop(&mut self) { if let Some(state) = self.state.upgrade() { - let rows = publish_partitioned( - self.pending_rows.take(), - &state.pending_rows, - self.shard_data, - ); + let mut rows = 0; + for shard_index in 0..self.pending_rows.n_ids() { + let shard = ShardId::from_usize(shard_index); + if let Some(buffer) = self.pending_rows.take(shard) { + rows += buffer.len(); + state.pending_rows[shard].push(buffer); + } + } state.total_rows.fetch_add(rows, Ordering::Relaxed); - let removals = publish_partitioned( - self.pending_removals.take(), - &state.pending_removals, - self.shard_data, - ) + publish_partitioned( - self.pending_known_removals.take(), - &state.pending_known_removals, - self.shard_data, - ); + let mut removals = 0; + for shard_index in 0..self.pending_removals.n_ids() { + let shard = ShardId::from_usize(shard_index); + if let Some(buffer) = self.pending_removals.take(shard) { + removals += buffer.len(); + state.pending_removals[shard].push(buffer); + } + } + for shard_index in 0..self.pending_known_removals.n_ids() { + let shard = ShardId::from_usize(shard_index); + if let Some(buffer) = self.pending_known_removals.take(shard) { + removals += buffer.len(); + state.pending_known_removals[shard].push(buffer); + } + } state.total_removals.fetch_add(removals, Ordering::Relaxed); } } } -fn publish_partitioned( - buffer: Option, - queues: &DenseIdMap>>, - shard_data: ShardData, -) -> usize { - let Some(buffer) = buffer else { - return 0; - }; - let len = buffer.len(); - let buffer = Arc::new(buffer.seal()); - for shard_index in 0..shard_data.n_shards() { - if !buffer.group_is_empty(shard_index) { - queues[ShardId::from_usize(shard_index)].push(buffer.clone()); - } - } - len -} - impl Table for SortedWritesTable { fn dyn_clone(&self) -> Box { Box::new(self.clone()) @@ -551,16 +656,15 @@ impl Table for SortedWritesTable { impl SortedWritesTable { fn new_table_buffer(&self) -> Buffer { + let n_shards = self.hash.shard_data().n_shards(); Buffer { - pending_rows: None, - pending_removals: None, - pending_known_removals: None, + pending_rows: DenseIdMap::with_capacity(n_shards), + pending_removals: DenseIdMap::with_capacity(n_shards), + pending_known_removals: DenseIdMap::with_capacity(n_shards), state: Arc::downgrade(&self.pending_state), n_keys: u32::try_from(self.n_keys).expect("n_keys should fit in u32"), n_cols: u32::try_from(self.n_columns).expect("n_columns should fit in u32"), shard_data: self.hash.shard_data(), - insert_partitioning: self.pending_state.insert_partitioning, - removal_partitioning: self.pending_state.removal_partitioning, } } @@ -593,7 +697,7 @@ impl SortedWritesTable { n_columns, sort_by, offsets: Default::default(), - pending_state: Arc::new(PendingState::new(shard_data, sort_by.is_some())), + pending_state: Arc::new(PendingState::new(shard_data)), merge: merge_fn.into(), to_rebuild, rebuild_index, @@ -605,11 +709,8 @@ impl SortedWritesTable { /// probes and erases all entries in its partition, then marks the removed /// rows stale as a separate write-only pass. fn parallel_delete(&mut self) -> bool { - let shard_data = self.hash.shard_data(); let pending_removals = &self.pending_state.pending_removals; let pending_known_removals = &self.pending_state.pending_known_removals; - let partitioning = self.pending_state.removal_partitioning; - let partition_count = partitioning.partitions_per_group(); let data = &self.data.data; let n_keys = self.n_keys; let stale_delta: usize = @@ -618,41 +719,28 @@ impl SortedWritesTable { for (shard_offset, shard) in shards.iter_mut().enumerate() { let shard_id = ShardId::from_usize(base_shard + shard_offset); let removal_buffers = drain_queue(&pending_removals[shard_id]); - for partition in 0..partition_count { - for buf in &removal_buffers { - for hashed in buf.partition( - partitioning.partition_index(shard_id.index(), partition), - ) { - let to_remove = hashed.row; - let hc = hashed.hash; - debug_assert_eq!(shard_data.shard_id(hc), shard_id); - debug_assert_eq!(buf.arity(), n_keys); - if let Ok(entry) = shard.find_entry(hc, |entry| { - entry.hashcode == (hc as HashCode) - && &data.get_row(entry.row)[0..n_keys] == to_remove - }) { - let (ent, _) = entry.remove(); - removed_rows.push(ent.row); - } + for buf in &removal_buffers { + debug_assert_eq!(buf.arity, n_keys); + for (hash, to_remove) in buf.rows_hashed() { + if let Ok(entry) = shard.find_entry(hash.probe().raw(), |entry| { + entry.hash == hash + && &data.get_row(entry.row)[0..n_keys] == to_remove + }) { + let (ent, _) = entry.remove(); + removed_rows.push(ent.row); } } } let known_buffers = drain_queue(&pending_known_removals[shard_id]); - for partition in 0..partition_count { - for buf in &known_buffers { - for hashed in buf.partition( - partitioning.partition_index(shard_id.index(), partition), - ) { - let hc = hashed.hash; - let row = RowId::new(hashed.row[0].rep()); - debug_assert_eq!(shard_data.shard_id(hc), shard_id); - debug_assert_eq!(buf.arity(), 1); - if let Ok(entry) = shard.find_entry(hc, |entry| { - entry.hashcode == hc as HashCode && entry.row == row - }) { - let (ent, _) = entry.remove(); - removed_rows.push(ent.row); - } + for buf in &known_buffers { + for removal in buf { + if let Ok(entry) = shard + .find_entry(removal.hash.probe().raw(), |entry| { + entry.hash == removal.hash && entry.row == removal.row + }) + { + let (ent, _) = entry.remove(); + removed_rows.push(ent.row); } } } @@ -669,9 +757,6 @@ impl SortedWritesTable { stale_delta > 0 } fn serial_delete(&mut self) -> bool { - let shard_data = self.hash.shard_data(); - let partitioning = self.pending_state.removal_partitioning; - let partition_count = partitioning.partitions_per_group(); let mut changed = false; self.hash .mut_shards() @@ -680,45 +765,30 @@ impl SortedWritesTable { .for_each(|(shard_id, shard)| { let shard_id = ShardId::from_usize(shard_id); let removal_buffers = drain_queue(&self.pending_state.pending_removals[shard_id]); - for partition in 0..partition_count { - for buf in &removal_buffers { - for hashed in - buf.partition(partitioning.partition_index(shard_id.index(), partition)) - { - let to_remove = hashed.row; - let hc = hashed.hash; - debug_assert_eq!(shard_data.shard_id(hc), shard_id); - debug_assert_eq!(buf.arity(), self.n_keys); - if let Ok(entry) = shard.find_entry(hc, |entry| { - entry.hashcode == (hc as HashCode) - && &self.data.get_row(entry.row).unwrap()[0..self.n_keys] - == to_remove - }) { - let (ent, _) = entry.remove(); - self.data.set_stale(ent.row); - changed = true; - } + for buf in &removal_buffers { + debug_assert_eq!(buf.arity, self.n_keys); + for (hash, to_remove) in buf.rows_hashed() { + if let Ok(entry) = shard.find_entry(hash.probe().raw(), |entry| { + entry.hash == hash + && &self.data.get_row(entry.row).unwrap()[0..self.n_keys] + == to_remove + }) { + let (ent, _) = entry.remove(); + self.data.set_stale(ent.row); + changed = true; } } } let known_buffers = drain_queue(&self.pending_state.pending_known_removals[shard_id]); - for partition in 0..partition_count { - for buf in &known_buffers { - for hashed in - buf.partition(partitioning.partition_index(shard_id.index(), partition)) - { - let hc = hashed.hash; - let row = RowId::new(hashed.row[0].rep()); - debug_assert_eq!(shard_data.shard_id(hc), shard_id); - debug_assert_eq!(buf.arity(), 1); - if let Ok(entry) = shard.find_entry(hc, |entry| { - entry.hashcode == hc as HashCode && entry.row == row - }) { - let (ent, _) = entry.remove(); - self.data.set_stale(ent.row); - changed = true; - } + for buf in &known_buffers { + for removal in buf { + if let Ok(entry) = shard.find_entry(removal.hash.probe().raw(), |entry| { + entry.hash == removal.hash && entry.row == removal.row + }) { + let (ent, _) = entry.remove(); + self.data.set_stale(ent.row); + changed = true; } } } @@ -760,69 +830,55 @@ impl SortedWritesTable { fn serial_insert(&mut self, exec_state: &mut ExecutionState) -> bool { let mut changed = false; let n_keys = self.n_keys; - let partitioning = self.pending_state.insert_partitioning; - let partition_count = partitioning.partitions_per_group(); let mut scratch = with_pool_set(|ps| ps.get::>()); for (outer_shard, queue) in self.pending_state.pending_rows.iter() { let buffers = drain_queue(queue); - let incoming = buffers - .iter() - .map(|buffer| buffer.group_len(outer_shard.index())) - .sum(); - self.hash.mut_shards()[outer_shard.index()].reserve(incoming, TableEntry::hashcode); - for partition in 0..partition_count { - for buf in &buffers { - for hashed in - buf.partition(partitioning.partition_index(outer_shard.index(), partition)) - { - let query = hashed.row; - if query.first().is_some_and(Value::is_stale) { - continue; - } - let hc = hashed.hash; - debug_assert_eq!(self.hash.shard_data().shard_id(hc), outer_shard); - let key = &query[0..n_keys]; - use hashbrown::hash_table::Entry; - match self.hash.mut_shards()[outer_shard.index()].entry( - hc, - |entry| { - entry.hashcode == hc as HashCode - && self - .data - .get_row(entry.row) - .is_some_and(|row| &row[0..n_keys] == key) - }, - TableEntry::hashcode, - ) { - Entry::Occupied(mut entry) => { - let old_row = entry.get().row; - let cur = self + let incoming = buffers.iter().map(HashedRowBuffer::len).sum(); + self.hash.mut_shards()[outer_shard.index()] + .reserve(incoming, TableEntry::raw_probe_hash); + for buf in &buffers { + for (hash, query) in buf.rows_hashed() { + if query.first().is_some_and(Value::is_stale) { + continue; + } + let key = &query[0..n_keys]; + use hashbrown::hash_table::Entry; + match self.hash.mut_shards()[outer_shard.index()].entry( + hash.probe().raw(), + |entry| { + entry.hash == hash + && self .data - .get_row(old_row) - .expect("table should not point to stale entry"); - if (self.merge)(exec_state, cur, query, &mut scratch) { - debug_assert_eq!(&scratch[0..n_keys], key); - let new = self.data.add_row(&scratch); - if let Some(sort_by) = self.sort_by { - update_sort_offsets(sort_by, query, new, &mut self.offsets); - } - self.data.set_stale(old_row); - entry.get_mut().row = new; - changed = true; - } - scratch.clear(); - } - Entry::Vacant(entry) => { - let new = self.data.add_row(query); + .get_row(entry.row) + .is_some_and(|row| &row[0..n_keys] == key) + }, + TableEntry::raw_probe_hash, + ) { + Entry::Occupied(mut entry) => { + let old_row = entry.get().row; + let cur = self + .data + .get_row(old_row) + .expect("table should not point to stale entry"); + if (self.merge)(exec_state, cur, query, &mut scratch) { + debug_assert_eq!(&scratch[0..n_keys], key); + let new = self.data.add_row(&scratch); if let Some(sort_by) = self.sort_by { update_sort_offsets(sort_by, query, new, &mut self.offsets); } - entry.insert(TableEntry { - hashcode: hc as _, - row: new, - }); + self.data.set_stale(old_row); + entry.get_mut().row = new; changed = true; } + scratch.clear(); + } + Entry::Vacant(entry) => { + let new = self.data.add_row(query); + if let Some(sort_by) = self.sort_by { + update_sort_offsets(sort_by, query, new, &mut self.offsets); + } + entry.insert(TableEntry { hash, row: new }); + changed = true; } } } @@ -840,14 +896,11 @@ impl SortedWritesTable { // pre-sharded, and one logical thread can process updates for each // shard independently. Updates happen in three phases, which comments // describe below. - let shard_data = self.hash.shard_data(); let n_keys = self.n_keys; let n_cols = self.n_columns; let next_offset = RowId::from_usize(self.data.data.len()); let row_writer = self.data.data.parallel_writer(); let pending_rows = &self.pending_state.pending_rows; - let partitioning = self.pending_state.insert_partitioning; - let partition_count = partitioning.partitions_per_group(); let merge = self.merge.clone(); let pending_adds = parallel::map_mut(self.hash.mut_shards(), |shard_id, shard| { let shard_id = ShardId::from_usize(shard_id); @@ -856,16 +909,12 @@ impl SortedWritesTable { let mut scratch = with_pool_set(|ps| ps.get::>()); let queue = &pending_rows[shard_id]; let buffers = drain_queue(queue); - let incoming = buffers - .iter() - .map(|buffer| buffer.group_len(shard_id.index())) - .sum(); - shard.reserve(incoming, TableEntry::hashcode); + let incoming = buffers.iter().map(HashedRowBuffer::len).sum(); + shard.reserve(incoming, TableEntry::raw_probe_hash); let mut marked_stale = 0usize; let mut staged = CoalescedInsertBatch::new(n_keys, n_cols, PARALLEL_INSERT_BATCH_SIZE); let mut changed = false; - // Flush after a bounded number of staged physical rows or at a - // destination-table partition boundary. + // Flush after a bounded number of staged physical rows. macro_rules! flush_insert_batch { () => {{ if !staged.is_empty() { @@ -883,7 +932,7 @@ impl SortedWritesTable { // contention. let mut cur_row = start_row; let read_handle = row_writer.read_handle(); - for (hc, row) in staged.rows_hashed() { + for (hash, row) in staged.rows_hashed() { if row.first().is_some_and(Value::is_stale) { cur_row = cur_row.inc(); continue; @@ -898,15 +947,14 @@ impl SortedWritesTable { assert_eq!(read_handle.get_row_unchecked(cur_row), row); } } - debug_assert_eq!(shard_data.shard_id(hc), shard_id); match shard.entry( - hc, + hash.probe().raw(), // SAFETY: `ent` must point to a valid row |ent| unsafe { - ent.hashcode == hc as HashCode + ent.hash == hash && &read_handle.get_row_unchecked(ent.row)[0..n_keys] == key }, - TableEntry::hashcode, + TableEntry::raw_probe_hash, ) { Entry::Occupied(mut occ) => { // SAFETY: `occ` must point to a valid row: we only insert valid rows @@ -957,7 +1005,7 @@ impl SortedWritesTable { checker.check_local(row); changed = true; v.insert(TableEntry { - hashcode: hc as HashCode, + hash, row: cur_row, }); } @@ -975,22 +1023,17 @@ impl SortedWritesTable { // reservation. // * Apply each live row to the exclusively owned destination // shard in `flush_insert_batch`. - for partition in 0..partition_count { - for buf in &buffers { - for hashed in - buf.partition(partitioning.partition_index(shard_id.index(), partition)) - { - staged.insert_hashed(hashed.hash, hashed.row, |cur, new, out| { - (merge)(&mut exec_state, cur, new, out) - }); - if staged.physical_len() >= PARALLEL_INSERT_BATCH_SIZE { - flush_insert_batch!(); - } + for buf in &buffers { + for (hash, row) in buf.rows_hashed() { + staged.insert_hashed(hash, row, |cur, new, out| { + (merge)(&mut exec_state, cur, new, out) + }); + if staged.physical_len() >= PARALLEL_INSERT_BATCH_SIZE { + flush_insert_batch!(); } } - // Do not carry rows into the next destination-table window. - flush_insert_batch!(); } + flush_insert_batch!(); (checker, marked_stale, changed) }); self.data.data = row_writer.finish(); @@ -1127,8 +1170,8 @@ impl SortedWritesTable { while cur_row < $end_row { if let Some(row) = self.data.get_row(cur_row) { count += 1; - let (shard, _) = hash_code(self.hash.shard_data(), row, self.n_keys); - *histogram.get_or_default(shard) += 1; + let hash = shard_hash(self.hash.shard_data(), row, self.n_keys); + *histogram.get_or_default(hash.shard) += 1; } cur_row = cur_row.inc(); } @@ -1351,11 +1394,11 @@ fn get_entry( table: &ShardedHashTable, test: impl Fn(RowId) -> bool, ) -> Option { - let (shard, hash) = hash_code(table.shard_data(), row, n_keys); + let hash = shard_hash(table.shard_data(), row, n_keys); table - .get_shard(shard) - .find(hash, |ent| { - ent.hashcode == hash as HashCode && test(ent.row) + .get_shard(hash.shard) + .find(hash.compact.probe().raw(), |ent| { + ent.hash == hash.compact && test(ent.row) }) .map(|ent| ent.row) } @@ -1366,23 +1409,20 @@ fn get_entry_mut<'a>( table: &'a mut ShardedHashTable, test: impl Fn(RowId) -> bool, ) -> Option<&'a mut RowId> { - let (shard, hash) = hash_code(table.shard_data(), row, n_keys); - table.mut_shards()[shard.index()] - .find_mut(hash, |ent| { - ent.hashcode == hash as HashCode && test(ent.row) + let hash = shard_hash(table.shard_data(), row, n_keys); + table.mut_shards()[hash.shard.index()] + .find_mut(hash.compact.probe().raw(), |ent| { + ent.hash == hash.compact && test(ent.row) }) .map(|ent| &mut ent.row) } -fn hash_code(shard_data: ShardData, row: &[Value], n_keys: usize) -> (ShardId, u64) { +fn shard_hash(shard_data: ShardData, row: &[Value], n_keys: usize) -> ShardHash { let mut hasher = FxHasher::default(); for val in &row[0..n_keys] { hasher.write_usize(val.index()); } - let full_code = hasher.finish(); - // We keep this cast here to allow for experimenting with HashCode=u32. - #[allow(clippy::unnecessary_cast)] - (shard_data.shard_id(full_code), full_code as HashCode as u64) + ShardHash::from_full(shard_data, FullHash(hasher.finish())) } fn update_sort_offsets( @@ -1405,18 +1445,6 @@ fn update_sort_offsets( } } -fn mutation_partitioning(shard_data: ShardData, ordered: bool) -> HashPartitioning { - if ordered || shard_data.n_shards() == 1 { - HashPartitioning::grouped(0, 0, shard_data.n_shards()) - } else { - HashPartitioning::grouped( - MUTATION_CACHE_WINDOW_BITS, - MUTATION_PARTITION_BITS, - shard_data.n_shards(), - ) - } -} - fn drain_queue(queue: &SegQueue) -> Vec { let mut drained = Vec::new(); while let Some(item) = queue.pop() { @@ -1436,20 +1464,16 @@ unsafe fn mark_rows_stale(data: &RowBuffer, rows: &[RowId]) { /// A simple struct for packaging up pending mutations to a `SortedWritesTable`. struct PendingState { - pending_rows: DenseIdMap>>, - pending_removals: DenseIdMap>>, - pending_known_removals: DenseIdMap>>, + pending_rows: DenseIdMap>, + pending_removals: DenseIdMap>, + pending_known_removals: DenseIdMap>>, total_removals: AtomicUsize, total_rows: AtomicUsize, - insert_partitioning: HashPartitioning, - removal_partitioning: HashPartitioning, } impl PendingState { - fn new(shard_data: ShardData, ordered: bool) -> PendingState { + fn new(shard_data: ShardData) -> PendingState { let n_shards = shard_data.n_shards(); - let insert_partitioning = mutation_partitioning(shard_data, ordered); - let removal_partitioning = mutation_partitioning(shard_data, false); let mut pending_rows = DenseIdMap::with_capacity(n_shards); let mut pending_removals = DenseIdMap::with_capacity(n_shards); let mut pending_known_removals = DenseIdMap::with_capacity(n_shards); @@ -1465,8 +1489,6 @@ impl PendingState { pending_known_removals, total_removals: AtomicUsize::new(0), total_rows: AtomicUsize::new(0), - insert_partitioning, - removal_partitioning, } } fn clear(&self) { @@ -1537,8 +1559,6 @@ impl PendingState { pending_known_removals, total_removals: AtomicUsize::new(self.total_removals.load(Ordering::Acquire)), total_rows: AtomicUsize::new(self.total_rows.load(Ordering::Acquire)), - insert_partitioning: self.insert_partitioning, - removal_partitioning: self.removal_partitioning, } } } @@ -1642,31 +1662,21 @@ impl OrderingChecker for SortChecker { /// /// Multiple rows with the same key are merged before the surviving value is /// combined with the destination table's current value. Physical rows and -/// their unchanged full hashes stay aligned, including stale intermediate +/// their compact hash fingerprints stay aligned, including stale intermediate /// merge outputs. struct CoalescedInsertBatch { n_keys: usize, hash: Pooled>, rows: RowBuffer, - hashes: Vec, + hashes: Vec, n_stale: usize, scratch: Pooled>, } impl CoalescedInsertBatch { - fn rows_hashed(&self) -> impl Iterator { + fn rows_hashed(&self) -> impl Iterator { debug_assert_eq!(self.hashes.len(), self.rows.len()); - self.hashes - .iter() - .copied() - .zip(self.rows.iter()) - .map(|(hash, row)| { - // Keep this conversion at the boundary so HashCode can still - // be changed independently. - #[allow(clippy::unnecessary_cast)] - let hash = hash as u64; - (hash, row) - }) + self.hashes.iter().copied().zip(self.rows.iter()) } fn new(n_keys: usize, n_cols: usize, capacity: usize) -> Self { @@ -1679,7 +1689,7 @@ impl CoalescedInsertBatch { scratch: ps.get(), }); res.hash - .reserve(capacity, |entry| coalesced_hash(entry.hashcode())); + .reserve(capacity, |entry| coalesced_hash(entry.hash)); res.rows.reserve(capacity); res } @@ -1701,7 +1711,7 @@ impl CoalescedInsertBatch { fn insert_hashed( &mut self, - hc: u64, + hash: CompactHash, row: &[Value], mut merge_fn: impl FnMut(&[Value], &[Value], &mut Vec) -> bool, ) { @@ -1712,12 +1722,12 @@ impl CoalescedInsertBatch { debug_assert_eq!(self.hashes.len(), self.rows.len()); use hashbrown::hash_table::Entry; let entry = self.hash.entry( - coalesced_hash(hc), + coalesced_hash(hash), |entry| { - entry.hashcode() == hc + entry.hash == hash && self.rows.get_row(entry.row)[0..self.n_keys] == row[0..self.n_keys] }, - |entry| coalesced_hash(entry.hashcode()), + |entry| coalesced_hash(entry.hash), ); match entry { Entry::Occupied(mut occupied) => { @@ -1729,7 +1739,7 @@ impl CoalescedInsertBatch { "merge functions must preserve table keys" ); let new = self.rows.add_row(&self.scratch); - self.hashes.push(hc as HashCode); + self.hashes.push(hash); self.rows.set_stale(occupied.get().row); self.n_stale += 1; occupied.get_mut().row = new; @@ -1738,11 +1748,8 @@ impl CoalescedInsertBatch { } Entry::Vacant(vacant) => { let next = self.rows.add_row(row); - self.hashes.push(hc as HashCode); - vacant.insert(TableEntry { - hashcode: hc as HashCode, - row: next, - }); + self.hashes.push(hash); + vacant.insert(TableEntry { hash, row: next }); } } debug_assert_eq!(self.hashes.len(), self.rows.len()); @@ -1756,9 +1763,9 @@ impl CoalescedInsertBatch { } #[inline] -fn coalesced_hash(hash: u64) -> u64 { - // Cache partitioning fixes a run of middle bits in the full hash. Mix - // higher bits down before probing the temporary table so those fixed bits - // do not artificially restrict its usable buckets. +fn coalesced_hash(hash: CompactHash) -> u64 { + // Mix the compact probe before using it in the temporary coalescing table + // so its fixed bucket/tag fields do not restrict the usable buckets. + let hash = hash.probe().raw(); hash ^ hash.wrapping_shr(17) } diff --git a/core-relations/src/table/tests.rs b/core-relations/src/table/tests.rs index ae275f33b..a32b8948f 100644 --- a/core-relations/src/table/tests.rs +++ b/core-relations/src/table/tests.rs @@ -1,3 +1,5 @@ +use std::mem::{align_of, size_of}; + use crate::numeric_id::NumericId; use rand::{Rng, rng}; @@ -5,14 +7,185 @@ use crate::{ common::{HashMap, ShardId, Value}, offsets::{RowId, SubsetRef}, row_buffer::TaggedRowBuffer, - table::{SortedWritesTable, TableEntry, hash_code}, + table::{ + CompactHash, FullHash, HashedRowBuffer, KnownRemoval, ProbeHash, ShardHash, + SortedWritesTable, TableEntry, shard_hash, + }, table_shortcuts::{fill_table, v}, - table_spec::{ColumnId, Constraint, Offset, Table, WrappedTable}, + table_spec::{ColumnId, Constraint, MutationBuffer, Offset, Table, WrappedTable}, }; use egglog_concurrency::ThreadPool; use super::sharded_hash_table::ShardedHashTable; +#[test] +fn compact_hash_preserves_bucket_bits_and_exact_h2_tag() { + const LOW_MASK: u64 = (1 << 25) - 1; + const H2_SHIFT: u32 = if usize::BITS < u64::BITS { + usize::BITS - 7 + } else { + u64::BITS - 7 + }; + const H2_MASK: u64 = 0x7f << H2_SHIFT; + let cases = [ + 0, + 1, + LOW_MASK, + 1 << 25, + 1 << (H2_SHIFT - 1), + 1 << H2_SHIFT, + 0x1234_5678_9abc_def0, + u64::MAX, + ]; + + for raw in cases { + let compact = CompactHash::from_full(FullHash(raw)); + let probe = compact.probe().raw(); + + assert_eq!(compact.0 as u64 & LOW_MASK, raw & LOW_MASK); + assert_eq!((compact.0 >> 25) as u64, (raw >> H2_SHIFT) & 0x7f); + assert_eq!(probe & u32::MAX as u64, compact.0 as u64); + assert_eq!(probe & LOW_MASK, raw & LOW_MASK); + assert_eq!(probe & H2_MASK, raw & H2_MASK); + assert_eq!(probe & !(u32::MAX as u64 | H2_MASK), 0); + } +} + +#[test] +fn compact_hash_layouts_remain_small() { + assert_eq!(size_of::(), 4); + assert_eq!(size_of::(), 8); + assert_eq!(size_of::(), 8); + assert_eq!(size_of::(), 8); + assert_eq!(align_of::(), 4); + assert_eq!(size_of::(), 8); + assert_eq!(align_of::(), 4); +} + +#[test] +fn full_hash_routes_before_omitted_shard_bits_are_compacted() { + let shard_data = crate::common::ShardData::new(16); + let left = ShardHash::from_full(shard_data, FullHash(0)); + let right = ShardHash::from_full(shard_data, FullHash(1 << 53)); + + assert_eq!(left.compact, right.compact); + assert_ne!(left.shard, right.shard); +} + +#[test] +fn hashed_row_buffer_round_trips_max_compact_hash_in_trailing_lane() { + let row = [v(7), v(11)]; + let mut buffer = HashedRowBuffer::new(row.len()); + buffer.add_row(CompactHash(u32::MAX), &row); + buffer.add_row(CompactHash(17), &[v(13), v(19)]); + + assert_eq!(buffer.values[row.len()].rep(), u32::MAX); + let buffered = buffer.rows_hashed().collect::>(); + assert_eq!(buffered[0], (CompactHash(u32::MAX), row.as_slice())); + assert_eq!(buffered[1], (CompactHash(17), [v(13), v(19)].as_slice())); +} + +#[test] +fn same_shard_compact_collision_survives_table_lifecycle() { + let shard_data = crate::common::ShardData::new(16); + let left_key = [Value::new(70_759_468), Value::new(2_418_062_243)]; + let right_key = [Value::new(3_427_223_041), Value::new(1_840_803_043)]; + let left_hash = shard_hash(shard_data, &left_key, left_key.len()); + let right_hash = shard_hash(shard_data, &right_key, right_key.len()); + assert_eq!(left_hash, right_hash); + assert_eq!(left_hash.shard, ShardId::new(14)); + assert_eq!(left_hash.compact, CompactHash(0x578d_834a)); + + empty_execution_state!(e); + let mut table = SortedWritesTable::new( + 2, + 3, + None, + vec![], + Box::new(|_, old, new, out: &mut Vec| { + out.extend_from_slice(new); + old != new + }), + ); + table.hash = ShardedHashTable::with_shards(16); + table.pending_state = std::sync::Arc::new(super::PendingState::new(shard_data)); + + table + .new_buffer() + .stage_insert(&[left_key[0], left_key[1], v(10)]); + table + .new_buffer() + .stage_insert(&[right_key[0], right_key[1], v(20)]); + table.merge(&mut e); + assert_eq!(table.get_row(&left_key).unwrap().vals[2], v(10)); + assert_eq!(table.get_row(&right_key).unwrap().vals[2], v(20)); + + table + .new_buffer() + .stage_insert(&[left_key[0], left_key[1], v(30)]); + table.merge(&mut e); + assert_eq!(table.get_row(&left_key).unwrap().vals[2], v(30)); + assert_eq!(table.get_row(&right_key).unwrap().vals[2], v(20)); + + table.rehash(); + assert_eq!(table.get_row(&left_key).unwrap().vals[2], v(30)); + assert_eq!(table.get_row(&right_key).unwrap().vals[2], v(20)); + + table.new_buffer().stage_remove(&right_key); + table.merge(&mut e); + assert_eq!(table.get_row(&left_key).unwrap().vals[2], v(30)); + assert!(table.get_row(&right_key).is_none()); + + table + .new_buffer() + .stage_insert(&[right_key[0], right_key[1], v(40)]); + table.merge(&mut e); + assert_eq!(table.get_row(&left_key).unwrap().vals[2], v(30)); + assert_eq!(table.get_row(&right_key).unwrap().vals[2], v(40)); +} + +#[test] +fn same_shard_compact_collision_survives_parallel_insert_and_delete() { + ThreadPool::new(4).install(|| { + let shard_data = crate::common::ShardData::new(16); + let left_key = [Value::new(70_759_468), Value::new(2_418_062_243)]; + let right_key = [Value::new(3_427_223_041), Value::new(1_840_803_043)]; + assert_eq!( + shard_hash(shard_data, &left_key, left_key.len()), + shard_hash(shard_data, &right_key, right_key.len()) + ); + + empty_execution_state!(e); + let mut table = SortedWritesTable::new( + 2, + 3, + None, + vec![], + Box::new(|_, old, new, out: &mut Vec| { + out.extend_from_slice(new); + old != new + }), + ); + table.hash = ShardedHashTable::with_shards(16); + table.pending_state = std::sync::Arc::new(super::PendingState::new(shard_data)); + + table + .new_buffer() + .stage_insert(&[left_key[0], left_key[1], v(10)]); + table + .new_buffer() + .stage_insert(&[right_key[0], right_key[1], v(20)]); + assert!(table.parallel_insert(&e, ())); + assert_eq!(table.get_row(&left_key).unwrap().vals[2], v(10)); + assert_eq!(table.get_row(&right_key).unwrap().vals[2], v(20)); + + table.new_buffer().stage_remove(&right_key); + assert!(table.parallel_delete()); + assert_eq!(table.get_row(&left_key).unwrap().vals[2], v(10)); + assert!(table.get_row(&right_key).is_none()); + }); +} + fn dump_buf(buf: &TaggedRowBuffer) -> Vec<(RowId, Vec)> { let mut res = Vec::new(); buf.iter() @@ -182,7 +355,7 @@ fn known_row_removals_match_the_staged_row_id() { } #[test] -fn pending_partitioned_mutations_survive_table_clone() { +fn pending_cached_hash_mutations_survive_table_clone() { empty_execution_state!(e); let mut table = fill_table( vec![vec![v(1), v(10)], vec![v(2), v(20)]], @@ -190,8 +363,10 @@ fn pending_partitioned_mutations_survive_table_clone() { None, |_, new| Some(new.to_vec()), ); - let mut pending = table.new_buffer(); + let known_row = table.get_row(&[v(2)]).unwrap().id; + let mut pending = table.new_table_buffer(); pending.stage_remove(&[v(1)]); + pending.stage_remove_row(known_row, &[v(2)]); pending.stage_insert(&[v(3), v(30)]); drop(pending); @@ -201,7 +376,7 @@ fn pending_partitioned_mutations_survive_table_clone() { for table in [&table, &cloned] { assert!(table.get_row(&[v(1)]).is_none()); - assert_eq!(&*table.get_row(&[v(2)]).unwrap().vals, &[v(2), v(20)]); + assert!(table.get_row(&[v(2)]).is_none()); assert_eq!(&*table.get_row(&[v(3)]).unwrap().vals, &[v(3), v(30)]); } } @@ -402,8 +577,9 @@ fn parallel_unsorted_rehash_compacts_heavy_stale_rows_and_repoints_entries() { for entry in shard.iter() { let row = table.data.data.get_row(entry.row); assert_eq!(row[1], v(UPDATES)); - assert_eq!(entry.hashcode, hash_code(shard_data, row, 1).1); - assert_eq!(shard_data.shard_id(entry.hashcode).index(), shard_index); + let expected_hash = shard_hash(shard_data, row, 1); + assert_eq!(entry.hash, expected_hash.compact); + assert_eq!(expected_hash.shard.index(), shard_index); assert!(entry.row.index() < KEYS); assert!(!seen_rows[entry.row.index()]); @@ -476,7 +652,7 @@ fn shard_math() { let mut hist = HashMap::default(); (0..100_000) .map(|_| { - hash_code( + shard_hash( table.shard_data(), &[ Value::new(rng.random()), @@ -485,7 +661,7 @@ fn shard_math() { ], 2, ) - .0 + .shard }) .for_each(|id| *hist.entry(id).or_insert(0) += 1); assert!(hist.iter().all(|(_, count)| *count > 100), "{hist:?}");