1use std::{
8 any::Any,
9 cmp,
10 hash::Hasher,
11 mem,
12 sync::{
13 Arc, Weak,
14 atomic::{AtomicUsize, Ordering},
15 },
16};
17
18use crate::numeric_id::{DenseIdMap, NumericId};
19use crossbeam_queue::SegQueue;
20use hashbrown::HashTable;
21use rustc_hash::FxHasher;
22use sharded_hash_table::ShardedHashTable;
23
24use crate::{
25 Pooled, TableChange, TableId,
26 action::ExecutionState,
27 common::{HashMap, ShardData, ShardId, SubsetTracker, Value},
28 hash_index::{ColumnIndex, Index},
29 offsets::{OffsetRange, Offsets, RowId, SortedOffsetVector, Subset, SubsetRef},
30 parallel,
31 parallel_heuristics::parallelize_table_op,
32 pool::with_pool_set,
33 row_buffer::{ParallelRowBufWriter, RowBuffer},
34 table_spec::{
35 ColumnId, Constraint, Generation, MutationBuffer, Offset, Row, Table, TableSpec,
36 TableVersion,
37 },
38};
39
40mod rebuild;
41mod sharded_hash_table;
42#[cfg(test)]
43mod tests;
44
45type HashCode = u64;
51
52#[derive(Clone, Debug)]
54pub(crate) struct TableEntry {
55 hashcode: HashCode,
56 row: RowId,
57}
58
59impl TableEntry {
60 fn hashcode(&self) -> u64 {
61 #[allow(clippy::unnecessary_cast)]
63 {
64 self.hashcode as u64
65 }
66 }
67}
68
69#[derive(Clone)]
74struct Rows {
75 data: RowBuffer,
76 scratch: RowBuffer,
77 stale_rows: usize,
78}
79
80impl Rows {
81 fn new(data: RowBuffer) -> Rows {
82 let arity = data.arity();
83 Rows {
84 data,
85 scratch: RowBuffer::new(arity),
86 stale_rows: 0,
87 }
88 }
89 fn clear(&mut self) {
90 self.data.clear();
91 self.stale_rows = 0;
92 }
93 fn next_row(&self) -> RowId {
94 RowId::from_usize(self.data.len())
95 }
96 fn set_stale(&mut self, row: RowId) {
97 if !self.data.set_stale(row) {
98 self.stale_rows += 1;
99 }
100 }
101
102 fn get_row(&self, row: RowId) -> Option<&[Value]> {
103 let row = self.data.get_row(row);
104 if row[0].is_stale() { None } else { Some(row) }
105 }
106
107 unsafe fn get_row_unchecked(&self, row: RowId) -> Option<&[Value]> {
109 let row = unsafe { self.data.get_row_unchecked(row) };
110 if row[0].is_stale() { None } else { Some(row) }
111 }
112
113 fn add_row(&mut self, row: &[Value]) -> RowId {
114 if row[0].is_stale() {
115 self.stale_rows += 1;
116 }
117 self.data.add_row(row)
118 }
119
120 fn remove_stale(&mut self, remap: impl FnMut(&[Value], RowId, RowId)) {
121 self.data.remove_stale(remap);
122 self.stale_rows = 0;
123 }
124}
125
126pub type MergeFn =
134 dyn Fn(&mut ExecutionState, &[Value], &[Value], &mut Vec<Value>) -> bool + Send + Sync;
135
136pub struct SortedWritesTable {
137 generation: Generation,
138 data: Rows,
139 hash: ShardedHashTable<TableEntry>,
140
141 n_keys: usize,
142 n_columns: usize,
143 sort_by: Option<ColumnId>,
144 offsets: Vec<(Value, RowId)>,
145
146 pending_state: Arc<PendingState>,
147 merge: Arc<MergeFn>,
148 to_rebuild: Vec<ColumnId>,
149 rebuild_index: Index<ColumnIndex>,
150 subset_tracker: SubsetTracker,
152}
153
154impl Clone for SortedWritesTable {
155 fn clone(&self) -> SortedWritesTable {
156 SortedWritesTable {
157 generation: self.generation,
158 data: self.data.clone(),
159 hash: self.hash.clone(),
160 n_keys: self.n_keys,
161 n_columns: self.n_columns,
162 sort_by: self.sort_by,
163 offsets: self.offsets.clone(),
164 pending_state: Arc::new(self.pending_state.deep_copy()),
165 merge: self.merge.clone(),
166 to_rebuild: self.to_rebuild.clone(),
167 rebuild_index: Index::new(self.to_rebuild.clone(), ColumnIndex::new()),
168 subset_tracker: Default::default(),
169 }
170 }
171}
172
173#[derive(Clone)]
179enum ArbitraryRowBuffer {
180 NonEmpty(RowBuffer),
181 Empty { rows: usize },
182}
183
184impl ArbitraryRowBuffer {
185 fn new(arity: usize) -> ArbitraryRowBuffer {
186 if arity == 0 {
187 ArbitraryRowBuffer::Empty { rows: 0 }
188 } else {
189 ArbitraryRowBuffer::NonEmpty(RowBuffer::new(arity))
190 }
191 }
192
193 fn add_row(&mut self, row: &[Value]) {
194 match self {
195 ArbitraryRowBuffer::NonEmpty(buf) => {
196 buf.add_row(row);
197 }
198 ArbitraryRowBuffer::Empty { rows } => {
199 *rows += 1;
200 }
201 }
202 }
203
204 fn len(&self) -> usize {
205 match self {
206 ArbitraryRowBuffer::NonEmpty(buf) => buf.len(),
207 ArbitraryRowBuffer::Empty { rows } => *rows,
208 }
209 }
210
211 fn for_each(&self, mut f: impl FnMut(&[Value])) {
212 match self {
213 ArbitraryRowBuffer::NonEmpty(buf) => {
214 for row in buf.iter() {
215 f(row);
216 }
217 }
218 ArbitraryRowBuffer::Empty { rows } => {
219 for _ in 0..*rows {
220 f(&[]);
221 }
222 }
223 }
224 }
225}
226
227struct Buffer {
228 pending_rows: DenseIdMap<ShardId, RowBuffer>,
229 pending_removals: DenseIdMap<ShardId, ArbitraryRowBuffer>,
230 state: Weak<PendingState>,
231 n_cols: u32,
232 n_keys: u32,
233 shard_data: ShardData,
234}
235
236impl MutationBuffer for Buffer {
237 fn stage_insert(&mut self, row: &[Value]) {
238 let (shard, _) = hash_code(self.shard_data, row, self.n_keys as _);
239 self.pending_rows
240 .get_or_insert(shard, || RowBuffer::new(self.n_cols as _))
241 .add_row(row);
242 }
243 fn stage_remove(&mut self, key: &[Value]) {
244 let (shard, _) = hash_code(self.shard_data, key, self.n_keys as _);
245 self.pending_removals
246 .get_or_insert(shard, || ArbitraryRowBuffer::new(self.n_keys as _))
247 .add_row(key);
248 }
249 fn fresh_handle(&self) -> Box<dyn MutationBuffer> {
250 Box::new(Buffer {
251 pending_rows: Default::default(),
252 pending_removals: Default::default(),
253 state: self.state.clone(),
254 n_cols: self.n_cols,
255 n_keys: self.n_keys,
256 shard_data: self.shard_data,
257 })
258 }
259}
260
261impl Drop for Buffer {
262 fn drop(&mut self) {
263 if let Some(state) = self.state.upgrade() {
264 let mut rows = 0;
265 for shard_id in 0..self.pending_rows.n_ids() {
266 let shard = ShardId::from_usize(shard_id);
267 let Some(buf) = self.pending_rows.take(shard) else {
268 continue;
269 };
270 rows += buf.len();
271 state.pending_rows[shard].push(buf);
272 }
273 state.total_rows.fetch_add(rows, Ordering::Relaxed);
274
275 let mut rows = 0;
276 for shard_id in 0..self.pending_removals.n_ids() {
277 let shard = ShardId::from_usize(shard_id);
278 let Some(buf) = self.pending_removals.take(shard) else {
279 continue;
280 };
281 rows += buf.len();
282 state.pending_removals[shard].push(buf);
283 }
284 state.total_removals.fetch_add(rows, Ordering::Relaxed);
285 }
286 }
287}
288
289impl Table for SortedWritesTable {
290 fn dyn_clone(&self) -> Box<dyn Table> {
291 Box::new(self.clone())
292 }
293 fn as_any(&self) -> &dyn Any {
294 self
295 }
296 fn clear(&mut self) {
297 self.pending_state.clear();
298 if self.data.data.len() == 0 {
299 return;
300 }
301 self.offsets.clear();
302 self.data.clear();
303 self.hash.clear();
304 self.generation = Generation::from_usize(self.version().major.index() + 1);
305 }
306
307 fn spec(&self) -> TableSpec {
308 TableSpec {
309 n_keys: self.n_keys,
310 n_vals: self.n_columns - self.n_keys,
311 uncacheable_columns: Default::default(),
312 allows_delete: true,
313 }
314 }
315
316 fn apply_rebuild(
317 &mut self,
318 table_id: TableId,
319 table: &crate::WrappedTable,
320 next_ts: Value,
321 exec_state: &mut ExecutionState,
322 ) -> bool {
323 self.do_rebuild(table_id, table, next_ts, exec_state)
324 }
325
326 fn refresh_rows_for_values(&mut self, dirty_ids: &[Value], next_ts: Value) -> bool {
327 SortedWritesTable::refresh_rows_for_values(self, dirty_ids, next_ts)
328 }
329
330 fn version(&self) -> TableVersion {
331 TableVersion {
332 major: self.generation,
333 minor: Offset::from_usize(self.data.next_row().index()),
334 }
335 }
336
337 fn updates_since(&self, offset: Offset) -> Subset {
338 Subset::Dense(OffsetRange::new(
339 RowId::from_usize(offset.index()),
340 self.data.next_row(),
341 ))
342 }
343
344 fn all(&self) -> Subset {
345 Subset::Dense(OffsetRange::new(RowId::new(0), self.data.next_row()))
346 }
347
348 fn has_stale_rows(&self) -> bool {
349 self.data.stale_rows > 0
350 }
351
352 fn len(&self) -> usize {
353 self.data.data.len() - self.data.stale_rows
354 }
355
356 fn scan_generic(&self, subset: SubsetRef, mut f: impl FnMut(RowId, &[Value]))
357 where
358 Self: Sized,
359 {
360 let Some((_low, hi)) = subset.bounds() else {
361 return;
363 };
364 assert!(
365 hi.index() <= self.data.data.len(),
366 "{} vs. {}",
367 hi.index(),
368 self.data.data.len()
369 );
370 if self.data.stale_rows == 0 {
371 subset.offsets(|row| unsafe { f(row, self.data.data.get_row_unchecked(row)) })
377 } else {
378 subset.offsets(|row| unsafe {
380 if let Some(vals) = self.data.get_row_unchecked(row) {
381 f(row, vals)
382 }
383 })
384 }
385 }
386
387 fn scan_generic_bounded(
391 &self,
392 subset: SubsetRef,
393 start: Offset,
394 n: usize,
395 cs: &[Constraint],
396 mut f: impl FnMut(RowId, &[Value]),
397 ) -> Option<Offset>
398 where
399 Self: Sized,
400 {
401 let Some((_low, hi)) = subset.bounds() else {
402 return None;
404 };
405 assert!(
406 hi.index() <= self.data.data.len(),
407 "{} vs. {}",
408 hi.index(),
409 self.data.data.len()
410 );
411 if cs.is_empty() {
412 if self.data.stale_rows == 0 {
413 subset
416 .iter_bounded(start.index(), start.index() + n, |row| {
417 let entry = unsafe { self.data.data.get_row_unchecked(row) };
418 f(row, entry);
419 })
420 .map(Offset::from_usize)
421 } else {
422 subset
423 .iter_bounded(start.index(), start.index() + n, |row| {
424 let Some(entry) = (unsafe { self.data.get_row_unchecked(row) }) else {
426 return;
427 };
428 f(row, entry);
429 })
430 .map(Offset::from_usize)
431 }
432 } else {
433 subset
434 .iter_bounded(start.index(), start.index() + n, |row| {
435 let Some(entry) = (unsafe { self.get_if_unchecked(cs, row) }) else {
437 return;
438 };
439 f(row, entry);
440 })
441 .map(Offset::from_usize)
442 }
443 }
444
445 fn fast_subset(&self, constraint: &Constraint) -> Option<Subset> {
446 let sort_by = self.sort_by?;
447 match constraint {
448 Constraint::Eq { .. } => None,
449 Constraint::EqConst { col, val } => {
450 if col == &sort_by {
451 match self.binary_search_sort_val(*val) {
452 Ok((found, bound)) => Some(Subset::Dense(OffsetRange::new(found, bound))),
453 Err(_) => Some(Subset::empty()),
454 }
455 } else {
456 None
457 }
458 }
459 Constraint::LtConst { col, val } => {
460 if col == &sort_by {
461 match self.binary_search_sort_val(*val) {
462 Ok((found, _)) => {
463 Some(Subset::Dense(OffsetRange::new(RowId::new(0), found)))
464 }
465 Err(next) => Some(Subset::Dense(OffsetRange::new(RowId::new(0), next))),
466 }
467 } else {
468 None
469 }
470 }
471 Constraint::GtConst { col, val } => {
472 if col == &sort_by {
473 match self.binary_search_sort_val(*val) {
474 Ok((_, bound)) => {
475 Some(Subset::Dense(OffsetRange::new(bound, self.data.next_row())))
476 }
477 Err(next) => {
478 Some(Subset::Dense(OffsetRange::new(next, self.data.next_row())))
479 }
480 }
481 } else {
482 None
483 }
484 }
485 Constraint::LeConst { col, val } => {
486 if col == &sort_by {
487 match self.binary_search_sort_val(*val) {
488 Ok((_, bound)) => {
489 Some(Subset::Dense(OffsetRange::new(RowId::new(0), bound)))
490 }
491 Err(next) => Some(Subset::Dense(OffsetRange::new(RowId::new(0), next))),
492 }
493 } else {
494 None
495 }
496 }
497 Constraint::GeConst { col, val } => {
498 if col == &sort_by {
499 match self.binary_search_sort_val(*val) {
500 Ok((found, _)) => {
501 Some(Subset::Dense(OffsetRange::new(found, self.data.next_row())))
502 }
503 Err(next) => {
504 Some(Subset::Dense(OffsetRange::new(next, self.data.next_row())))
505 }
506 }
507 } else {
508 None
509 }
510 }
511 }
512 }
513
514 fn refine_one(&self, mut subset: Subset, c: &Constraint) -> Subset {
515 subset.retain(|row| self.eval(std::slice::from_ref(c), row));
518 subset
519 }
520
521 fn refine_ref(&self, subset: SubsetRef, cs: &[Constraint], _check_live: bool) -> Subset {
522 let n = subset.size();
526 let mut vec: Pooled<SortedOffsetVector> = with_pool_set(|ps| ps.get());
527 subset.offsets(|row| {
528 if self.eval(cs, row) {
529 unsafe { vec.push_unchecked(row) }
531 }
532 });
533 if let SubsetRef::Dense(range) = subset
534 && vec.slice().inner().len() == n
535 {
536 return Subset::Dense(range);
537 }
538 Subset::Sparse(vec)
539 }
540
541 fn new_buffer(&self) -> Box<dyn MutationBuffer> {
542 let n_shards = self.hash.shard_data().n_shards();
543 Box::new(Buffer {
544 pending_rows: DenseIdMap::with_capacity(n_shards),
545 pending_removals: DenseIdMap::with_capacity(n_shards),
546 state: Arc::downgrade(&self.pending_state),
547 n_keys: u32::try_from(self.n_keys).expect("n_keys should fit in u32"),
548 n_cols: u32::try_from(self.n_columns).expect("n_columns should fit in u32"),
549 shard_data: self.hash.shard_data(),
550 })
551 }
552
553 fn merge(&mut self, exec_state: &mut ExecutionState) -> TableChange {
554 let removed = self.do_delete();
555 let added = self.do_insert(exec_state);
556 self.maybe_rehash();
557 TableChange { removed, added }
558 }
559
560 fn get_row(&self, key: &[Value]) -> Option<Row> {
561 let id = get_entry(key, self.n_keys, &self.hash, |row| {
562 &self.data.get_row(row).unwrap()[0..self.n_keys] == key
563 })?;
564 let mut vals = with_pool_set(|ps| ps.get::<Vec<Value>>());
565 vals.extend_from_slice(self.data.get_row(id).unwrap());
566 Some(Row { id, vals })
567 }
568
569 fn get_row_column(&self, key: &[Value], col: ColumnId) -> Option<Value> {
570 let id = get_entry(key, self.n_keys, &self.hash, |row| {
571 &self.data.get_row(row).unwrap()[0..self.n_keys] == key
572 })?;
573 Some(self.data.get_row(id).unwrap()[col.index()])
574 }
575}
576
577impl SortedWritesTable {
578 pub fn new(
590 n_keys: usize,
591 n_columns: usize,
592 sort_by: Option<ColumnId>,
593 to_rebuild: Vec<ColumnId>,
594 merge_fn: Box<MergeFn>,
595 ) -> Self {
596 let hash = ShardedHashTable::<TableEntry>::default();
597 let shard_data = hash.shard_data();
598 let rebuild_index = Index::new(to_rebuild.clone(), ColumnIndex::new());
599 SortedWritesTable {
600 generation: Generation::new(0),
601 data: Rows::new(RowBuffer::new(n_columns)),
602 hash,
603 n_keys,
604 n_columns,
605 sort_by,
606 offsets: Default::default(),
607 pending_state: Arc::new(PendingState::new(shard_data)),
608 merge: merge_fn.into(),
609 to_rebuild,
610 rebuild_index,
611 subset_tracker: Default::default(),
612 }
613 }
614
615 fn parallel_delete(&mut self) -> bool {
617 let shard_data = self.hash.shard_data();
618 let pending_removals = &self.pending_state.pending_removals;
619 let data = &self.data.data;
620 let n_keys = self.n_keys;
621 let stale_delta: usize = parallel::map_mut(self.hash.mut_shards(), |shard_id, shard| {
622 let shard_id = ShardId::from_usize(shard_id);
623 if pending_removals[shard_id].is_empty() {
624 return 0;
625 }
626 let queue = &pending_removals[shard_id];
627 let mut marked_stale = 0;
628 while let Some(buf) = queue.pop() {
629 buf.for_each(|to_remove| {
630 let (actual_shard, hc) = hash_code(shard_data, to_remove, n_keys);
631 assert_eq!(actual_shard, shard_id);
632 if let Ok(entry) = shard.find_entry(hc, |entry| {
633 entry.hashcode == (hc as _)
634 && &data.get_row(entry.row)[0..n_keys] == to_remove
635 }) {
636 let (ent, _) = entry.remove();
637 marked_stale += unsafe { !data.set_stale_shared(ent.row) } as usize;
645 }
646 });
647 }
648 marked_stale
649 })
650 .into_iter()
651 .sum();
652 self.data.stale_rows += stale_delta;
654 stale_delta > 0
655 }
656 fn serial_delete(&mut self) -> bool {
657 let shard_data = self.hash.shard_data();
658 let mut changed = false;
659 self.hash
660 .mut_shards()
661 .iter_mut()
662 .enumerate()
663 .for_each(|(shard_id, shard)| {
664 let shard_id = ShardId::from_usize(shard_id);
665 let queue = &self.pending_state.pending_removals[shard_id];
666 while let Some(buf) = queue.pop() {
667 buf.for_each(|to_remove| {
668 let (actual_shard, hc) = hash_code(shard_data, to_remove, self.n_keys);
669 assert_eq!(actual_shard, shard_id);
670 if let Ok(entry) = shard.find_entry(hc, |entry| {
671 entry.hashcode == (hc as _)
672 && &self.data.get_row(entry.row).unwrap()[0..self.n_keys]
673 == to_remove
674 }) {
675 let (ent, _) = entry.remove();
676 self.data.set_stale(ent.row);
677 changed = true;
678 }
679 })
680 }
681 });
682 changed
683 }
684
685 fn do_delete(&mut self) -> bool {
686 let total = self.pending_state.total_removals.swap(0, Ordering::Relaxed);
687
688 if parallelize_table_op(total) {
689 self.parallel_delete()
690 } else {
691 self.serial_delete()
692 }
693 }
694
695 fn do_insert(&mut self, exec_state: &mut ExecutionState) -> bool {
696 let total = self.pending_state.total_rows.swap(0, Ordering::Relaxed);
697 self.data.data.reserve(total);
698 if parallelize_table_op(total) {
699 if let Some(col) = self.sort_by {
700 self.parallel_insert(
701 exec_state,
702 SortChecker {
703 col,
704 current: None,
705 baseline: self.offsets.last().map(|(v, _)| *v),
706 },
707 )
708 } else {
709 self.parallel_insert(exec_state, ())
710 }
711 } else {
712 self.serial_insert(exec_state)
713 }
714 }
715
716 fn serial_insert(&mut self, exec_state: &mut ExecutionState) -> bool {
717 let mut changed = false;
718 let n_keys = self.n_keys;
719 let mut scratch = with_pool_set(|ps| ps.get::<Vec<Value>>());
720 for (_outer_shard, queue) in self.pending_state.pending_rows.iter() {
721 if let Some(sort_by) = self.sort_by {
722 while let Some(buf) = queue.pop() {
723 for query in buf.non_stale() {
724 let key = &query[0..n_keys];
725 let entry = get_entry_mut(query, n_keys, &mut self.hash, |row| {
726 let Some(row) = self.data.get_row(row) else {
727 return false;
728 };
729 &row[0..n_keys] == key
730 });
731
732 if let Some(row) = entry {
733 let cur = self
737 .data
738 .get_row(*row)
739 .expect("table should not point to stale entry");
740 if (self.merge)(exec_state, cur, query, &mut scratch) {
741 let sort_val = query[sort_by.index()];
742 let new = self.data.add_row(&scratch);
743 if let Some(largest) = self.offsets.last().map(|(v, _)| *v) {
744 assert!(
745 sort_val >= largest,
746 "inserting row that violates sort order ({sort_val:?} vs. {largest:?})"
747 );
748 if sort_val > largest {
749 self.offsets.push((sort_val, new));
750 }
751 } else {
752 self.offsets.push((sort_val, new));
753 }
754 self.data.set_stale(*row);
755 *row = new;
756 changed = true;
757 }
758 scratch.clear();
759 } else {
760 let sort_val = query[sort_by.index()];
761 let new = self.data.add_row(query);
763 if let Some(largest) = self.offsets.last().map(|(v, _)| *v) {
764 assert!(
765 sort_val >= largest,
766 "inserting row that violates sort order {sort_val:?} vs. {largest:?}"
767 );
768 if sort_val > largest {
769 self.offsets.push((sort_val, new));
770 }
771 } else {
772 self.offsets.push((sort_val, new));
773 }
774 let (shard, hc) = hash_code(self.hash.shard_data(), query, self.n_keys);
775 debug_assert_eq!(shard, _outer_shard);
776 self.hash.mut_shards()[shard.index()].insert_unique(
777 hc as _,
778 TableEntry {
779 hashcode: hc as _,
780 row: new,
781 },
782 TableEntry::hashcode,
783 );
784 changed = true;
785 }
786 }
787 }
788 } else {
789 while let Some(buf) = queue.pop() {
791 for query in buf.non_stale() {
792 let key = &query[0..n_keys];
793 let entry = get_entry_mut(query, n_keys, &mut self.hash, |row| {
794 let Some(row) = self.data.get_row(row) else {
795 return false;
796 };
797 &row[0..n_keys] == key
798 });
799
800 if let Some(row) = entry {
801 let cur = self
802 .data
803 .get_row(*row)
804 .expect("table should not point to stale entry");
805 if (self.merge)(exec_state, cur, query, &mut scratch) {
806 let new = self.data.add_row(&scratch);
807 self.data.set_stale(*row);
808 *row = new;
809 changed = true;
810 }
811 scratch.clear();
812 } else {
813 let new = self.data.add_row(query);
815 let (shard, hc) = hash_code(self.hash.shard_data(), query, self.n_keys);
816 debug_assert_eq!(shard, _outer_shard);
817 self.hash.mut_shards()[shard.index()].insert_unique(
818 hc as _,
819 TableEntry {
820 hashcode: hc as _,
821 row: new,
822 },
823 TableEntry::hashcode,
824 );
825 changed = true;
826 }
827 }
828 }
829 };
830 }
831 changed
832 }
833
834 fn parallel_insert<C: OrderingChecker>(
835 &mut self,
836 exec_state: &ExecutionState,
837 checker: C,
838 ) -> bool {
839 const BATCH_SIZE: usize = 1 << 18;
840 let shard_data = self.hash.shard_data();
845 let n_keys = self.n_keys;
846 let n_cols = self.n_columns;
847 let next_offset = RowId::from_usize(self.data.data.len());
848 let row_writer = self.data.data.parallel_writer();
849 let pending_rows = &self.pending_state.pending_rows;
850 let merge = self.merge.clone();
851 let pending_adds = parallel::map_mut(self.hash.mut_shards(), |shard_id, shard| {
852 let shard_id = ShardId::from_usize(shard_id);
853 let mut checker = checker.clone();
854 let mut exec_state = exec_state.clone();
855 let mut scratch = with_pool_set(|ps| ps.get::<Vec<Value>>());
856 let queue = &pending_rows[shard_id];
857 let mut marked_stale = 0usize;
858 let mut staged = StagedOutputs::new(n_keys, n_cols, BATCH_SIZE);
859 let mut changed = false;
860 macro_rules! flush_staged_outputs {
863 () => {{
864 let (start_row, stale) = staged.write_output(&row_writer);
867 marked_stale += stale;
868 let mut cur_row = start_row;
877 let read_handle = row_writer.read_handle();
878 for row in staged.rows() {
879 if row.first().map(Value::is_stale).unwrap_or(false) {
880 cur_row = cur_row.inc();
881 continue;
882 }
883 use hashbrown::hash_table::Entry;
884 checker.check_local(row);
885 changed = true;
886 let key = &row[0..n_keys];
887 let (_actual_shard, hc) = hash_code(shard_data, row, n_keys);
888 #[cfg(any(debug_assertions, test))]
889 {
890 unsafe {
891 assert_eq!(read_handle.get_row_unchecked(cur_row), row);
894 }
895 }
896 debug_assert_eq!(_actual_shard, shard_id);
897 match shard.entry(
898 hc,
899 |ent| unsafe {
901 ent.hashcode == hc as HashCode
902 && &read_handle.get_row_unchecked(ent.row)[0..n_keys] == key
903 },
904 TableEntry::hashcode,
905 ) {
906 Entry::Occupied(mut occ) => {
907 let cur = unsafe { read_handle.get_row_unchecked(occ.get().row) };
910
911 if (merge)(&mut exec_state, cur, row, &mut scratch) {
917 unsafe {
918 let _was_stale = read_handle.set_stale_shared(occ.get().row);
919 debug_assert!(!_was_stale);
920 }
921 occ.get_mut().row = cur_row;
922 changed = true;
923 } else {
924 unsafe {
926 let _was_stale = read_handle.set_stale_shared(cur_row);
927 debug_assert!(!_was_stale);
928 }
929 }
930 marked_stale += 1;
931 scratch.clear();
932 }
933 Entry::Vacant(v) => {
934 changed = true;
935 v.insert(TableEntry {
936 hashcode: hc as HashCode,
937 row: cur_row,
938 });
939 }
940 }
941
942 cur_row = cur_row.inc();
943 }
944 staged.clear();
945 }};
946 }
947 while let Some(buf) = queue.pop() {
952 for row in buf.non_stale() {
956 staged.insert(row, |cur, new, out| (merge)(&mut exec_state, cur, new, out));
957 if staged.len() >= BATCH_SIZE {
958 flush_staged_outputs!();
959 }
960 }
961 }
962 flush_staged_outputs!();
963 (checker, marked_stale, changed)
964 });
965 self.data.data = row_writer.finish();
966 let checker = C::check_global(pending_adds.iter().map(|(checker, _, _)| checker));
971 checker.update_offsets(next_offset, &mut self.offsets);
972
973 self.data.stale_rows += pending_adds
975 .iter()
976 .map(|(_, stale, _)| *stale)
977 .sum::<usize>();
978
979 pending_adds.iter().any(|(_, _, changed)| *changed)
981 }
982
983 fn binary_search_sort_val(&self, val: Value) -> Result<(RowId, RowId), RowId> {
984 debug_assert!(
985 self.offsets.windows(2).all(|x| x[0].1 < x[1].1),
986 "{:?}",
987 self.offsets
988 );
989
990 debug_assert!(
991 self.offsets.windows(2).all(|x| x[0].0 < x[1].0),
992 "{:?}",
993 self.offsets
994 );
995 match self.offsets.binary_search_by_key(&val, |(v, _)| *v) {
996 Ok(got) => Ok((
997 self.offsets[got].1,
998 self.offsets
999 .get(got + 1)
1000 .map(|(_, r)| *r)
1001 .unwrap_or(self.data.next_row()),
1002 )),
1003 Err(next) => Err(self
1004 .offsets
1005 .get(next)
1006 .map(|(_, id)| *id)
1007 .unwrap_or(self.data.next_row())),
1008 }
1009 }
1010 fn eval(&self, cs: &[Constraint], row: RowId) -> bool {
1011 self.get_if(cs, row).is_some()
1012 }
1013
1014 fn eval_constraints(cs: &[Constraint], row: &[Value]) -> bool {
1015 cs.iter().all(|constraint| match constraint {
1016 Constraint::Eq { l_col, r_col } => row[l_col.index()] == row[r_col.index()],
1017 Constraint::EqConst { col, val } => row[col.index()] == *val,
1018 Constraint::LtConst { col, val } => row[col.index()] < *val,
1019 Constraint::GtConst { col, val } => row[col.index()] > *val,
1020 Constraint::LeConst { col, val } => row[col.index()] <= *val,
1021 Constraint::GeConst { col, val } => row[col.index()] >= *val,
1022 })
1023 }
1024
1025 unsafe fn get_if_unchecked(&self, cs: &[Constraint], row: RowId) -> Option<&[Value]> {
1026 let row = unsafe { self.data.get_row_unchecked(row) }?;
1027 if Self::eval_constraints(cs, row) {
1028 Some(row)
1029 } else {
1030 None
1031 }
1032 }
1033
1034 fn get_if(&self, cs: &[Constraint], row: RowId) -> Option<&[Value]> {
1035 let row = self.data.get_row(row)?;
1036 if Self::eval_constraints(cs, row) {
1037 Some(row)
1038 } else {
1039 None
1040 }
1041 }
1042
1043 fn maybe_rehash(&mut self) {
1044 if self.data.stale_rows <= cmp::max(16, self.data.data.len() / 2) {
1045 return;
1046 }
1047
1048 if parallelize_table_op(self.data.data.len()) {
1049 self.parallel_rehash();
1050 } else {
1051 self.rehash();
1052 }
1053 }
1054 fn parallel_rehash(&mut self) {
1055 let Some(sort_by) = self.sort_by else {
1059 self.rehash();
1066 return;
1067 };
1068 self.generation = self.generation.inc();
1069 assert!(!self.offsets.is_empty());
1070 struct TimestampStats {
1071 value: Value,
1072 count: usize,
1073 histogram: Pooled<DenseIdMap<ShardId, usize>>,
1074 }
1075 impl Default for TimestampStats {
1076 fn default() -> TimestampStats {
1077 TimestampStats {
1078 value: Value::stale(),
1079 count: 0,
1080 histogram: with_pool_set(|ps| ps.get()),
1081 }
1082 }
1083 }
1084 let mut results = Vec::<TimestampStats>::with_capacity(self.offsets.len());
1085 let offset_windows = self
1086 .offsets
1087 .windows(2)
1088 .map(|xs| {
1089 let [(start_val, start_row), (_, end_row)] = xs else {
1090 unreachable!()
1091 };
1092 (*start_val, *start_row, *end_row)
1093 })
1094 .collect::<Vec<_>>();
1095 macro_rules! compute_hist {
1097 ($start_val: expr, $start_row: expr, $end_row: expr) => {{
1098 let mut histogram: Pooled<DenseIdMap<ShardId, usize>> =
1099 with_pool_set(|ps| ps.get());
1100 let mut cur_row = $start_row;
1101 let mut count = 0;
1102 while cur_row < $end_row {
1103 if let Some(row) = self.data.get_row(cur_row) {
1104 count += 1;
1105 let (shard, _) = hash_code(self.hash.shard_data(), row, self.n_keys);
1106 *histogram.get_or_default(shard) += 1;
1107 }
1108 cur_row = cur_row.inc();
1109 }
1110 TimestampStats {
1111 value: $start_val,
1112 count,
1113 histogram,
1114 }
1115 }};
1116 }
1117 results.extend(parallel::map(
1118 &offset_windows,
1119 |_, (start_val, start_row, end_row)| compute_hist!(*start_val, *start_row, *end_row),
1120 ));
1121 let (start_val, start_row) = self.offsets.last().unwrap();
1123 let end_row = self.data.next_row();
1124 let last = compute_hist!(*start_val, *start_row, end_row);
1125 results.push(last);
1126 let mut prev_count = 0;
1131 self.offsets.clear();
1132 for stats in results.iter_mut() {
1133 if stats.count == 0 {
1134 continue;
1135 }
1136 self.offsets
1137 .push((stats.value, RowId::from_usize(prev_count)));
1138 let mut inner = prev_count;
1139 for (_, count) in stats.histogram.iter_mut() {
1140 let tmp = *count;
1143 *count = inner;
1144 inner += tmp;
1145 }
1146 prev_count += stats.count;
1147 debug_assert_eq!(inner, prev_count)
1148 }
1149
1150 self.data.scratch.clear();
1163 self.data.scratch.reserve(prev_count);
1164 let scratch_ptr = self.data.scratch.raw_rows() as usize;
1165 let data = &self.data.data;
1166 let n_columns = self.n_columns;
1167 parallel::for_each_mut(self.hash.mut_shards(), |shard_id, shard| {
1168 let shard_id = ShardId::from_usize(shard_id);
1169 let scratch_ptr = scratch_ptr as *const Value;
1170 let mut progress = HashMap::<Value , RowId >::default();
1171 progress.reserve(results.len());
1172 for stats in &results {
1173 let Some(start) = stats.histogram.get(shard_id) else {
1174 continue;
1175 };
1176 progress.insert(stats.value, RowId::from_usize(*start));
1177 }
1178 for TableEntry { row: row_id, .. } in shard.iter_mut() {
1179 let row = data.get_row(*row_id);
1180 debug_assert!(!row[0].is_stale(), "shard should not map to a stale value");
1181 let val = row[sort_by.index()];
1182 let next = progress[&val];
1183 unsafe {
1185 std::ptr::copy_nonoverlapping(
1186 row.as_ptr(),
1187 scratch_ptr.add(next.index() * n_columns) as *mut Value,
1188 n_columns,
1189 )
1190 }
1191 *row_id = next;
1192 progress.insert(val, next.inc());
1193 }
1194 });
1195 unsafe { self.data.scratch.set_len(prev_count) };
1197 mem::swap(&mut self.data.data, &mut self.data.scratch);
1198 self.data.stale_rows = 0;
1199 }
1200 fn rehash_impl(
1201 sort_by: Option<ColumnId>,
1202 n_keys: usize,
1203 rows: &mut Rows,
1204 offsets: &mut Vec<(Value, RowId)>,
1205 hash: &mut ShardedHashTable<TableEntry>,
1206 ) {
1207 if let Some(sort_by) = sort_by {
1208 offsets.clear();
1209 rows.remove_stale(|row, old, new| {
1210 let stale_entry = get_entry_mut(row, n_keys, hash, |x| x == old)
1211 .expect("non-stale entry not mapped in hash");
1212 *stale_entry = new;
1213 let sort_col = row[sort_by.index()];
1214 if let Some((max, _)) = offsets.last() {
1215 if sort_col > *max {
1216 offsets.push((sort_col, new));
1217 }
1218 } else {
1219 offsets.push((sort_col, new));
1220 }
1221 })
1222 } else {
1223 rows.remove_stale(|row, old, new| {
1224 let stale_entry = get_entry_mut(row, n_keys, hash, |x| x == old)
1225 .expect("non-stale entry not mapped in hash");
1226 *stale_entry = new;
1227 })
1228 }
1229 }
1230
1231 fn rehash(&mut self) {
1232 self.generation = self.generation.inc();
1233 Self::rehash_impl(
1234 self.sort_by,
1235 self.n_keys,
1236 &mut self.data,
1237 &mut self.offsets,
1238 &mut self.hash,
1239 )
1240 }
1241}
1242
1243fn get_entry(
1244 row: &[Value],
1245 n_keys: usize,
1246 table: &ShardedHashTable<TableEntry>,
1247 test: impl Fn(RowId) -> bool,
1248) -> Option<RowId> {
1249 let (shard, hash) = hash_code(table.shard_data(), row, n_keys);
1250 table
1251 .get_shard(shard)
1252 .find(hash, |ent| {
1253 ent.hashcode == hash as HashCode && test(ent.row)
1254 })
1255 .map(|ent| ent.row)
1256}
1257
1258fn get_entry_mut<'a>(
1259 row: &[Value],
1260 n_keys: usize,
1261 table: &'a mut ShardedHashTable<TableEntry>,
1262 test: impl Fn(RowId) -> bool,
1263) -> Option<&'a mut RowId> {
1264 let (shard, hash) = hash_code(table.shard_data(), row, n_keys);
1265 table.mut_shards()[shard.index()]
1266 .find_mut(hash, |ent| {
1267 ent.hashcode == hash as HashCode && test(ent.row)
1268 })
1269 .map(|ent| &mut ent.row)
1270}
1271
1272fn hash_code(shard_data: ShardData, row: &[Value], n_keys: usize) -> (ShardId, u64) {
1273 let mut hasher = FxHasher::default();
1274 for val in &row[0..n_keys] {
1275 hasher.write_usize(val.index());
1276 }
1277 let full_code = hasher.finish();
1278 #[allow(clippy::unnecessary_cast)]
1280 (shard_data.shard_id(full_code), full_code as HashCode as u64)
1281}
1282
1283struct PendingState {
1285 pending_rows: DenseIdMap<ShardId, SegQueue<RowBuffer>>,
1286 pending_removals: DenseIdMap<ShardId, SegQueue<ArbitraryRowBuffer>>,
1287 total_removals: AtomicUsize,
1288 total_rows: AtomicUsize,
1289}
1290
1291impl PendingState {
1292 fn new(shard_data: ShardData) -> PendingState {
1293 let n_shards = shard_data.n_shards();
1294 let mut pending_rows = DenseIdMap::with_capacity(n_shards);
1295 let mut pending_removals = DenseIdMap::with_capacity(n_shards);
1296 for i in 0..n_shards {
1297 pending_rows.insert(ShardId::from_usize(i), SegQueue::default());
1298 pending_removals.insert(ShardId::from_usize(i), SegQueue::default());
1299 }
1300
1301 PendingState {
1302 pending_rows,
1303 pending_removals,
1304 total_removals: AtomicUsize::new(0),
1305 total_rows: AtomicUsize::new(0),
1306 }
1307 }
1308 fn clear(&self) {
1309 for (_, queue) in self.pending_rows.iter() {
1310 while queue.pop().is_some() {}
1311 }
1312
1313 for (_, queue) in self.pending_removals.iter() {
1314 while queue.pop().is_some() {}
1315 }
1316 }
1317
1318 fn deep_copy(&self) -> PendingState {
1324 let mut pending_rows = DenseIdMap::new();
1325 let mut pending_removals = DenseIdMap::new();
1326 fn drain_queue<T>(queue: &SegQueue<T>) -> Vec<T> {
1327 let mut res = Vec::new();
1328 while let Some(x) = queue.pop() {
1329 res.push(x);
1330 }
1331 res
1332 }
1333 for (shard, queue) in self.pending_rows.iter() {
1334 let contents = drain_queue(queue);
1335 let new_queue = SegQueue::default();
1336 for x in contents {
1337 new_queue.push(x.clone());
1338 queue.push(x);
1339 }
1340 pending_rows.insert(shard, new_queue);
1341 }
1342
1343 for (shard, queue) in self.pending_removals.iter() {
1344 let contents = drain_queue(queue);
1345 let new_queue = SegQueue::default();
1346 for x in contents {
1347 new_queue.push(x.clone());
1348 queue.push(x);
1349 }
1350 pending_removals.insert(shard, new_queue);
1351 }
1352
1353 PendingState {
1354 pending_rows,
1355 pending_removals,
1356 total_removals: AtomicUsize::new(self.total_removals.load(Ordering::Acquire)),
1357 total_rows: AtomicUsize::new(self.total_rows.load(Ordering::Acquire)),
1358 }
1359 }
1360}
1361
1362trait OrderingChecker: Clone + Send + Sync {
1370 fn check_local(&mut self, row: &[Value]);
1373 fn check_global<'a>(checkers: impl Iterator<Item = &'a Self>) -> Self
1377 where
1378 Self: 'a;
1379 fn update_offsets(&self, start: RowId, offsets: &mut Vec<(Value, RowId)>);
1381}
1382
1383impl OrderingChecker for () {
1384 fn check_local(&mut self, _: &[Value]) {}
1385 fn check_global<'a>(_: impl Iterator<Item = &'a ()>) {}
1386 fn update_offsets(&self, _: RowId, _: &mut Vec<(Value, RowId)>) {}
1387}
1388
1389#[derive(Copy, Clone)]
1390struct SortChecker {
1391 col: ColumnId,
1392 baseline: Option<Value>,
1393 current: Option<Value>,
1394}
1395
1396impl OrderingChecker for SortChecker {
1397 fn check_local(&mut self, row: &[Value]) {
1398 let val = row[self.col.index()];
1399 if let Some(cur) = self.current {
1400 assert_eq!(
1401 cur, val,
1402 "concurrently inserting rows with different sort keys"
1403 );
1404 } else {
1405 self.current = Some(val);
1406 if let Some(baseline) = self.baseline {
1407 assert!(val >= baseline, "inserted row violates sort order");
1408 }
1409 }
1410 }
1411
1412 fn check_global<'a>(mut checkers: impl Iterator<Item = &'a Self>) -> Self {
1413 let Some(start) = checkers.next() else {
1414 return SortChecker {
1415 col: ColumnId::new(!0),
1416 baseline: None,
1417 current: None,
1418 };
1419 };
1420 let mut expected = start.current;
1421 for checker in checkers {
1422 assert_eq!(checker.baseline, start.baseline);
1423 match (&mut expected, checker.current) {
1424 (None, None) => {}
1425 (cur @ None, Some(x)) => {
1426 *cur = Some(x);
1427 }
1428 (Some(_), None) => {}
1429 (Some(x), Some(y)) => {
1430 assert_eq!(
1431 *x, y,
1432 "concurrently inserting rows with different sort keys"
1433 );
1434 }
1435 }
1436 }
1437 SortChecker {
1438 col: start.col,
1439 baseline: start.baseline,
1440 current: expected,
1441 }
1442 }
1443
1444 fn update_offsets(&self, start: RowId, offsets: &mut Vec<(Value, RowId)>) {
1445 if let Some(cur) = self.current {
1446 if let Some((max, _)) = offsets.last() {
1447 if cur > *max {
1448 offsets.push((cur, start));
1449 }
1450 } else {
1451 offsets.push((cur, start));
1452 }
1453 }
1454 }
1455}
1456
1457struct StagedOutputs {
1462 shard_data: ShardData,
1463 n_keys: usize,
1464 hash: Pooled<HashTable<TableEntry>>,
1465 rows: RowBuffer,
1466 n_stale: usize,
1467 scratch: Pooled<Vec<Value>>,
1468}
1469
1470impl StagedOutputs {
1471 fn rows(&self) -> impl Iterator<Item = &[Value]> {
1472 self.rows.iter()
1473 }
1474 fn new(n_keys: usize, n_cols: usize, capacity: usize) -> Self {
1475 let mut res = with_pool_set(|ps| StagedOutputs {
1476 shard_data: ShardData::new(1),
1477 n_keys,
1478 n_stale: 0,
1479 hash: ps.get(),
1480 rows: RowBuffer::new(n_cols),
1481 scratch: ps.get(),
1482 });
1483 res.hash.reserve(capacity, TableEntry::hashcode);
1484 res.rows.reserve(capacity);
1485 res
1486 }
1487 fn clear(&mut self) {
1488 self.hash.clear();
1489 self.rows.clear();
1490 self.n_stale = 0;
1491 }
1492 fn len(&self) -> usize {
1493 self.rows.len() - self.n_stale
1494 }
1495
1496 fn insert(
1497 &mut self,
1498 row: &[Value],
1499 mut merge_fn: impl FnMut(&[Value], &[Value], &mut Vec<Value>) -> bool,
1500 ) {
1501 if row[0].is_stale() {
1502 return;
1503 }
1504 use hashbrown::hash_table::Entry;
1505 let (_, hc) = hash_code(self.shard_data, row, self.n_keys);
1506 let entry = self.hash.entry(
1507 hc,
1508 |te| {
1509 te.hashcode() == hc
1510 && self.rows.get_row(te.row)[0..self.n_keys] == row[0..self.n_keys]
1511 },
1512 TableEntry::hashcode,
1513 );
1514 match entry {
1515 Entry::Occupied(mut occupied_entry) => {
1516 let cur = self.rows.get_row(occupied_entry.get().row);
1517 if merge_fn(cur, row, &mut self.scratch) {
1518 let new = self.rows.add_row(&self.scratch);
1519 self.rows.set_stale(occupied_entry.get().row);
1520 self.n_stale += 1;
1521 occupied_entry.get_mut().row = new;
1522 }
1523 self.scratch.clear();
1524 }
1525 Entry::Vacant(vacant_entry) => {
1526 let next = self.rows.add_row(row);
1527 vacant_entry.insert(TableEntry {
1528 hashcode: hc as _,
1529 row: next,
1530 });
1531 }
1532 }
1533 }
1534
1535 fn write_output(&self, output: &ParallelRowBufWriter) -> (RowId, usize) {
1538 (output.append_contents(&self.rows), self.n_stale)
1539 }
1540}