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(
388 &self,
389 subset: SubsetRef,
390 start: Offset,
391 n: usize,
392 cs: &[Constraint],
393 mut f: impl FnMut(RowId, &[Value]),
394 ) -> Option<Offset>
395 where
396 Self: Sized,
397 {
398 let Some((_low, hi)) = subset.bounds() else {
399 return None;
401 };
402 assert!(
403 hi.index() <= self.data.data.len(),
404 "{} vs. {}",
405 hi.index(),
406 self.data.data.len()
407 );
408 if cs.is_empty() {
409 if self.data.stale_rows == 0 {
410 subset
413 .iter_bounded(start.index(), start.index() + n, |row| {
414 let entry = unsafe { self.data.data.get_row_unchecked(row) };
415 f(row, entry);
416 })
417 .map(Offset::from_usize)
418 } else {
419 subset
420 .iter_bounded(start.index(), start.index() + n, |row| {
421 let Some(entry) = (unsafe { self.data.get_row_unchecked(row) }) else {
423 return;
424 };
425 f(row, entry);
426 })
427 .map(Offset::from_usize)
428 }
429 } else {
430 subset
431 .iter_bounded(start.index(), start.index() + n, |row| {
432 let Some(entry) = (unsafe { self.get_if_unchecked(cs, row) }) else {
434 return;
435 };
436 f(row, entry);
437 })
438 .map(Offset::from_usize)
439 }
440 }
441
442 fn fast_subset(&self, constraint: &Constraint) -> Option<Subset> {
443 let sort_by = self.sort_by?;
444 match constraint {
445 Constraint::Eq { .. } => None,
446 Constraint::EqConst { col, val } => {
447 if col == &sort_by {
448 match self.binary_search_sort_val(*val) {
449 Ok((found, bound)) => Some(Subset::Dense(OffsetRange::new(found, bound))),
450 Err(_) => Some(Subset::empty()),
451 }
452 } else {
453 None
454 }
455 }
456 Constraint::LtConst { col, val } => {
457 if col == &sort_by {
458 match self.binary_search_sort_val(*val) {
459 Ok((found, _)) => {
460 Some(Subset::Dense(OffsetRange::new(RowId::new(0), found)))
461 }
462 Err(next) => Some(Subset::Dense(OffsetRange::new(RowId::new(0), next))),
463 }
464 } else {
465 None
466 }
467 }
468 Constraint::GtConst { col, val } => {
469 if col == &sort_by {
470 match self.binary_search_sort_val(*val) {
471 Ok((_, bound)) => {
472 Some(Subset::Dense(OffsetRange::new(bound, self.data.next_row())))
473 }
474 Err(next) => {
475 Some(Subset::Dense(OffsetRange::new(next, self.data.next_row())))
476 }
477 }
478 } else {
479 None
480 }
481 }
482 Constraint::LeConst { col, val } => {
483 if col == &sort_by {
484 match self.binary_search_sort_val(*val) {
485 Ok((_, bound)) => {
486 Some(Subset::Dense(OffsetRange::new(RowId::new(0), bound)))
487 }
488 Err(next) => Some(Subset::Dense(OffsetRange::new(RowId::new(0), next))),
489 }
490 } else {
491 None
492 }
493 }
494 Constraint::GeConst { col, val } => {
495 if col == &sort_by {
496 match self.binary_search_sort_val(*val) {
497 Ok((found, _)) => {
498 Some(Subset::Dense(OffsetRange::new(found, self.data.next_row())))
499 }
500 Err(next) => {
501 Some(Subset::Dense(OffsetRange::new(next, self.data.next_row())))
502 }
503 }
504 } else {
505 None
506 }
507 }
508 }
509 }
510
511 fn refine_one(&self, mut subset: Subset, c: &Constraint) -> Subset {
512 subset.retain(|row| self.eval(std::slice::from_ref(c), row));
515 subset
516 }
517
518 fn refine_ref(&self, subset: SubsetRef, cs: &[Constraint], _check_live: bool) -> Subset {
519 let n = subset.size();
523 let mut vec: Pooled<SortedOffsetVector> = with_pool_set(|ps| ps.get());
524 subset.offsets(|row| {
525 if self.eval(cs, row) {
526 unsafe { vec.push_unchecked(row) }
528 }
529 });
530 if let SubsetRef::Dense(range) = subset
531 && vec.slice().inner().len() == n
532 {
533 return Subset::Dense(range);
534 }
535 Subset::Sparse(vec)
536 }
537
538 fn new_buffer(&self) -> Box<dyn MutationBuffer> {
539 let n_shards = self.hash.shard_data().n_shards();
540 Box::new(Buffer {
541 pending_rows: DenseIdMap::with_capacity(n_shards),
542 pending_removals: DenseIdMap::with_capacity(n_shards),
543 state: Arc::downgrade(&self.pending_state),
544 n_keys: u32::try_from(self.n_keys).expect("n_keys should fit in u32"),
545 n_cols: u32::try_from(self.n_columns).expect("n_columns should fit in u32"),
546 shard_data: self.hash.shard_data(),
547 })
548 }
549
550 fn merge(&mut self, exec_state: &mut ExecutionState) -> TableChange {
551 let removed = self.do_delete();
552 let added = self.do_insert(exec_state);
553 self.maybe_rehash();
554 TableChange { removed, added }
555 }
556
557 fn get_row(&self, key: &[Value]) -> Option<Row> {
558 let id = get_entry(key, self.n_keys, &self.hash, |row| {
559 &self.data.get_row(row).unwrap()[0..self.n_keys] == key
560 })?;
561 let mut vals = with_pool_set(|ps| ps.get::<Vec<Value>>());
562 vals.extend_from_slice(self.data.get_row(id).unwrap());
563 Some(Row { id, vals })
564 }
565
566 fn get_row_column(&self, key: &[Value], col: ColumnId) -> Option<Value> {
567 let id = get_entry(key, self.n_keys, &self.hash, |row| {
568 &self.data.get_row(row).unwrap()[0..self.n_keys] == key
569 })?;
570 Some(self.data.get_row(id).unwrap()[col.index()])
571 }
572}
573
574impl SortedWritesTable {
575 pub fn new(
587 n_keys: usize,
588 n_columns: usize,
589 sort_by: Option<ColumnId>,
590 to_rebuild: Vec<ColumnId>,
591 merge_fn: Box<MergeFn>,
592 ) -> Self {
593 let hash = ShardedHashTable::<TableEntry>::default();
594 let shard_data = hash.shard_data();
595 let rebuild_index = Index::new(to_rebuild.clone(), ColumnIndex::new());
596 SortedWritesTable {
597 generation: Generation::new(0),
598 data: Rows::new(RowBuffer::new(n_columns)),
599 hash,
600 n_keys,
601 n_columns,
602 sort_by,
603 offsets: Default::default(),
604 pending_state: Arc::new(PendingState::new(shard_data)),
605 merge: merge_fn.into(),
606 to_rebuild,
607 rebuild_index,
608 subset_tracker: Default::default(),
609 }
610 }
611
612 fn parallel_delete(&mut self) -> bool {
614 let shard_data = self.hash.shard_data();
615 let pending_removals = &self.pending_state.pending_removals;
616 let data = &self.data.data;
617 let n_keys = self.n_keys;
618 let stale_delta: usize = parallel::map_mut(self.hash.mut_shards(), |shard_id, shard| {
619 let shard_id = ShardId::from_usize(shard_id);
620 if pending_removals[shard_id].is_empty() {
621 return 0;
622 }
623 let queue = &pending_removals[shard_id];
624 let mut marked_stale = 0;
625 while let Some(buf) = queue.pop() {
626 buf.for_each(|to_remove| {
627 let (actual_shard, hc) = hash_code(shard_data, to_remove, n_keys);
628 assert_eq!(actual_shard, shard_id);
629 if let Ok(entry) = shard.find_entry(hc, |entry| {
630 entry.hashcode == (hc as _)
631 && &data.get_row(entry.row)[0..n_keys] == to_remove
632 }) {
633 let (ent, _) = entry.remove();
634 marked_stale += unsafe { !data.set_stale_shared(ent.row) } as usize;
642 }
643 });
644 }
645 marked_stale
646 })
647 .into_iter()
648 .sum();
649 self.data.stale_rows += stale_delta;
651 stale_delta > 0
652 }
653 fn serial_delete(&mut self) -> bool {
654 let shard_data = self.hash.shard_data();
655 let mut changed = false;
656 self.hash
657 .mut_shards()
658 .iter_mut()
659 .enumerate()
660 .for_each(|(shard_id, shard)| {
661 let shard_id = ShardId::from_usize(shard_id);
662 let queue = &self.pending_state.pending_removals[shard_id];
663 while let Some(buf) = queue.pop() {
664 buf.for_each(|to_remove| {
665 let (actual_shard, hc) = hash_code(shard_data, to_remove, self.n_keys);
666 assert_eq!(actual_shard, shard_id);
667 if let Ok(entry) = shard.find_entry(hc, |entry| {
668 entry.hashcode == (hc as _)
669 && &self.data.get_row(entry.row).unwrap()[0..self.n_keys]
670 == to_remove
671 }) {
672 let (ent, _) = entry.remove();
673 self.data.set_stale(ent.row);
674 changed = true;
675 }
676 })
677 }
678 });
679 changed
680 }
681
682 fn do_delete(&mut self) -> bool {
683 let total = self.pending_state.total_removals.swap(0, Ordering::Relaxed);
684
685 if parallelize_table_op(total) {
686 self.parallel_delete()
687 } else {
688 self.serial_delete()
689 }
690 }
691
692 fn do_insert(&mut self, exec_state: &mut ExecutionState) -> bool {
693 let total = self.pending_state.total_rows.swap(0, Ordering::Relaxed);
694 self.data.data.reserve(total);
695 if parallelize_table_op(total) {
696 if let Some(col) = self.sort_by {
697 self.parallel_insert(
698 exec_state,
699 SortChecker {
700 col,
701 current: None,
702 baseline: self.offsets.last().map(|(v, _)| *v),
703 },
704 )
705 } else {
706 self.parallel_insert(exec_state, ())
707 }
708 } else {
709 self.serial_insert(exec_state)
710 }
711 }
712
713 fn serial_insert(&mut self, exec_state: &mut ExecutionState) -> bool {
714 let mut changed = false;
715 let n_keys = self.n_keys;
716 let mut scratch = with_pool_set(|ps| ps.get::<Vec<Value>>());
717 for (_outer_shard, queue) in self.pending_state.pending_rows.iter() {
718 if let Some(sort_by) = self.sort_by {
719 while let Some(buf) = queue.pop() {
720 for query in buf.non_stale() {
721 let key = &query[0..n_keys];
722 let entry = get_entry_mut(query, n_keys, &mut self.hash, |row| {
723 let Some(row) = self.data.get_row(row) else {
724 return false;
725 };
726 &row[0..n_keys] == key
727 });
728
729 if let Some(row) = entry {
730 let cur = self
734 .data
735 .get_row(*row)
736 .expect("table should not point to stale entry");
737 if (self.merge)(exec_state, cur, query, &mut scratch) {
738 let sort_val = query[sort_by.index()];
739 let new = self.data.add_row(&scratch);
740 if let Some(largest) = self.offsets.last().map(|(v, _)| *v) {
741 assert!(
742 sort_val >= largest,
743 "inserting row that violates sort order ({sort_val:?} vs. {largest:?})"
744 );
745 if sort_val > largest {
746 self.offsets.push((sort_val, new));
747 }
748 } else {
749 self.offsets.push((sort_val, new));
750 }
751 self.data.set_stale(*row);
752 *row = new;
753 changed = true;
754 }
755 scratch.clear();
756 } else {
757 let sort_val = query[sort_by.index()];
758 let new = self.data.add_row(query);
760 if let Some(largest) = self.offsets.last().map(|(v, _)| *v) {
761 assert!(
762 sort_val >= largest,
763 "inserting row that violates sort order {sort_val:?} vs. {largest:?}"
764 );
765 if sort_val > largest {
766 self.offsets.push((sort_val, new));
767 }
768 } else {
769 self.offsets.push((sort_val, new));
770 }
771 let (shard, hc) = hash_code(self.hash.shard_data(), query, self.n_keys);
772 debug_assert_eq!(shard, _outer_shard);
773 self.hash.mut_shards()[shard.index()].insert_unique(
774 hc as _,
775 TableEntry {
776 hashcode: hc as _,
777 row: new,
778 },
779 TableEntry::hashcode,
780 );
781 changed = true;
782 }
783 }
784 }
785 } else {
786 while let Some(buf) = queue.pop() {
788 for query in buf.non_stale() {
789 let key = &query[0..n_keys];
790 let entry = get_entry_mut(query, n_keys, &mut self.hash, |row| {
791 let Some(row) = self.data.get_row(row) else {
792 return false;
793 };
794 &row[0..n_keys] == key
795 });
796
797 if let Some(row) = entry {
798 let cur = self
799 .data
800 .get_row(*row)
801 .expect("table should not point to stale entry");
802 if (self.merge)(exec_state, cur, query, &mut scratch) {
803 let new = self.data.add_row(&scratch);
804 self.data.set_stale(*row);
805 *row = new;
806 changed = true;
807 }
808 scratch.clear();
809 } else {
810 let new = self.data.add_row(query);
812 let (shard, hc) = hash_code(self.hash.shard_data(), query, self.n_keys);
813 debug_assert_eq!(shard, _outer_shard);
814 self.hash.mut_shards()[shard.index()].insert_unique(
815 hc as _,
816 TableEntry {
817 hashcode: hc as _,
818 row: new,
819 },
820 TableEntry::hashcode,
821 );
822 changed = true;
823 }
824 }
825 }
826 };
827 }
828 changed
829 }
830
831 fn parallel_insert<C: OrderingChecker>(
832 &mut self,
833 exec_state: &ExecutionState,
834 checker: C,
835 ) -> bool {
836 const BATCH_SIZE: usize = 1 << 18;
837 let shard_data = self.hash.shard_data();
842 let n_keys = self.n_keys;
843 let n_cols = self.n_columns;
844 let next_offset = RowId::from_usize(self.data.data.len());
845 let row_writer = self.data.data.parallel_writer();
846 let pending_rows = &self.pending_state.pending_rows;
847 let merge = self.merge.clone();
848 let pending_adds = parallel::map_mut(self.hash.mut_shards(), |shard_id, shard| {
849 let shard_id = ShardId::from_usize(shard_id);
850 let mut checker = checker.clone();
851 let mut exec_state = exec_state.clone();
852 let mut scratch = with_pool_set(|ps| ps.get::<Vec<Value>>());
853 let queue = &pending_rows[shard_id];
854 let mut marked_stale = 0usize;
855 let mut staged = StagedOutputs::new(n_keys, n_cols, BATCH_SIZE);
856 let mut changed = false;
857 macro_rules! flush_staged_outputs {
860 () => {{
861 let (start_row, stale) = staged.write_output(&row_writer);
864 marked_stale += stale;
865 let mut cur_row = start_row;
874 let read_handle = row_writer.read_handle();
875 for row in staged.rows() {
876 if row.first().map(Value::is_stale).unwrap_or(false) {
877 cur_row = cur_row.inc();
878 continue;
879 }
880 use hashbrown::hash_table::Entry;
881 checker.check_local(row);
882 changed = true;
883 let key = &row[0..n_keys];
884 let (_actual_shard, hc) = hash_code(shard_data, row, n_keys);
885 #[cfg(any(debug_assertions, test))]
886 {
887 unsafe {
888 assert_eq!(read_handle.get_row_unchecked(cur_row), row);
891 }
892 }
893 debug_assert_eq!(_actual_shard, shard_id);
894 match shard.entry(
895 hc,
896 |ent| unsafe {
898 ent.hashcode == hc as HashCode
899 && &read_handle.get_row_unchecked(ent.row)[0..n_keys] == key
900 },
901 TableEntry::hashcode,
902 ) {
903 Entry::Occupied(mut occ) => {
904 let cur = unsafe { read_handle.get_row_unchecked(occ.get().row) };
907
908 if (merge)(&mut exec_state, cur, row, &mut scratch) {
914 unsafe {
915 let _was_stale = read_handle.set_stale_shared(occ.get().row);
916 debug_assert!(!_was_stale);
917 }
918 occ.get_mut().row = cur_row;
919 changed = true;
920 } else {
921 unsafe {
923 let _was_stale = read_handle.set_stale_shared(cur_row);
924 debug_assert!(!_was_stale);
925 }
926 }
927 marked_stale += 1;
928 scratch.clear();
929 }
930 Entry::Vacant(v) => {
931 changed = true;
932 v.insert(TableEntry {
933 hashcode: hc as HashCode,
934 row: cur_row,
935 });
936 }
937 }
938
939 cur_row = cur_row.inc();
940 }
941 staged.clear();
942 }};
943 }
944 while let Some(buf) = queue.pop() {
949 for row in buf.non_stale() {
953 staged.insert(row, |cur, new, out| (merge)(&mut exec_state, cur, new, out));
954 if staged.len() >= BATCH_SIZE {
955 flush_staged_outputs!();
956 }
957 }
958 }
959 flush_staged_outputs!();
960 (checker, marked_stale, changed)
961 });
962 self.data.data = row_writer.finish();
963 let checker = C::check_global(pending_adds.iter().map(|(checker, _, _)| checker));
968 checker.update_offsets(next_offset, &mut self.offsets);
969
970 self.data.stale_rows += pending_adds
972 .iter()
973 .map(|(_, stale, _)| *stale)
974 .sum::<usize>();
975
976 pending_adds.iter().any(|(_, _, changed)| *changed)
978 }
979
980 fn binary_search_sort_val(&self, val: Value) -> Result<(RowId, RowId), RowId> {
981 debug_assert!(
982 self.offsets.windows(2).all(|x| x[0].1 < x[1].1),
983 "{:?}",
984 self.offsets
985 );
986
987 debug_assert!(
988 self.offsets.windows(2).all(|x| x[0].0 < x[1].0),
989 "{:?}",
990 self.offsets
991 );
992 match self.offsets.binary_search_by_key(&val, |(v, _)| *v) {
993 Ok(got) => Ok((
994 self.offsets[got].1,
995 self.offsets
996 .get(got + 1)
997 .map(|(_, r)| *r)
998 .unwrap_or(self.data.next_row()),
999 )),
1000 Err(next) => Err(self
1001 .offsets
1002 .get(next)
1003 .map(|(_, id)| *id)
1004 .unwrap_or(self.data.next_row())),
1005 }
1006 }
1007 fn eval(&self, cs: &[Constraint], row: RowId) -> bool {
1008 self.get_if(cs, row).is_some()
1009 }
1010
1011 fn eval_constraints(cs: &[Constraint], row: &[Value]) -> bool {
1012 cs.iter().all(|constraint| match constraint {
1013 Constraint::Eq { l_col, r_col } => row[l_col.index()] == row[r_col.index()],
1014 Constraint::EqConst { col, val } => row[col.index()] == *val,
1015 Constraint::LtConst { col, val } => row[col.index()] < *val,
1016 Constraint::GtConst { col, val } => row[col.index()] > *val,
1017 Constraint::LeConst { col, val } => row[col.index()] <= *val,
1018 Constraint::GeConst { col, val } => row[col.index()] >= *val,
1019 })
1020 }
1021
1022 unsafe fn get_if_unchecked(&self, cs: &[Constraint], row: RowId) -> Option<&[Value]> {
1023 let row = unsafe { self.data.data.get_row_unchecked(row) };
1024 if Self::eval_constraints(cs, row) {
1025 Some(row)
1026 } else {
1027 None
1028 }
1029 }
1030
1031 fn get_if(&self, cs: &[Constraint], row: RowId) -> Option<&[Value]> {
1032 let row = self.data.get_row(row)?;
1033 if Self::eval_constraints(cs, row) {
1034 Some(row)
1035 } else {
1036 None
1037 }
1038 }
1039
1040 fn maybe_rehash(&mut self) {
1041 if self.data.stale_rows <= cmp::max(16, self.data.data.len() / 2) {
1042 return;
1043 }
1044
1045 if parallelize_table_op(self.data.data.len()) {
1046 self.parallel_rehash();
1047 } else {
1048 self.rehash();
1049 }
1050 }
1051 fn parallel_rehash(&mut self) {
1052 let Some(sort_by) = self.sort_by else {
1056 self.rehash();
1063 return;
1064 };
1065 self.generation = self.generation.inc();
1066 assert!(!self.offsets.is_empty());
1067 struct TimestampStats {
1068 value: Value,
1069 count: usize,
1070 histogram: Pooled<DenseIdMap<ShardId, usize>>,
1071 }
1072 impl Default for TimestampStats {
1073 fn default() -> TimestampStats {
1074 TimestampStats {
1075 value: Value::stale(),
1076 count: 0,
1077 histogram: with_pool_set(|ps| ps.get()),
1078 }
1079 }
1080 }
1081 let mut results = Vec::<TimestampStats>::with_capacity(self.offsets.len());
1082 let offset_windows = self
1083 .offsets
1084 .windows(2)
1085 .map(|xs| {
1086 let [(start_val, start_row), (_, end_row)] = xs else {
1087 unreachable!()
1088 };
1089 (*start_val, *start_row, *end_row)
1090 })
1091 .collect::<Vec<_>>();
1092 macro_rules! compute_hist {
1094 ($start_val: expr, $start_row: expr, $end_row: expr) => {{
1095 let mut histogram: Pooled<DenseIdMap<ShardId, usize>> =
1096 with_pool_set(|ps| ps.get());
1097 let mut cur_row = $start_row;
1098 let mut count = 0;
1099 while cur_row < $end_row {
1100 if let Some(row) = self.data.get_row(cur_row) {
1101 count += 1;
1102 let (shard, _) = hash_code(self.hash.shard_data(), row, self.n_keys);
1103 *histogram.get_or_default(shard) += 1;
1104 }
1105 cur_row = cur_row.inc();
1106 }
1107 TimestampStats {
1108 value: $start_val,
1109 count,
1110 histogram,
1111 }
1112 }};
1113 }
1114 results.extend(parallel::map(
1115 &offset_windows,
1116 |_, (start_val, start_row, end_row)| compute_hist!(*start_val, *start_row, *end_row),
1117 ));
1118 let (start_val, start_row) = self.offsets.last().unwrap();
1120 let end_row = self.data.next_row();
1121 let last = compute_hist!(*start_val, *start_row, end_row);
1122 results.push(last);
1123 let mut prev_count = 0;
1128 self.offsets.clear();
1129 for stats in results.iter_mut() {
1130 if stats.count == 0 {
1131 continue;
1132 }
1133 self.offsets
1134 .push((stats.value, RowId::from_usize(prev_count)));
1135 let mut inner = prev_count;
1136 for (_, count) in stats.histogram.iter_mut() {
1137 let tmp = *count;
1140 *count = inner;
1141 inner += tmp;
1142 }
1143 prev_count += stats.count;
1144 debug_assert_eq!(inner, prev_count)
1145 }
1146
1147 self.data.scratch.clear();
1160 self.data.scratch.reserve(prev_count);
1161 let scratch_ptr = self.data.scratch.raw_rows() as usize;
1162 let data = &self.data.data;
1163 let n_columns = self.n_columns;
1164 parallel::for_each_mut(self.hash.mut_shards(), |shard_id, shard| {
1165 let shard_id = ShardId::from_usize(shard_id);
1166 let scratch_ptr = scratch_ptr as *const Value;
1167 let mut progress = HashMap::<Value , RowId >::default();
1168 progress.reserve(results.len());
1169 for stats in &results {
1170 let Some(start) = stats.histogram.get(shard_id) else {
1171 continue;
1172 };
1173 progress.insert(stats.value, RowId::from_usize(*start));
1174 }
1175 for TableEntry { row: row_id, .. } in shard.iter_mut() {
1176 let row = data.get_row(*row_id);
1177 debug_assert!(!row[0].is_stale(), "shard should not map to a stale value");
1178 let val = row[sort_by.index()];
1179 let next = progress[&val];
1180 unsafe {
1182 std::ptr::copy_nonoverlapping(
1183 row.as_ptr(),
1184 scratch_ptr.add(next.index() * n_columns) as *mut Value,
1185 n_columns,
1186 )
1187 }
1188 *row_id = next;
1189 progress.insert(val, next.inc());
1190 }
1191 });
1192 unsafe { self.data.scratch.set_len(prev_count) };
1194 mem::swap(&mut self.data.data, &mut self.data.scratch);
1195 self.data.stale_rows = 0;
1196 }
1197 fn rehash_impl(
1198 sort_by: Option<ColumnId>,
1199 n_keys: usize,
1200 rows: &mut Rows,
1201 offsets: &mut Vec<(Value, RowId)>,
1202 hash: &mut ShardedHashTable<TableEntry>,
1203 ) {
1204 if let Some(sort_by) = sort_by {
1205 offsets.clear();
1206 rows.remove_stale(|row, old, new| {
1207 let stale_entry = get_entry_mut(row, n_keys, hash, |x| x == old)
1208 .expect("non-stale entry not mapped in hash");
1209 *stale_entry = new;
1210 let sort_col = row[sort_by.index()];
1211 if let Some((max, _)) = offsets.last() {
1212 if sort_col > *max {
1213 offsets.push((sort_col, new));
1214 }
1215 } else {
1216 offsets.push((sort_col, new));
1217 }
1218 })
1219 } else {
1220 rows.remove_stale(|row, old, new| {
1221 let stale_entry = get_entry_mut(row, n_keys, hash, |x| x == old)
1222 .expect("non-stale entry not mapped in hash");
1223 *stale_entry = new;
1224 })
1225 }
1226 }
1227
1228 fn rehash(&mut self) {
1229 self.generation = self.generation.inc();
1230 Self::rehash_impl(
1231 self.sort_by,
1232 self.n_keys,
1233 &mut self.data,
1234 &mut self.offsets,
1235 &mut self.hash,
1236 )
1237 }
1238}
1239
1240fn get_entry(
1241 row: &[Value],
1242 n_keys: usize,
1243 table: &ShardedHashTable<TableEntry>,
1244 test: impl Fn(RowId) -> bool,
1245) -> Option<RowId> {
1246 let (shard, hash) = hash_code(table.shard_data(), row, n_keys);
1247 table
1248 .get_shard(shard)
1249 .find(hash, |ent| {
1250 ent.hashcode == hash as HashCode && test(ent.row)
1251 })
1252 .map(|ent| ent.row)
1253}
1254
1255fn get_entry_mut<'a>(
1256 row: &[Value],
1257 n_keys: usize,
1258 table: &'a mut ShardedHashTable<TableEntry>,
1259 test: impl Fn(RowId) -> bool,
1260) -> Option<&'a mut RowId> {
1261 let (shard, hash) = hash_code(table.shard_data(), row, n_keys);
1262 table.mut_shards()[shard.index()]
1263 .find_mut(hash, |ent| {
1264 ent.hashcode == hash as HashCode && test(ent.row)
1265 })
1266 .map(|ent| &mut ent.row)
1267}
1268
1269fn hash_code(shard_data: ShardData, row: &[Value], n_keys: usize) -> (ShardId, u64) {
1270 let mut hasher = FxHasher::default();
1271 for val in &row[0..n_keys] {
1272 hasher.write_usize(val.index());
1273 }
1274 let full_code = hasher.finish();
1275 #[allow(clippy::unnecessary_cast)]
1277 (shard_data.shard_id(full_code), full_code as HashCode as u64)
1278}
1279
1280struct PendingState {
1282 pending_rows: DenseIdMap<ShardId, SegQueue<RowBuffer>>,
1283 pending_removals: DenseIdMap<ShardId, SegQueue<ArbitraryRowBuffer>>,
1284 total_removals: AtomicUsize,
1285 total_rows: AtomicUsize,
1286}
1287
1288impl PendingState {
1289 fn new(shard_data: ShardData) -> PendingState {
1290 let n_shards = shard_data.n_shards();
1291 let mut pending_rows = DenseIdMap::with_capacity(n_shards);
1292 let mut pending_removals = DenseIdMap::with_capacity(n_shards);
1293 for i in 0..n_shards {
1294 pending_rows.insert(ShardId::from_usize(i), SegQueue::default());
1295 pending_removals.insert(ShardId::from_usize(i), SegQueue::default());
1296 }
1297
1298 PendingState {
1299 pending_rows,
1300 pending_removals,
1301 total_removals: AtomicUsize::new(0),
1302 total_rows: AtomicUsize::new(0),
1303 }
1304 }
1305 fn clear(&self) {
1306 for (_, queue) in self.pending_rows.iter() {
1307 while queue.pop().is_some() {}
1308 }
1309
1310 for (_, queue) in self.pending_removals.iter() {
1311 while queue.pop().is_some() {}
1312 }
1313 }
1314
1315 fn deep_copy(&self) -> PendingState {
1321 let mut pending_rows = DenseIdMap::new();
1322 let mut pending_removals = DenseIdMap::new();
1323 fn drain_queue<T>(queue: &SegQueue<T>) -> Vec<T> {
1324 let mut res = Vec::new();
1325 while let Some(x) = queue.pop() {
1326 res.push(x);
1327 }
1328 res
1329 }
1330 for (shard, queue) in self.pending_rows.iter() {
1331 let contents = drain_queue(queue);
1332 let new_queue = SegQueue::default();
1333 for x in contents {
1334 new_queue.push(x.clone());
1335 queue.push(x);
1336 }
1337 pending_rows.insert(shard, new_queue);
1338 }
1339
1340 for (shard, queue) in self.pending_removals.iter() {
1341 let contents = drain_queue(queue);
1342 let new_queue = SegQueue::default();
1343 for x in contents {
1344 new_queue.push(x.clone());
1345 queue.push(x);
1346 }
1347 pending_removals.insert(shard, new_queue);
1348 }
1349
1350 PendingState {
1351 pending_rows,
1352 pending_removals,
1353 total_removals: AtomicUsize::new(self.total_removals.load(Ordering::Acquire)),
1354 total_rows: AtomicUsize::new(self.total_rows.load(Ordering::Acquire)),
1355 }
1356 }
1357}
1358
1359trait OrderingChecker: Clone + Send + Sync {
1367 fn check_local(&mut self, row: &[Value]);
1370 fn check_global<'a>(checkers: impl Iterator<Item = &'a Self>) -> Self
1374 where
1375 Self: 'a;
1376 fn update_offsets(&self, start: RowId, offsets: &mut Vec<(Value, RowId)>);
1378}
1379
1380impl OrderingChecker for () {
1381 fn check_local(&mut self, _: &[Value]) {}
1382 fn check_global<'a>(_: impl Iterator<Item = &'a ()>) {}
1383 fn update_offsets(&self, _: RowId, _: &mut Vec<(Value, RowId)>) {}
1384}
1385
1386#[derive(Copy, Clone)]
1387struct SortChecker {
1388 col: ColumnId,
1389 baseline: Option<Value>,
1390 current: Option<Value>,
1391}
1392
1393impl OrderingChecker for SortChecker {
1394 fn check_local(&mut self, row: &[Value]) {
1395 let val = row[self.col.index()];
1396 if let Some(cur) = self.current {
1397 assert_eq!(
1398 cur, val,
1399 "concurrently inserting rows with different sort keys"
1400 );
1401 } else {
1402 self.current = Some(val);
1403 if let Some(baseline) = self.baseline {
1404 assert!(val >= baseline, "inserted row violates sort order");
1405 }
1406 }
1407 }
1408
1409 fn check_global<'a>(mut checkers: impl Iterator<Item = &'a Self>) -> Self {
1410 let Some(start) = checkers.next() else {
1411 return SortChecker {
1412 col: ColumnId::new(!0),
1413 baseline: None,
1414 current: None,
1415 };
1416 };
1417 let mut expected = start.current;
1418 for checker in checkers {
1419 assert_eq!(checker.baseline, start.baseline);
1420 match (&mut expected, checker.current) {
1421 (None, None) => {}
1422 (cur @ None, Some(x)) => {
1423 *cur = Some(x);
1424 }
1425 (Some(_), None) => {}
1426 (Some(x), Some(y)) => {
1427 assert_eq!(
1428 *x, y,
1429 "concurrently inserting rows with different sort keys"
1430 );
1431 }
1432 }
1433 }
1434 SortChecker {
1435 col: start.col,
1436 baseline: start.baseline,
1437 current: expected,
1438 }
1439 }
1440
1441 fn update_offsets(&self, start: RowId, offsets: &mut Vec<(Value, RowId)>) {
1442 if let Some(cur) = self.current {
1443 if let Some((max, _)) = offsets.last() {
1444 if cur > *max {
1445 offsets.push((cur, start));
1446 }
1447 } else {
1448 offsets.push((cur, start));
1449 }
1450 }
1451 }
1452}
1453
1454struct StagedOutputs {
1459 shard_data: ShardData,
1460 n_keys: usize,
1461 hash: Pooled<HashTable<TableEntry>>,
1462 rows: RowBuffer,
1463 n_stale: usize,
1464 scratch: Pooled<Vec<Value>>,
1465}
1466
1467impl StagedOutputs {
1468 fn rows(&self) -> impl Iterator<Item = &[Value]> {
1469 self.rows.iter()
1470 }
1471 fn new(n_keys: usize, n_cols: usize, capacity: usize) -> Self {
1472 let mut res = with_pool_set(|ps| StagedOutputs {
1473 shard_data: ShardData::new(1),
1474 n_keys,
1475 n_stale: 0,
1476 hash: ps.get(),
1477 rows: RowBuffer::new(n_cols),
1478 scratch: ps.get(),
1479 });
1480 res.hash.reserve(capacity, TableEntry::hashcode);
1481 res.rows.reserve(capacity);
1482 res
1483 }
1484 fn clear(&mut self) {
1485 self.hash.clear();
1486 self.rows.clear();
1487 self.n_stale = 0;
1488 }
1489 fn len(&self) -> usize {
1490 self.rows.len() - self.n_stale
1491 }
1492
1493 fn insert(
1494 &mut self,
1495 row: &[Value],
1496 mut merge_fn: impl FnMut(&[Value], &[Value], &mut Vec<Value>) -> bool,
1497 ) {
1498 if row[0].is_stale() {
1499 return;
1500 }
1501 use hashbrown::hash_table::Entry;
1502 let (_, hc) = hash_code(self.shard_data, row, self.n_keys);
1503 let entry = self.hash.entry(
1504 hc,
1505 |te| {
1506 te.hashcode() == hc
1507 && self.rows.get_row(te.row)[0..self.n_keys] == row[0..self.n_keys]
1508 },
1509 TableEntry::hashcode,
1510 );
1511 match entry {
1512 Entry::Occupied(mut occupied_entry) => {
1513 let cur = self.rows.get_row(occupied_entry.get().row);
1514 if merge_fn(cur, row, &mut self.scratch) {
1515 let new = self.rows.add_row(&self.scratch);
1516 self.rows.set_stale(occupied_entry.get().row);
1517 self.n_stale += 1;
1518 occupied_entry.get_mut().row = new;
1519 }
1520 self.scratch.clear();
1521 }
1522 Entry::Vacant(vacant_entry) => {
1523 let next = self.rows.add_row(row);
1524 vacant_entry.insert(TableEntry {
1525 hashcode: hc as _,
1526 row: next,
1527 });
1528 }
1529 }
1530 }
1531
1532 fn write_output(&self, output: &ParallelRowBufWriter) -> (RowId, usize) {
1535 (output.append_contents(&self.rows), self.n_stale)
1536 }
1537}