1#[cfg(feature = "parallel")]
39use rayon::prelude::*;
40use rustc_hash::FxHashSet;
41use std::collections::VecDeque;
42use std::sync::atomic::AtomicUsize;
43use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
44
45use radixdb_core::value::NULL_VALUE;
46use radixdb_core::{CompactArc, CompactVec};
47use radixdb_core::{Result, Row, RowVec, Value};
48use radixdb_functions::FunctionRegistry;
49use radixdb_sql::ast::Expression;
50
51use super::context::{CancellationHandle, ExecutionContext};
52use super::expression::ExpressionEval;
53#[cfg(feature = "parallel")]
54use super::expression::RowFilter;
55use super::hash_table::{hash_keys_with, JoinMemoryReservation};
56use super::operator::ColumnSource;
57use super::operator::{ColumnInfo, Operator, OrderingProperty, RowRef};
58use super::utils::{hash_composite_key, verify_composite_key_equality, RetainedRowsBudget};
59
60#[cfg(feature = "parallel")]
62#[inline]
63fn num_threads() -> usize {
64 rayon::current_num_threads()
65}
66
67pub use super::operators::hash_join::JoinType;
69
70#[doc(hidden)]
71pub static PARALLEL_JOIN_CANCELLATION_OBSERVED: AtomicUsize = AtomicUsize::new(0);
72
73#[inline]
74fn check_parallel_cancellation(cancellation: Option<&CancellationHandle>) -> Result<()> {
75 if cancellation.is_some_and(CancellationHandle::is_cancelled) {
76 PARALLEL_JOIN_CANCELLATION_OBSERVED.fetch_add(1, Ordering::Relaxed);
77 Err(radixdb_core::Error::QueryCancelled)
78 } else {
79 Ok(())
80 }
81}
82
83pub const DEFAULT_PARALLEL_FILTER_THRESHOLD: usize = 10_000;
86pub const DEFAULT_PARALLEL_JOIN_THRESHOLD: usize = 10_000;
87pub const DEFAULT_PARALLEL_CHUNK_SIZE: usize = 2048;
88pub const DEFAULT_PARALLEL_JOIN_OUTPUT_BATCH_ROWS: usize = 2048;
89pub const DEFAULT_PARALLEL_JOIN_OUTPUT_BATCH_BYTES: usize = 16 * 1024 * 1024;
90
91#[derive(Clone, Debug)]
93pub struct ParallelConfig {
94 pub enabled: bool,
96 pub min_rows_for_parallel_filter: usize,
98 pub min_rows_for_parallel_join: usize,
100 pub chunk_size: usize,
102 pub join_output_batch_rows: usize,
104 pub join_output_batch_bytes: usize,
106}
107
108impl Default for ParallelConfig {
109 fn default() -> Self {
110 Self {
111 enabled: true,
112 min_rows_for_parallel_filter: DEFAULT_PARALLEL_FILTER_THRESHOLD,
113 min_rows_for_parallel_join: DEFAULT_PARALLEL_JOIN_THRESHOLD,
114 chunk_size: DEFAULT_PARALLEL_CHUNK_SIZE,
119 join_output_batch_rows: DEFAULT_PARALLEL_JOIN_OUTPUT_BATCH_ROWS,
120 join_output_batch_bytes: DEFAULT_PARALLEL_JOIN_OUTPUT_BATCH_BYTES,
121 }
122 }
123}
124
125impl ParallelConfig {
126 #[inline]
128 pub fn should_parallel_filter(&self, row_count: usize) -> bool {
129 #[cfg(any(test, feature = "test-failpoints"))]
130 {
131 if radixdb_storage::test_failpoints::force_serial_execution() {
132 return false;
133 }
134 if radixdb_storage::test_failpoints::force_parallel_execution() {
135 return self.enabled && row_count > 0;
136 }
137 }
138 self.enabled && row_count >= self.min_rows_for_parallel_filter
139 }
140
141 #[inline]
143 pub fn should_parallel_join(&self, build_rows: usize) -> bool {
144 self.enabled && build_rows >= self.min_rows_for_parallel_join
145 }
146}
147
148pub fn parallel_filter(
168 rows: RowVec,
169 filter_expr: &Expression,
170 columns: &[String],
171 _function_registry: &FunctionRegistry,
172 config: &ParallelConfig,
173 ctx: &ExecutionContext,
174) -> Result<RowVec> {
175 #[cfg(feature = "parallel")]
176 {
177 let row_count = rows.len();
178 if config.should_parallel_filter(row_count) {
179 #[cfg(any(test, feature = "test-failpoints"))]
180 radixdb_storage::test_failpoints::record_execution_path(7);
181 let columns_vec: Vec<String> = columns.to_vec();
183 let filter = RowFilter::new(filter_expr, &columns_vec)?.with_context(ctx);
184
185 let mut row_vec: Vec<(i64, Row)> = rows.into_vec();
189 let keep: Result<Vec<bool>> = row_vec
190 .par_iter()
191 .map(|(_, row)| filter.matches_checked(row))
192 .collect();
193 let keep = keep?;
194
195 let kept_count = keep.iter().filter(|&&b| b).count();
197 let mut filtered = Vec::with_capacity(kept_count);
198 for (i, entry) in row_vec.drain(..).enumerate() {
199 if keep[i] {
200 filtered.push(entry);
201 }
202 }
203
204 return Ok(RowVec::from_vec(filtered));
206 }
207 }
208
209 #[cfg(any(test, feature = "test-failpoints"))]
211 radixdb_storage::test_failpoints::record_execution_path(6);
212 let _ = config; sequential_filter(rows, filter_expr, columns, ctx)
214}
215
216fn sequential_filter(
218 rows: RowVec,
219 filter_expr: &Expression,
220 columns: &[String],
221 ctx: &ExecutionContext,
222) -> Result<RowVec> {
223 let columns_vec: Vec<String> = columns.to_vec();
224 let mut eval = ExpressionEval::compile(filter_expr, &columns_vec)?.with_context(ctx);
225
226 let mut result = RowVec::with_capacity(rows.len());
227 for (id, row) in rows {
228 if eval.eval_bool_checked(&row)? {
229 result.push((id, row));
230 }
231 }
232 Ok(result)
233}
234
235const EMPTY_PARALLEL_BUCKET: u32 = u32::MAX;
240
241#[repr(C)]
244struct ParallelHashEntry {
245 hash: AtomicU64,
246 next: AtomicU32,
247}
248
249impl ParallelHashEntry {
250 fn empty() -> Self {
251 Self {
252 hash: AtomicU64::new(0),
253 next: AtomicU32::new(EMPTY_PARALLEL_BUCKET),
254 }
255 }
256}
257
258struct BuildMatchedTracker {
264 matched: Vec<AtomicBool>,
265}
266
267impl BuildMatchedTracker {
268 fn new(size: usize) -> Self {
270 BuildMatchedTracker {
271 matched: (0..size).map(|_| AtomicBool::new(false)).collect(),
272 }
273 }
274
275 #[inline]
281 fn mark_matched(&self, idx: usize) {
282 self.matched[idx].store(true, Ordering::Release);
283 }
284
285 #[inline]
289 fn was_matched(&self, idx: usize) -> bool {
290 self.matched[idx].load(Ordering::Acquire)
291 }
292}
293
294struct ParallelHashTable {
296 bucket_heads: Vec<AtomicU32>,
297 entries: Vec<ParallelHashEntry>,
298 bucket_mask: u64,
299}
300
301impl ParallelHashTable {
302 fn retained_bytes(row_count: usize) -> Option<usize> {
303 crate::hash_table::JoinHashTable::estimated_retained_bytes(row_count)
304 }
305
306 fn with_capacity(row_count: usize, max_bytes: usize) -> Result<Self> {
307 let retained_bytes = Self::retained_bytes(row_count).ok_or_else(|| {
308 radixdb_core::Error::invalid_argument(
309 "parallel join hash state exceeds its addressable row range",
310 )
311 })?;
312 if retained_bytes > max_bytes {
313 return Err(radixdb_core::Error::invalid_argument(format!(
314 "parallel join hash state exceeds memory budget ({retained_bytes}/{max_bytes} bytes)"
315 )));
316 }
317
318 let entry_bytes = row_count
319 .checked_mul(std::mem::size_of::<ParallelHashEntry>())
320 .ok_or_else(|| {
321 radixdb_core::Error::invalid_argument(
322 "parallel join hash entry size exceeds addressable memory",
323 )
324 })?;
325 let bucket_bytes = retained_bytes.checked_sub(entry_bytes).ok_or_else(|| {
326 radixdb_core::Error::internal("parallel join hash retained-size contract mismatch")
327 })?;
328 let bucket_count = bucket_bytes / std::mem::size_of::<AtomicU32>();
329 let bucket_heads = (0..bucket_count)
330 .map(|_| AtomicU32::new(EMPTY_PARALLEL_BUCKET))
331 .collect();
332 let entries = (0..row_count).map(|_| ParallelHashEntry::empty()).collect();
333
334 Ok(Self {
335 bucket_heads,
336 entries,
337 bucket_mask: (bucket_count - 1) as u64,
338 })
339 }
340
341 #[inline]
342 fn insert(&self, hash: u64, row_idx: usize) {
343 let row_idx = row_idx as u32;
344 let bucket = (hash & self.bucket_mask) as usize;
345 let previous = self.bucket_heads[bucket].swap(row_idx, Ordering::Relaxed);
346 let entry = &self.entries[row_idx as usize];
347 entry.hash.store(hash, Ordering::Relaxed);
348 entry.next.store(previous, Ordering::Relaxed);
349 }
350
351 #[inline]
352 fn for_each_match(&self, key: &u64, mut visit: impl FnMut(usize)) {
353 let mut entry_idx =
354 self.bucket_heads[(*key & self.bucket_mask) as usize].load(Ordering::Relaxed);
355 while entry_idx != EMPTY_PARALLEL_BUCKET {
356 let entry = &self.entries[entry_idx as usize];
357 if entry.hash.load(Ordering::Relaxed) == *key {
358 visit(entry_idx as usize);
359 }
360 entry_idx = entry.next.load(Ordering::Relaxed);
361 }
362 }
363
364 #[inline]
365 fn bucket_head(&self, hash: u64) -> u32 {
366 self.bucket_heads[(hash & self.bucket_mask) as usize].load(Ordering::Relaxed)
367 }
368}
369
370pub fn parallel_join_state_retained_bytes(
371 build_rows: usize,
372 track_build_matches: bool,
373) -> Option<usize> {
374 let hash_bytes = ParallelHashTable::retained_bytes(build_rows)?;
375 let tracker_bytes = if track_build_matches {
376 build_rows.checked_mul(std::mem::size_of::<AtomicBool>())?
377 } else {
378 0
379 };
380 hash_bytes.checked_add(tracker_bytes)
381}
382
383#[derive(Debug)]
384struct ParallelProbeState {
385 row: Option<RowRef>,
386 hash: u64,
387 next_entry: u32,
388 matched: bool,
389 done: bool,
390}
391
392#[derive(Debug, Clone, Copy)]
393enum ParallelProbeCandidate {
394 Match {
395 probe_index: usize,
396 build_index: usize,
397 },
398 UnmatchedProbe {
399 probe_index: usize,
400 },
401}
402
403#[derive(Debug)]
404struct ParallelProbeRound {
405 candidate: Option<ParallelProbeCandidate>,
406 examined: u64,
407}
408
409pub struct ParallelHashJoinOperator {
417 probe: Box<dyn Operator>,
418 build_rows: CompactArc<Vec<Row>>,
419 build_key_indices: Vec<usize>,
420 probe_key_indices: Vec<usize>,
421 join_type: JoinType,
422 build_is_left: bool,
423 left_col_count: usize,
424 right_col_count: usize,
425 schema: Vec<ColumnInfo>,
426 projection_columns: CompactArc<[ColumnSource]>,
427 config: ParallelConfig,
428 cancellation: CancellationHandle,
429 hash_table: Option<ParallelHashTable>,
430 build_matched: Option<BuildMatchedTracker>,
431 probe_batch: Vec<ParallelProbeState>,
432 ready: VecDeque<ParallelProbeCandidate>,
433 probe_batch_budget: RetainedRowsBudget,
434 probe_exhausted: bool,
435 unmatched_build_index: usize,
436 hash_state_bytes: usize,
437 _hash_state_reservation: Option<JoinMemoryReservation>,
438 _batch_reservation: Option<JoinMemoryReservation>,
439 opened: bool,
440 closed: bool,
441 observed_probe_rows: u64,
442 observed_candidate_rows: u64,
443 #[cfg(test)]
444 peak_probe_batch_rows: usize,
445}
446
447#[allow(clippy::too_many_arguments)]
448impl ParallelHashJoinOperator {
449 pub fn new(
450 probe: Box<dyn Operator>,
451 build_rows: CompactArc<Vec<Row>>,
452 build_key_indices: Vec<usize>,
453 probe_key_indices: Vec<usize>,
454 join_type: JoinType,
455 build_is_left: bool,
456 left_col_count: usize,
457 right_col_count: usize,
458 columns: Vec<String>,
459 projection: Option<Vec<ColumnSource>>,
460 config: ParallelConfig,
461 cancellation: CancellationHandle,
462 hash_state_bytes: usize,
463 hash_state_reservation: JoinMemoryReservation,
464 batch_reservation: JoinMemoryReservation,
465 ) -> Result<Self> {
466 let projection = projection.unwrap_or_else(|| {
467 let mut sources = Vec::with_capacity(left_col_count + right_col_count);
468 sources.extend((0..left_col_count).map(ColumnSource::Outer));
469 sources.extend((0..right_col_count).map(ColumnSource::Inner));
470 sources
471 });
472 super::operator::JoinProjection {
473 columns: projection.clone(),
474 }
475 .validate(left_col_count, right_col_count, columns.len())?;
476
477 let batch_rows = config.join_output_batch_rows.max(1);
478 let batch_bytes = config.join_output_batch_bytes.max(1);
479 Ok(Self {
480 probe,
481 build_rows,
482 build_key_indices,
483 probe_key_indices,
484 join_type,
485 build_is_left,
486 left_col_count,
487 right_col_count,
488 schema: columns.into_iter().map(ColumnInfo::new).collect(),
489 projection_columns: CompactArc::from(projection),
490 config,
491 cancellation,
492 hash_table: None,
493 build_matched: None,
494 probe_batch: Vec::with_capacity(batch_rows),
495 ready: VecDeque::with_capacity(batch_rows),
496 probe_batch_budget: RetainedRowsBudget::with_limits(
497 "parallel JOIN probe batch",
498 batch_rows,
499 batch_bytes,
500 ),
501 probe_exhausted: false,
502 unmatched_build_index: 0,
503 hash_state_bytes,
504 _hash_state_reservation: Some(hash_state_reservation),
505 _batch_reservation: Some(batch_reservation),
506 opened: false,
507 closed: false,
508 observed_probe_rows: 0,
509 observed_candidate_rows: 0,
510 #[cfg(test)]
511 peak_probe_batch_rows: 0,
512 })
513 }
514
515 pub fn observed_probe_rows(&self) -> u64 {
516 self.observed_probe_rows
517 }
518
519 pub fn observed_candidate_rows(&self) -> u64 {
520 self.observed_candidate_rows
521 }
522
523 fn fill_probe_batch(&mut self) -> Result<()> {
524 debug_assert!(self.probe_batch.is_empty());
525 while self.probe_batch.len() < self.config.join_output_batch_rows.max(1) {
526 check_parallel_cancellation(Some(&self.cancellation))?;
527 let Some(row) = self.probe.next()? else {
528 self.probe_exhausted = true;
529 break;
530 };
531 self.probe_batch_budget
532 .admit_estimated_bytes(row.estimated_retained_bytes())?;
533 let hash = hash_keys_with(&self.probe_key_indices, |index| row.get(index));
534 let next_entry = self
535 .hash_table
536 .as_ref()
537 .expect("opened parallel JOIN must own a hash table")
538 .bucket_head(hash);
539 self.probe_batch.push(ParallelProbeState {
540 row: Some(row),
541 hash,
542 next_entry,
543 matched: false,
544 done: false,
545 });
546 #[cfg(test)]
547 {
548 self.peak_probe_batch_rows = self.peak_probe_batch_rows.max(self.probe_batch.len());
549 }
550 self.observed_probe_rows = self.observed_probe_rows.saturating_add(1);
551 }
552 Ok(())
553 }
554
555 fn advance_probe_state(
556 state: &mut ParallelProbeState,
557 probe_index: usize,
558 hash_table: &ParallelHashTable,
559 build_rows: &[Row],
560 probe_key_indices: &[usize],
561 build_key_indices: &[usize],
562 build_matched: Option<&BuildMatchedTracker>,
563 needs_unmatched_probe: bool,
564 cancellation: &CancellationHandle,
565 ) -> ParallelProbeRound {
566 if state.done || cancellation.is_cancelled() {
567 return ParallelProbeRound {
568 candidate: None,
569 examined: 0,
570 };
571 }
572
573 let probe_row = state
574 .row
575 .as_ref()
576 .expect("live parallel probe state must retain its row");
577 let mut examined = 0_u64;
578 while state.next_entry != EMPTY_PARALLEL_BUCKET {
579 let build_index = state.next_entry as usize;
580 let entry = &hash_table.entries[build_index];
581 state.next_entry = entry.next.load(Ordering::Relaxed);
582 examined = examined.saturating_add(1);
583 if entry.hash.load(Ordering::Relaxed) != state.hash {
584 continue;
585 }
586 let build_row = &build_rows[build_index];
587 let matches = probe_key_indices.iter().zip(build_key_indices.iter()).all(
588 |(&probe_index, &build_index)| {
589 let (Some(probe_value), Some(build_value)) =
590 (probe_row.get(probe_index), build_row.get(build_index))
591 else {
592 return false;
593 };
594 !probe_value.is_null() && !build_value.is_null() && probe_value == build_value
595 },
596 );
597 if matches {
598 state.matched = true;
599 if let Some(tracker) = build_matched {
600 tracker.mark_matched(build_index);
601 }
602 if state.next_entry == EMPTY_PARALLEL_BUCKET {
603 state.done = true;
604 }
605 return ParallelProbeRound {
606 candidate: Some(ParallelProbeCandidate::Match {
607 probe_index,
608 build_index,
609 }),
610 examined,
611 };
612 }
613 }
614
615 state.done = true;
616 ParallelProbeRound {
617 candidate: (!state.matched && needs_unmatched_probe)
618 .then_some(ParallelProbeCandidate::UnmatchedProbe { probe_index }),
619 examined,
620 }
621 }
622
623 fn run_probe_round(&mut self) -> Result<()> {
624 check_parallel_cancellation(Some(&self.cancellation))?;
625 let hash_table = self
626 .hash_table
627 .as_ref()
628 .expect("opened parallel JOIN must own a hash table");
629 let build_rows = self.build_rows.as_slice();
630 let probe_key_indices = self.probe_key_indices.as_slice();
631 let build_key_indices = self.build_key_indices.as_slice();
632 let build_matched = self.build_matched.as_ref();
633 let needs_unmatched_probe = self.join_type.needs_unmatched_probe(self.build_is_left);
634 let cancellation = &self.cancellation;
635
636 #[cfg(feature = "parallel")]
637 let rounds: Vec<ParallelProbeRound> = self
638 .probe_batch
639 .par_iter_mut()
640 .enumerate()
641 .map(|(probe_index, state)| {
642 Self::advance_probe_state(
643 state,
644 probe_index,
645 hash_table,
646 build_rows,
647 probe_key_indices,
648 build_key_indices,
649 build_matched,
650 needs_unmatched_probe,
651 cancellation,
652 )
653 })
654 .collect();
655 #[cfg(not(feature = "parallel"))]
656 let rounds: Vec<ParallelProbeRound> = self
657 .probe_batch
658 .iter_mut()
659 .enumerate()
660 .map(|(probe_index, state)| {
661 Self::advance_probe_state(
662 state,
663 probe_index,
664 hash_table,
665 build_rows,
666 probe_key_indices,
667 build_key_indices,
668 build_matched,
669 needs_unmatched_probe,
670 cancellation,
671 )
672 })
673 .collect();
674
675 check_parallel_cancellation(Some(&self.cancellation))?;
676 for round in rounds {
677 self.observed_candidate_rows =
678 self.observed_candidate_rows.saturating_add(round.examined);
679 if let Some(candidate) = round.candidate {
680 self.ready.push_back(candidate);
681 }
682 }
683 Ok(())
684 }
685
686 fn projected_row(&mut self, candidate: ParallelProbeCandidate) -> RowRef {
687 match candidate {
688 ParallelProbeCandidate::Match {
689 probe_index,
690 build_index,
691 } => {
692 let state = &mut self.probe_batch[probe_index];
693 let probe = if state.done {
694 state
695 .row
696 .take()
697 .expect("completed parallel probe must retain its result row")
698 } else {
699 state
700 .row
701 .as_ref()
702 .expect("live parallel probe must retain its row")
703 .clone()
704 };
705 let build = RowRef::shared(CompactArc::clone(&self.build_rows), build_index);
706 let (left, right) = if self.build_is_left {
707 (build, probe)
708 } else {
709 (probe, build)
710 };
711 RowRef::projected(left, right, CompactArc::clone(&self.projection_columns))
712 }
713 ParallelProbeCandidate::UnmatchedProbe { probe_index } => {
714 let probe = self.probe_batch[probe_index]
715 .row
716 .take()
717 .expect("unmatched parallel probe must retain its row");
718 let build_width = if self.build_is_left {
719 self.left_col_count
720 } else {
721 self.right_col_count
722 };
723 let null_build = RowRef::owned(Row::from_values(vec![NULL_VALUE; build_width]));
724 let (left, right) = if self.build_is_left {
725 (null_build, probe)
726 } else {
727 (probe, null_build)
728 };
729 RowRef::projected(left, right, CompactArc::clone(&self.projection_columns))
730 }
731 }
732 }
733
734 fn next_unmatched_build(&mut self) -> Result<Option<RowRef>> {
735 let Some(tracker) = self.build_matched.as_ref() else {
736 return Ok(None);
737 };
738 while self.unmatched_build_index < self.build_rows.len() {
739 if self.unmatched_build_index & 0xff == 0 {
740 check_parallel_cancellation(Some(&self.cancellation))?;
741 }
742 let build_index = self.unmatched_build_index;
743 self.unmatched_build_index += 1;
744 if tracker.was_matched(build_index) {
745 continue;
746 }
747 let build = RowRef::shared(CompactArc::clone(&self.build_rows), build_index);
748 let probe_width = if self.build_is_left {
749 self.right_col_count
750 } else {
751 self.left_col_count
752 };
753 let null_probe = RowRef::owned(Row::from_values(vec![NULL_VALUE; probe_width]));
754 let (left, right) = if self.build_is_left {
755 (build, null_probe)
756 } else {
757 (null_probe, build)
758 };
759 return Ok(Some(RowRef::projected(
760 left,
761 right,
762 CompactArc::clone(&self.projection_columns),
763 )));
764 }
765 Ok(None)
766 }
767
768 fn clear_probe_batch(&mut self) {
769 self.probe_batch.clear();
770 self.ready.clear();
771 self.probe_batch_budget = RetainedRowsBudget::with_limits(
772 "parallel JOIN probe batch",
773 self.config.join_output_batch_rows.max(1),
774 self.config.join_output_batch_bytes.max(1),
775 );
776 }
777}
778
779impl Operator for ParallelHashJoinOperator {
780 fn open(&mut self) -> Result<()> {
781 if self.opened {
782 return Ok(());
783 }
784 check_parallel_cancellation(Some(&self.cancellation))?;
785 let tracker_bytes = if self.join_type.needs_unmatched_build(self.build_is_left) {
786 self.build_rows
787 .len()
788 .checked_mul(std::mem::size_of::<AtomicBool>())
789 .ok_or_else(|| {
790 radixdb_core::Error::invalid_argument(
791 "parallel join match tracker exceeds addressable memory",
792 )
793 })?
794 } else {
795 0
796 };
797 self.hash_table = Some(parallel_hash_build_inner(
798 &self.build_rows,
799 &self.build_key_indices,
800 &self.config,
801 Some(&self.cancellation),
802 self.hash_state_bytes.saturating_sub(tracker_bytes),
803 )?);
804 if tracker_bytes > 0 {
805 self.build_matched = Some(BuildMatchedTracker::new(self.build_rows.len()));
806 }
807 if let Err(error) = self.probe.open() {
808 self.hash_table = None;
809 self.build_matched = None;
810 return Err(error);
811 }
812 self.opened = true;
813 Ok(())
814 }
815
816 fn next(&mut self) -> Result<Option<RowRef>> {
817 if !self.opened || self.closed {
818 return Ok(None);
819 }
820 loop {
821 check_parallel_cancellation(Some(&self.cancellation))?;
822 if let Some(candidate) = self.ready.pop_front() {
823 return Ok(Some(self.projected_row(candidate)));
824 }
825
826 if !self.probe_batch.is_empty() {
827 if self.probe_batch.iter().all(|state| state.done) {
828 self.clear_probe_batch();
829 } else {
830 self.run_probe_round()?;
831 continue;
832 }
833 }
834
835 if !self.probe_exhausted {
836 self.fill_probe_batch()?;
837 if !self.probe_batch.is_empty() {
838 continue;
839 }
840 }
841
842 return self.next_unmatched_build();
843 }
844 }
845
846 fn close(&mut self) -> Result<()> {
847 if self.closed {
848 return Ok(());
849 }
850 let _ = check_parallel_cancellation(Some(&self.cancellation));
854 self.closed = true;
855 self.clear_probe_batch();
856 self.hash_table = None;
857 self.build_matched = None;
858 self._hash_state_reservation = None;
859 self._batch_reservation = None;
860 self.probe.close()
861 }
862
863 fn schema(&self) -> &[ColumnInfo] {
864 &self.schema
865 }
866
867 fn estimated_rows(&self) -> Option<usize> {
868 None
869 }
870
871 fn ordering(&self) -> OrderingProperty {
872 OrderingProperty::Unknown
873 }
874
875 fn name(&self) -> &str {
876 "ParallelHashJoin"
877 }
878}
879
880#[cfg(test)]
881pub fn parallel_join_state_fits_budget(
882 build_rows: usize,
883 track_build_matches: bool,
884 max_bytes: usize,
885) -> bool {
886 parallel_join_state_retained_bytes(build_rows, track_build_matches)
887 .is_some_and(|bytes| bytes <= max_bytes)
888}
889
890fn parallel_hash_build_inner(
891 build_rows: &[Row],
892 key_indices: &[usize],
893 config: &ParallelConfig,
894 cancellation: Option<&CancellationHandle>,
895 max_bytes: usize,
896) -> Result<ParallelHashTable> {
897 check_parallel_cancellation(cancellation)?;
898 let row_count = build_rows.len();
899 let table = ParallelHashTable::with_capacity(row_count, max_bytes)?;
900
901 #[cfg(feature = "parallel")]
902 if config.should_parallel_join(row_count) {
903 let n_threads = num_threads();
904 let chunk_size = config.chunk_size.max(row_count / n_threads).max(1000);
905
906 build_rows
909 .par_chunks(chunk_size)
910 .enumerate()
911 .for_each(|(chunk_idx, chunk)| {
912 if cancellation.is_some_and(CancellationHandle::is_cancelled) {
913 return;
914 }
915 let base_idx = chunk_idx * chunk_size;
916 for (local_idx, row) in chunk.iter().enumerate() {
917 if local_idx & 0xff == 0
918 && cancellation.is_some_and(CancellationHandle::is_cancelled)
919 {
920 return;
921 }
922 debug_assert!(
925 base_idx.checked_add(local_idx).is_some(),
926 "Index overflow in parallel hash build: base_idx={} + local_idx={}",
927 base_idx,
928 local_idx
929 );
930 let global_idx = base_idx + local_idx;
931 let hash = hash_composite_key(row, key_indices);
932 table.insert(hash, global_idx);
933 }
934 });
935
936 check_parallel_cancellation(cancellation)?;
937 return Ok(table);
938 }
939
940 let _ = config; for (idx, row) in build_rows.iter().enumerate() {
944 if idx & 0xff == 0 {
945 check_parallel_cancellation(cancellation)?;
946 }
947 let hash = hash_composite_key(row, key_indices);
948 table.insert(hash, idx);
949 }
950 Ok(table)
951}
952
953#[inline]
955fn hash_row_by_keys(row: &Row, key_indices: &[usize]) -> u64 {
956 hash_composite_key(row, key_indices)
957}
958
959#[inline]
961fn verify_key_match(
962 probe_row: &Row,
963 build_row: &Row,
964 probe_key_indices: &[usize],
965 build_key_indices: &[usize],
966) -> bool {
967 verify_composite_key_equality(probe_row, build_row, probe_key_indices, build_key_indices)
968}
969
970pub struct ParallelJoinResult {
974 pub rows: Vec<Row>,
976}
977
978#[allow(clippy::too_many_arguments)]
980fn sequential_probe(
981 probe_rows: &[Row],
982 build_rows: &[Row],
983 hash_table: &ParallelHashTable,
984 probe_key_indices: &[usize],
985 build_key_indices: &[usize],
986 join_type: &JoinType,
987 probe_col_count: usize,
988 build_col_count: usize,
989 swapped: bool,
990 projection: Option<&[ColumnSource]>,
991 build_matched: &Option<BuildMatchedTracker>,
992 cancellation: Option<&CancellationHandle>,
993) -> Result<(Vec<Row>, Vec<Row>)> {
994 let mut matched_rows = Vec::new();
995 let needs_unmatched_probe = join_type.needs_unmatched_probe(swapped);
996
997 for (probe_index, probe_row) in probe_rows.iter().enumerate() {
998 if probe_index & 0xff == 0 {
999 check_parallel_cancellation(cancellation)?;
1000 }
1001 let hash = hash_row_by_keys(probe_row, probe_key_indices);
1002 let mut matched = false;
1003
1004 hash_table.for_each_match(&hash, |build_idx| {
1005 let build_row = &build_rows[build_idx];
1006 if verify_key_match(probe_row, build_row, probe_key_indices, build_key_indices) {
1007 matched = true;
1008 if let Some(ref tracker) = build_matched {
1009 tracker.mark_matched(build_idx);
1010 }
1011 let combined = combine_join_rows(
1012 probe_row,
1013 build_row,
1014 probe_col_count,
1015 build_col_count,
1016 swapped,
1017 projection,
1018 );
1019 matched_rows.push(Row::from_compact_vec(combined));
1020 }
1021 });
1022
1023 if !matched && needs_unmatched_probe {
1024 let values = combine_with_nulls(
1025 probe_row,
1026 probe_col_count,
1027 build_col_count,
1028 swapped,
1029 projection,
1030 );
1031 matched_rows.push(Row::from_compact_vec(values));
1032 }
1033 }
1034
1035 Ok((matched_rows, Vec::new()))
1036}
1037
1038#[allow(clippy::too_many_arguments)]
1040#[cfg(test)]
1041fn parallel_hash_join(
1042 probe_rows: &[Row],
1043 build_rows: &[Row],
1044 probe_key_indices: &[usize],
1045 build_key_indices: &[usize],
1046 join_type: JoinType,
1047 probe_col_count: usize,
1048 build_col_count: usize,
1049 swapped: bool,
1050 config: &ParallelConfig,
1051) -> ParallelJoinResult {
1052 parallel_hash_join_inner(
1053 probe_rows,
1054 build_rows,
1055 probe_key_indices,
1056 build_key_indices,
1057 join_type,
1058 probe_col_count,
1059 build_col_count,
1060 swapped,
1061 None,
1062 config,
1063 None,
1064 usize::MAX,
1065 )
1066 .expect("uncancellable parallel hash join")
1067}
1068
1069#[allow(clippy::too_many_arguments)]
1070pub fn parallel_hash_join_cancellable(
1071 probe_rows: &[Row],
1072 build_rows: &[Row],
1073 probe_key_indices: &[usize],
1074 build_key_indices: &[usize],
1075 join_type: JoinType,
1076 probe_col_count: usize,
1077 build_col_count: usize,
1078 swapped: bool,
1079 projection: Option<&[ColumnSource]>,
1080 config: &ParallelConfig,
1081 cancellation: &CancellationHandle,
1082 max_state_bytes: usize,
1083) -> Result<ParallelJoinResult> {
1084 parallel_hash_join_inner(
1085 probe_rows,
1086 build_rows,
1087 probe_key_indices,
1088 build_key_indices,
1089 join_type,
1090 probe_col_count,
1091 build_col_count,
1092 swapped,
1093 projection,
1094 config,
1095 Some(cancellation),
1096 max_state_bytes,
1097 )
1098}
1099
1100#[allow(clippy::too_many_arguments)]
1101fn parallel_hash_join_inner(
1102 probe_rows: &[Row],
1103 build_rows: &[Row],
1104 probe_key_indices: &[usize],
1105 build_key_indices: &[usize],
1106 join_type: JoinType,
1107 probe_col_count: usize,
1108 build_col_count: usize,
1109 swapped: bool,
1110 projection: Option<&[ColumnSource]>,
1111 config: &ParallelConfig,
1112 cancellation: Option<&CancellationHandle>,
1113 max_state_bytes: usize,
1114) -> Result<ParallelJoinResult> {
1115 check_parallel_cancellation(cancellation)?;
1116 #[cfg(feature = "parallel")]
1117 let probe_count = probe_rows.len();
1118 let build_count = build_rows.len();
1119
1120 #[cfg(feature = "parallel")]
1122 let use_parallel =
1123 config.should_parallel_join(build_count) || config.should_parallel_join(probe_count);
1124
1125 let track_build_matches = join_type.needs_unmatched_build(swapped);
1128 let retained_bytes = parallel_join_state_retained_bytes(build_count, track_build_matches)
1129 .ok_or_else(|| {
1130 radixdb_core::Error::invalid_argument(
1131 "parallel join state exceeds its addressable memory range",
1132 )
1133 })?;
1134 if retained_bytes > max_state_bytes {
1135 return Err(radixdb_core::Error::invalid_argument(format!(
1136 "parallel join state exceeds memory budget ({retained_bytes}/{max_state_bytes} bytes)"
1137 )));
1138 }
1139 let tracker_bytes = if track_build_matches {
1140 build_count
1141 .checked_mul(std::mem::size_of::<AtomicBool>())
1142 .ok_or_else(|| {
1143 radixdb_core::Error::invalid_argument(
1144 "parallel join match tracker exceeds addressable memory",
1145 )
1146 })?
1147 } else {
1148 0
1149 };
1150
1151 let hash_table = parallel_hash_build_inner(
1154 build_rows,
1155 build_key_indices,
1156 config,
1157 cancellation,
1158 max_state_bytes.saturating_sub(tracker_bytes),
1159 )?;
1160
1161 let build_matched: Option<BuildMatchedTracker> = if track_build_matches {
1162 Some(BuildMatchedTracker::new(build_count))
1163 } else {
1164 None
1165 };
1166
1167 #[allow(unused_variables)]
1169 let (matched_rows, unmatched_probe_rows) = {
1170 #[cfg(feature = "parallel")]
1171 {
1172 if use_parallel && join_type == JoinType::Inner {
1173 let matches: Vec<Row> = probe_rows
1175 .par_chunks(config.chunk_size.max(1000))
1176 .flat_map(|chunk| {
1177 let mut local_results = Vec::new();
1178 for (probe_index, probe_row) in chunk.iter().enumerate() {
1179 if probe_index & 0xff == 0
1180 && cancellation.is_some_and(CancellationHandle::is_cancelled)
1181 {
1182 break;
1183 }
1184 let hash = hash_row_by_keys(probe_row, probe_key_indices);
1185 hash_table.for_each_match(&hash, |build_idx| {
1186 let build_row = &build_rows[build_idx];
1187 if verify_key_match(
1188 probe_row,
1189 build_row,
1190 probe_key_indices,
1191 build_key_indices,
1192 ) {
1193 let combined = combine_join_rows(
1194 probe_row,
1195 build_row,
1196 probe_col_count,
1197 build_col_count,
1198 swapped,
1199 projection,
1200 );
1201 local_results.push(Row::from_compact_vec(combined));
1202 }
1203 });
1204 }
1205 local_results
1206 })
1207 .collect();
1208 check_parallel_cancellation(cancellation)?;
1209 (matches, Vec::new())
1210 } else if use_parallel {
1211 let needs_unmatched_probe = join_type.needs_unmatched_probe(swapped);
1214
1215 let chunk_results: Vec<(Vec<Row>, Vec<Row>)> = probe_rows
1217 .par_chunks(config.chunk_size.max(1000))
1218 .map(|chunk| {
1219 let mut matched_results = Vec::new();
1220 let mut unmatched_results = Vec::new();
1221
1222 for (probe_index, probe_row) in chunk.iter().enumerate() {
1223 if probe_index & 0xff == 0
1224 && cancellation.is_some_and(CancellationHandle::is_cancelled)
1225 {
1226 break;
1227 }
1228 let mut matched = false;
1229 let hash = hash_row_by_keys(probe_row, probe_key_indices);
1230
1231 hash_table.for_each_match(&hash, |build_idx| {
1232 let build_row = &build_rows[build_idx];
1233 if verify_key_match(
1234 probe_row,
1235 build_row,
1236 probe_key_indices,
1237 build_key_indices,
1238 ) {
1239 matched = true;
1240 if let Some(ref tracker) = build_matched {
1241 tracker.mark_matched(build_idx);
1242 }
1243 let combined = combine_join_rows(
1244 probe_row,
1245 build_row,
1246 probe_col_count,
1247 build_col_count,
1248 swapped,
1249 projection,
1250 );
1251 matched_results.push(Row::from_compact_vec(combined));
1252 }
1253 });
1254
1255 if !matched && needs_unmatched_probe {
1256 let values = combine_with_nulls(
1257 probe_row,
1258 probe_col_count,
1259 build_col_count,
1260 swapped,
1261 projection,
1262 );
1263 unmatched_results.push(Row::from_compact_vec(values));
1264 }
1265 }
1266
1267 (matched_results, unmatched_results)
1268 })
1269 .collect();
1270 check_parallel_cancellation(cancellation)?;
1271
1272 let total_matched: usize = chunk_results.iter().map(|(m, _)| m.len()).sum();
1273 let total_unmatched: usize = chunk_results.iter().map(|(_, u)| u.len()).sum();
1274
1275 let mut matched_rows = Vec::with_capacity(total_matched);
1276 let mut unmatched_rows = Vec::with_capacity(total_unmatched);
1277
1278 for (matched, unmatched) in chunk_results {
1279 matched_rows.extend(matched);
1280 unmatched_rows.extend(unmatched);
1281 }
1282
1283 std::sync::atomic::fence(Ordering::Acquire);
1287
1288 (matched_rows, unmatched_rows)
1289 } else {
1290 sequential_probe(
1292 probe_rows,
1293 build_rows,
1294 &hash_table,
1295 probe_key_indices,
1296 build_key_indices,
1297 &join_type,
1298 probe_col_count,
1299 build_col_count,
1300 swapped,
1301 projection,
1302 &build_matched,
1303 cancellation,
1304 )?
1305 }
1306 }
1307 #[cfg(not(feature = "parallel"))]
1308 {
1309 sequential_probe(
1310 probe_rows,
1311 build_rows,
1312 &hash_table,
1313 probe_key_indices,
1314 build_key_indices,
1315 &join_type,
1316 probe_col_count,
1317 build_col_count,
1318 swapped,
1319 projection,
1320 &build_matched,
1321 cancellation,
1322 )?
1323 }
1324 };
1325
1326 let mut result_rows = matched_rows;
1327 result_rows.extend(unmatched_probe_rows);
1328
1329 if let Some(ref tracker) = build_matched {
1332 for (build_idx, build_row) in build_rows.iter().enumerate() {
1333 if build_idx & 0xff == 0 {
1334 check_parallel_cancellation(cancellation)?;
1335 }
1336 if !tracker.was_matched(build_idx) {
1337 let values = combine_build_with_nulls(
1338 build_row,
1339 build_col_count,
1340 probe_col_count,
1341 swapped,
1342 projection,
1343 );
1344 result_rows.push(Row::from_compact_vec(values));
1345 }
1346 }
1347 }
1348
1349 Ok(ParallelJoinResult { rows: result_rows })
1350}
1351
1352#[inline]
1354fn combine_join_rows(
1355 probe_row: &Row,
1356 build_row: &Row,
1357 probe_col_count: usize,
1358 build_col_count: usize,
1359 swapped: bool,
1360 projection: Option<&[ColumnSource]>,
1361) -> CompactVec<Value> {
1362 if let Some(projection) = projection {
1363 let (left, right) = if swapped {
1364 (build_row, probe_row)
1365 } else {
1366 (probe_row, build_row)
1367 };
1368 return project_join_rows(Some(left), Some(right), projection);
1369 }
1370 let mut combined: CompactVec<Value> =
1372 CompactVec::with_capacity(probe_col_count + build_col_count);
1373 if swapped {
1374 for i in 0..build_col_count {
1376 combined.push(build_row.get(i).cloned().unwrap_or(NULL_VALUE));
1377 }
1378 for i in 0..probe_col_count {
1379 combined.push(probe_row.get(i).cloned().unwrap_or(NULL_VALUE));
1380 }
1381 } else {
1382 for i in 0..probe_col_count {
1384 combined.push(probe_row.get(i).cloned().unwrap_or(NULL_VALUE));
1385 }
1386 for i in 0..build_col_count {
1387 combined.push(build_row.get(i).cloned().unwrap_or(NULL_VALUE));
1388 }
1389 }
1390 combined
1391}
1392
1393#[inline]
1395fn combine_with_nulls(
1396 probe_row: &Row,
1397 probe_col_count: usize,
1398 build_col_count: usize,
1399 swapped: bool,
1400 projection: Option<&[ColumnSource]>,
1401) -> CompactVec<Value> {
1402 if let Some(projection) = projection {
1403 let (left, right) = if swapped {
1404 (None, Some(probe_row))
1405 } else {
1406 (Some(probe_row), None)
1407 };
1408 return project_join_rows(left, right, projection);
1409 }
1410 let mut combined: CompactVec<Value> =
1412 CompactVec::with_capacity(probe_col_count + build_col_count);
1413 if swapped {
1414 combined.extend(std::iter::repeat_n(NULL_VALUE, build_col_count));
1416 for i in 0..probe_col_count {
1417 combined.push(probe_row.get(i).cloned().unwrap_or(NULL_VALUE));
1418 }
1419 } else {
1420 for i in 0..probe_col_count {
1422 combined.push(probe_row.get(i).cloned().unwrap_or(NULL_VALUE));
1423 }
1424 combined.extend(std::iter::repeat_n(NULL_VALUE, build_col_count));
1425 }
1426 combined
1427}
1428
1429#[inline]
1431fn combine_build_with_nulls(
1432 build_row: &Row,
1433 build_col_count: usize,
1434 probe_col_count: usize,
1435 swapped: bool,
1436 projection: Option<&[ColumnSource]>,
1437) -> CompactVec<Value> {
1438 if let Some(projection) = projection {
1439 let (left, right) = if swapped {
1440 (Some(build_row), None)
1441 } else {
1442 (None, Some(build_row))
1443 };
1444 return project_join_rows(left, right, projection);
1445 }
1446 let mut combined: CompactVec<Value> =
1448 CompactVec::with_capacity(probe_col_count + build_col_count);
1449 if swapped {
1450 for i in 0..build_col_count {
1452 combined.push(build_row.get(i).cloned().unwrap_or(NULL_VALUE));
1453 }
1454 combined.extend(std::iter::repeat_n(NULL_VALUE, probe_col_count));
1455 } else {
1456 combined.extend(std::iter::repeat_n(NULL_VALUE, probe_col_count));
1458 for i in 0..build_col_count {
1459 combined.push(build_row.get(i).cloned().unwrap_or(NULL_VALUE));
1460 }
1461 }
1462 combined
1463}
1464
1465#[inline]
1466fn project_join_rows(
1467 left: Option<&Row>,
1468 right: Option<&Row>,
1469 projection: &[ColumnSource],
1470) -> CompactVec<Value> {
1471 let mut values = CompactVec::with_capacity(projection.len());
1472 for source in projection {
1473 let value = match source {
1474 ColumnSource::Outer(index) => left.and_then(|row| row.get(*index)),
1475 ColumnSource::Inner(index) => right.and_then(|row| row.get(*index)),
1476 };
1477 values.push(value.cloned().unwrap_or(NULL_VALUE));
1478 }
1479 values
1480}
1481
1482#[derive(Debug, Clone, Copy)]
1484pub enum DistanceMetric {
1485 L2,
1486 Cosine,
1487 InnerProduct,
1488}
1489
1490pub fn parallel_topn_vector_search(
1497 rows: RowVec,
1498 vector_col_idx: usize,
1499 query_bytes: &[u8],
1500 k: usize,
1501 metric: DistanceMetric,
1502 config: &ParallelConfig,
1503) -> radixdb_core::Result<Vec<(i64, Row, f64)>> {
1504 use std::collections::BinaryHeap;
1505
1506 if k == 0 || rows.is_empty() {
1507 return Ok(Vec::new());
1508 }
1509
1510 let distance_fn: fn(&[u8], &[u8]) -> radixdb_core::Result<f64> = match metric {
1511 DistanceMetric::L2 => radixdb_functions::scalar::vector::l2_distance_bytes,
1512 DistanceMetric::Cosine => radixdb_functions::scalar::vector::cosine_distance_bytes,
1513 DistanceMetric::InnerProduct => radixdb_functions::scalar::vector::ip_distance_bytes,
1514 };
1515
1516 #[inline]
1518 fn get_vector_bytes(row: &Row, col_idx: usize) -> Option<&[u8]> {
1519 match row.get(col_idx)? {
1520 Value::Extension(data)
1521 if data.first() == Some(&(radixdb_core::DataType::Vector as u8)) =>
1522 {
1523 Some(&data[1..])
1524 }
1525 _ => None,
1526 }
1527 }
1528
1529 struct HeapEntry {
1531 distance: f64,
1532 idx: usize, }
1534
1535 impl PartialEq for HeapEntry {
1536 fn eq(&self, other: &Self) -> bool {
1537 self.distance == other.distance
1538 }
1539 }
1540 impl Eq for HeapEntry {}
1541 impl PartialOrd for HeapEntry {
1542 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
1543 Some(self.cmp(other))
1544 }
1545 }
1546 impl Ord for HeapEntry {
1547 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
1548 self.distance.total_cmp(&other.distance)
1549 }
1550 }
1551
1552 let row_vec: Vec<(i64, Row)> = rows.into_vec();
1553 distance_fn(query_bytes, query_bytes)?;
1556 for (_, row) in &row_vec {
1557 if let Some(vector) = get_vector_bytes(row, vector_col_idx) {
1558 distance_fn(vector, query_bytes)?;
1559 }
1560 }
1561
1562 #[cfg(feature = "parallel")]
1563 let use_parallel = config.should_parallel_filter(row_vec.len());
1564 #[cfg(not(feature = "parallel"))]
1565 let use_parallel = false;
1566 let _ = config;
1567
1568 if use_parallel {
1569 #[cfg(feature = "parallel")]
1570 {
1571 let n_threads = num_threads();
1572 let chunk_size = (row_vec.len() / (n_threads * 4)).max(1024);
1573
1574 let chunk_results: Vec<Vec<(i64, Row, f64)>> = row_vec
1576 .into_par_iter()
1577 .chunks(chunk_size)
1578 .map(|chunk| {
1579 let mut heap: BinaryHeap<HeapEntry> = BinaryHeap::with_capacity(k + 1);
1580 let mut entries: Vec<(i64, Row, f64)> = Vec::with_capacity(k + 1);
1581
1582 for (row_id, row) in chunk {
1583 let dist = if let Some(vec_bytes) = get_vector_bytes(&row, vector_col_idx) {
1584 if vec_bytes.len() == query_bytes.len() {
1585 distance_fn(vec_bytes, query_bytes)
1586 .expect("vector shapes were validated before parallel search")
1587 } else {
1588 f64::INFINITY }
1590 } else {
1591 f64::INFINITY };
1593 if entries.len() < k {
1594 let idx = entries.len();
1595 entries.push((row_id, row, dist));
1596 heap.push(HeapEntry {
1597 distance: dist,
1598 idx,
1599 });
1600 } else if let Some(worst) = heap.peek() {
1601 if dist < worst.distance {
1602 let evict_idx = worst.idx;
1603 heap.pop();
1604 entries[evict_idx] = (row_id, row, dist);
1605 heap.push(HeapEntry {
1606 distance: dist,
1607 idx: evict_idx,
1608 });
1609 }
1610 }
1611 }
1612
1613 let live_indices: FxHashSet<usize> = heap.into_iter().map(|e| e.idx).collect();
1615 entries
1616 .into_iter()
1617 .enumerate()
1618 .filter_map(|(i, e)| {
1619 if live_indices.contains(&i) {
1620 Some(e)
1621 } else {
1622 None
1623 }
1624 })
1625 .collect()
1626 })
1627 .collect();
1628
1629 let mut merged: Vec<(i64, Row, f64)> = chunk_results.into_iter().flatten().collect();
1631 merged.sort_unstable_by(|a, b| a.2.total_cmp(&b.2));
1632 merged.truncate(k);
1633 return Ok(merged);
1634 }
1635 }
1636
1637 let mut heap: BinaryHeap<HeapEntry> = BinaryHeap::with_capacity(k + 1);
1639 let mut entries: Vec<(i64, Row, f64)> = Vec::with_capacity(k + 1);
1640
1641 for (row_id, row) in row_vec {
1642 let dist = if let Some(vec_bytes) = get_vector_bytes(&row, vector_col_idx) {
1643 if vec_bytes.len() == query_bytes.len() {
1644 distance_fn(vec_bytes, query_bytes)
1645 .expect("vector shapes were validated before sequential search")
1646 } else {
1647 f64::INFINITY }
1649 } else {
1650 f64::INFINITY };
1652 if entries.len() < k {
1653 let idx = entries.len();
1654 entries.push((row_id, row, dist));
1655 heap.push(HeapEntry {
1656 distance: dist,
1657 idx,
1658 });
1659 } else if let Some(worst) = heap.peek() {
1660 if dist < worst.distance {
1661 let evict_idx = worst.idx;
1662 heap.pop();
1663 entries[evict_idx] = (row_id, row, dist);
1664 heap.push(HeapEntry {
1665 distance: dist,
1666 idx: evict_idx,
1667 });
1668 }
1669 }
1670 }
1671
1672 let live_indices: FxHashSet<usize> = heap.into_iter().map(|e| e.idx).collect();
1673 let mut result: Vec<(i64, Row, f64)> = entries
1674 .into_iter()
1675 .enumerate()
1676 .filter_map(|(i, e)| {
1677 if live_indices.contains(&i) {
1678 Some(e)
1679 } else {
1680 None
1681 }
1682 })
1683 .collect();
1684 result.sort_unstable_by(|a, b| a.2.total_cmp(&b.2));
1685 Ok(result)
1686}
1687
1688#[cfg(test)]
1689mod tests {
1690 use super::*;
1691 use radixdb_core::Value;
1692
1693 #[test]
1694 fn test_parallel_config_thresholds() {
1695 let config = ParallelConfig::default();
1696
1697 assert!(!config.should_parallel_filter(1000)); assert!(config.should_parallel_filter(20_000)); assert!(!config.should_parallel_join(1000)); assert!(config.should_parallel_join(20_000)); }
1703
1704 #[test]
1705 fn test_parallel_hash_build() {
1706 let build_rows: Vec<Row> = (0..10_000)
1708 .map(|i| {
1709 Row::from_values(vec![
1710 Value::Integer(i),
1711 Value::Text(format!("build_{}", i).into()),
1712 ])
1713 })
1714 .collect();
1715
1716 let config = ParallelConfig {
1717 min_rows_for_parallel_join: 1000,
1718 ..Default::default()
1719 };
1720
1721 let retained = ParallelHashTable::retained_bytes(build_rows.len()).unwrap();
1722 let hash_table =
1723 parallel_hash_build_inner(&build_rows, &[0], &config, None, retained).unwrap();
1724 let test_hash = hash_row_by_keys(&build_rows[500], &[0]);
1726 let mut found = false;
1727 hash_table.for_each_match(&test_hash, |index| found |= index == 500);
1728 assert!(found);
1729 }
1730
1731 #[test]
1732 fn parallel_hash_state_has_one_fixed_entry_per_row_and_is_admitted_up_front() {
1733 assert_eq!(
1734 std::mem::size_of::<ParallelHashEntry>(),
1735 16,
1736 "parallel hash entries must remain fixed width"
1737 );
1738
1739 let build_rows: Vec<Row> = (0..128)
1740 .map(|i| Row::from_values(vec![Value::Integer(i)]))
1741 .collect();
1742 let retained = ParallelHashTable::retained_bytes(build_rows.len()).unwrap();
1743 assert_eq!(
1744 retained,
1745 crate::hash_table::JoinHashTable::estimated_retained_bytes(build_rows.len()).unwrap()
1746 );
1747 assert!(parallel_join_state_fits_budget(
1748 build_rows.len(),
1749 false,
1750 retained
1751 ));
1752 assert!(!parallel_join_state_fits_budget(
1753 build_rows.len(),
1754 false,
1755 retained - 1
1756 ));
1757
1758 let config = ParallelConfig {
1759 min_rows_for_parallel_join: 1,
1760 chunk_size: 8,
1761 ..Default::default()
1762 };
1763 let error = parallel_hash_build_inner(&build_rows, &[0], &config, None, retained - 1)
1764 .err()
1765 .expect("over-budget table must be rejected before build");
1766 assert!(error.to_string().contains("exceeds memory budget"));
1767 }
1768
1769 #[test]
1770 fn outer_match_tracker_is_part_of_parallel_join_budget() {
1771 let rows = 128;
1772 let hash_bytes = ParallelHashTable::retained_bytes(rows).unwrap();
1773 let total = parallel_join_state_retained_bytes(rows, true).unwrap();
1774 assert_eq!(total, hash_bytes + rows * std::mem::size_of::<AtomicBool>());
1775 assert!(!parallel_join_state_fits_budget(rows, true, hash_bytes));
1776 assert!(parallel_join_state_fits_budget(rows, true, total));
1777 }
1778
1779 #[test]
1780 fn parallel_pull_join_bounds_probe_batches_and_releases_request_memory() {
1781 let build_rows = CompactArc::new(vec![
1782 Row::from_values(vec![Value::Integer(1)]),
1783 Row::from_values(vec![Value::Integer(1)]),
1784 Row::from_values(vec![Value::Integer(1)]),
1785 Row::from_values(vec![Value::Integer(2)]),
1786 ]);
1787 let probe_rows = vec![
1788 Row::from_values(vec![Value::Integer(1)]),
1789 Row::from_values(vec![Value::Integer(2)]),
1790 Row::from_values(vec![Value::Integer(3)]),
1791 Row::from_values(vec![Value::Integer(1)]),
1792 Row::from_values(vec![Value::Integer(4)]),
1793 ];
1794 let probe = Box::new(crate::operator::MaterializedOperator::new(
1795 probe_rows,
1796 vec![ColumnInfo::new("probe.id")],
1797 ));
1798 let config = ParallelConfig {
1799 min_rows_for_parallel_join: 1,
1800 chunk_size: 1,
1801 join_output_batch_rows: 2,
1802 join_output_batch_bytes: 4096,
1803 ..Default::default()
1804 };
1805 let ctx = ExecutionContext::new();
1806 let state_bytes = parallel_join_state_retained_bytes(build_rows.len(), false).unwrap();
1807 let state_reservation = ctx.reserve_join_memory(state_bytes).unwrap();
1808 let batch_reservation = ctx
1809 .reserve_join_memory(config.join_output_batch_bytes)
1810 .unwrap();
1811 let cancellation = ctx.cancellation_handle();
1812 let mut operator = ParallelHashJoinOperator::new(
1813 probe,
1814 build_rows,
1815 vec![0],
1816 vec![0],
1817 JoinType::Left,
1818 false,
1819 1,
1820 1,
1821 vec!["probe.id".to_string(), "build.id".to_string()],
1822 None,
1823 config,
1824 cancellation,
1825 state_bytes,
1826 state_reservation,
1827 batch_reservation,
1828 )
1829 .unwrap();
1830
1831 operator.open().unwrap();
1832 assert_eq!(operator.peak_probe_batch_rows, 0);
1833 let first = operator.next().unwrap().unwrap().into_owned();
1834 assert_eq!(first.get(0), Some(&Value::Integer(1)));
1835 assert_eq!(operator.peak_probe_batch_rows, 2);
1836
1837 let mut output_rows = 1;
1838 while operator.next().unwrap().is_some() {
1839 output_rows += 1;
1840 assert!(operator.probe_batch.len() <= 2);
1841 assert!(operator.ready.len() <= 2);
1842 }
1843 assert_eq!(output_rows, 9);
1844 assert_eq!(operator.peak_probe_batch_rows, 2);
1845 assert!(ctx.retained_join_memory_bytes() > 0);
1846 operator.close().unwrap();
1847 assert_eq!(ctx.retained_join_memory_bytes(), 0);
1848 }
1849
1850 #[test]
1851 fn parallel_pull_join_cancellation_closes_without_retained_state() {
1852 let build_rows = CompactArc::new(
1853 (0..128)
1854 .map(|value| Row::from_values(vec![Value::Integer(value)]))
1855 .collect::<Vec<_>>(),
1856 );
1857 let probe = Box::new(crate::operator::MaterializedOperator::new(
1858 vec![Row::from_values(vec![Value::Integer(1)])],
1859 vec![ColumnInfo::new("probe.id")],
1860 ));
1861 let config = ParallelConfig {
1862 min_rows_for_parallel_join: 1,
1863 join_output_batch_rows: 2,
1864 join_output_batch_bytes: 4096,
1865 ..Default::default()
1866 };
1867 let ctx = ExecutionContext::new();
1868 let state_bytes = parallel_join_state_retained_bytes(build_rows.len(), false).unwrap();
1869 let state_reservation = ctx.reserve_join_memory(state_bytes).unwrap();
1870 let batch_reservation = ctx
1871 .reserve_join_memory(config.join_output_batch_bytes)
1872 .unwrap();
1873 let cancellation = ctx.cancellation_handle();
1874 let mut operator = ParallelHashJoinOperator::new(
1875 probe,
1876 build_rows,
1877 vec![0],
1878 vec![0],
1879 JoinType::Inner,
1880 false,
1881 1,
1882 1,
1883 vec!["probe.id".to_string(), "build.id".to_string()],
1884 None,
1885 config,
1886 cancellation.clone(),
1887 state_bytes,
1888 state_reservation,
1889 batch_reservation,
1890 )
1891 .unwrap();
1892
1893 operator.open().unwrap();
1894 cancellation.cancel();
1895 assert!(matches!(
1896 operator.next(),
1897 Err(radixdb_core::Error::QueryCancelled)
1898 ));
1899 operator.close().unwrap();
1900 assert_eq!(ctx.retained_join_memory_bytes(), 0);
1901 }
1902
1903 #[test]
1904 fn test_verify_key_match() {
1905 let row1 = Row::from_values(vec![Value::Integer(1), Value::Text("a".to_string().into())]);
1906 let row2 = Row::from_values(vec![Value::Integer(1), Value::Text("b".to_string().into())]);
1907 let row3 = Row::from_values(vec![Value::Integer(2), Value::Text("a".to_string().into())]);
1908
1909 assert!(verify_key_match(&row1, &row2, &[0], &[0]));
1911
1912 assert!(!verify_key_match(&row1, &row3, &[0], &[0]));
1914
1915 assert!(verify_key_match(&row1, &row3, &[1], &[1]));
1917 }
1918
1919 #[test]
1920 fn test_hash_join_numeric_key_contract_sequential_and_parallel() {
1921 const EXACT: i64 = 1_i64 << 53;
1922 let build_rows = vec![
1923 Row::from_values(vec![Value::Integer(EXACT)]),
1924 Row::from_values(vec![Value::Integer(EXACT + 1)]),
1925 Row::from_values(vec![Value::Float(-0.0)]),
1926 Row::from_values(vec![Value::Float(f64::NAN)]),
1927 ];
1928 let probe_rows = vec![
1929 Row::from_values(vec![Value::Float(EXACT as f64)]),
1930 Row::from_values(vec![Value::Integer(0)]),
1931 Row::from_values(vec![Value::Float(f64::from_bits(0x7ff8_0000_0000_0042))]),
1932 ];
1933
1934 for (name, config) in [
1935 (
1936 "sequential",
1937 ParallelConfig {
1938 min_rows_for_parallel_join: usize::MAX,
1939 ..Default::default()
1940 },
1941 ),
1942 (
1943 "parallel",
1944 ParallelConfig {
1945 min_rows_for_parallel_join: 1,
1946 chunk_size: 1,
1947 ..Default::default()
1948 },
1949 ),
1950 ] {
1951 for swapped in [false, true] {
1952 let result = parallel_hash_join(
1953 &probe_rows,
1954 &build_rows,
1955 &[0],
1956 &[0],
1957 JoinType::Inner,
1958 1,
1959 1,
1960 swapped,
1961 &config,
1962 );
1963 assert_eq!(
1964 result.rows.len(),
1965 3,
1966 "{name} numeric join with swapped={swapped}"
1967 );
1968 let _ = name;
1969 }
1970 }
1971 }
1972
1973 #[test]
1979 fn test_parallel_hash_join_collision_handling() {
1980 let build_rows: Vec<Row> = vec![
1982 Row::from_values(vec![Value::Integer(1), Value::Text("build_a".into())]),
1983 Row::from_values(vec![Value::Integer(2), Value::Text("build_b".into())]),
1984 Row::from_values(vec![Value::Integer(3), Value::Text("build_c".into())]),
1985 ];
1986
1987 let probe_rows: Vec<Row> = vec![
1989 Row::from_values(vec![Value::Integer(1), Value::Text("probe_x".into())]),
1990 Row::from_values(vec![Value::Integer(2), Value::Text("probe_y".into())]),
1991 Row::from_values(vec![Value::Integer(4), Value::Text("probe_z".into())]), ];
1993
1994 let config = ParallelConfig {
1995 min_rows_for_parallel_join: 1, ..Default::default()
1997 };
1998
1999 let result = parallel_hash_join(
2001 &probe_rows,
2002 &build_rows,
2003 &[0], &[0], JoinType::Inner,
2006 2, 2, false,
2009 &config,
2010 );
2011
2012 assert_eq!(result.rows.len(), 2, "INNER JOIN should have 2 matches");
2014
2015 for row in &result.rows {
2017 assert_eq!(row.len(), 4);
2019 }
2020 }
2021
2022 #[test]
2024 fn test_parallel_left_join_unmatched() {
2025 let build_rows: Vec<Row> = vec![Row::from_values(vec![
2026 Value::Integer(1),
2027 Value::Text("match".into()),
2028 ])];
2029
2030 let probe_rows: Vec<Row> = vec![
2031 Row::from_values(vec![Value::Integer(1), Value::Text("p1".into())]), Row::from_values(vec![Value::Integer(2), Value::Text("p2".into())]), Row::from_values(vec![Value::Integer(3), Value::Text("p3".into())]), ];
2035
2036 let config = ParallelConfig {
2037 min_rows_for_parallel_join: 1,
2038 ..Default::default()
2039 };
2040
2041 let result = parallel_hash_join(
2042 &probe_rows,
2043 &build_rows,
2044 &[0],
2045 &[0],
2046 JoinType::Left,
2047 2,
2048 2,
2049 false,
2050 &config,
2051 );
2052
2053 assert_eq!(result.rows.len(), 3, "LEFT JOIN should have 3 rows");
2055
2056 let null_count = result
2058 .rows
2059 .iter()
2060 .filter(|r| {
2061 r.get(2).map(|v| v.is_null()).unwrap_or(false)
2062 && r.get(3).map(|v| v.is_null()).unwrap_or(false)
2063 })
2064 .count();
2065 assert_eq!(
2066 null_count, 2,
2067 "Should have 2 unmatched rows with NULL build columns"
2068 );
2069 }
2070
2071 #[test]
2072 fn test_parallel_left_join_fuses_projection_before_output_materialization() {
2073 let build_rows = vec![Row::from_values(vec![
2074 Value::Integer(1),
2075 Value::Text("build".into()),
2076 Value::Text("unused-build-payload".repeat(128).into()),
2077 ])];
2078 let probe_rows = vec![
2079 Row::from_values(vec![
2080 Value::Integer(1),
2081 Value::Text("matched".into()),
2082 Value::Text("unused-probe-payload".repeat(128).into()),
2083 ]),
2084 Row::from_values(vec![
2085 Value::Integer(2),
2086 Value::Text("unmatched".into()),
2087 Value::Text("unused-probe-payload".repeat(128).into()),
2088 ]),
2089 ];
2090 let projection = [ColumnSource::Outer(1), ColumnSource::Inner(1)];
2091 let config = ParallelConfig {
2092 min_rows_for_parallel_join: 1,
2093 ..Default::default()
2094 };
2095
2096 let result = parallel_hash_join_inner(
2097 &probe_rows,
2098 &build_rows,
2099 &[0],
2100 &[0],
2101 JoinType::Left,
2102 3,
2103 3,
2104 false,
2105 Some(&projection),
2106 &config,
2107 None,
2108 usize::MAX,
2109 )
2110 .unwrap();
2111
2112 assert_eq!(result.rows.len(), 2);
2113 assert!(result.rows.iter().all(|row| row.len() == 2));
2114 assert!(result.rows.iter().any(|row| {
2115 row.get(0) == Some(&Value::text("matched")) && row.get(1) == Some(&Value::text("build"))
2116 }));
2117 assert!(result.rows.iter().any(|row| {
2118 row.get(0) == Some(&Value::text("unmatched")) && row.get(1).is_some_and(Value::is_null)
2119 }));
2120 }
2121
2122 #[test]
2124 fn test_parallel_right_join_unmatched() {
2125 let build_rows: Vec<Row> = vec![
2126 Row::from_values(vec![Value::Integer(1), Value::Text("b1".into())]), Row::from_values(vec![Value::Integer(2), Value::Text("b2".into())]), Row::from_values(vec![Value::Integer(3), Value::Text("b3".into())]), ];
2130
2131 let probe_rows: Vec<Row> = vec![Row::from_values(vec![
2132 Value::Integer(1),
2133 Value::Text("p1".into()),
2134 ])];
2135
2136 let config = ParallelConfig {
2137 min_rows_for_parallel_join: 1,
2138 ..Default::default()
2139 };
2140
2141 let result = parallel_hash_join(
2142 &probe_rows,
2143 &build_rows,
2144 &[0],
2145 &[0],
2146 JoinType::Right,
2147 2,
2148 2,
2149 false,
2150 &config,
2151 );
2152
2153 assert_eq!(result.rows.len(), 3, "RIGHT JOIN should have 3 rows");
2155
2156 let null_count = result
2158 .rows
2159 .iter()
2160 .filter(|r| {
2161 r.get(0).map(|v| v.is_null()).unwrap_or(false)
2162 && r.get(1).map(|v| v.is_null()).unwrap_or(false)
2163 })
2164 .count();
2165 assert_eq!(
2166 null_count, 2,
2167 "Should have 2 unmatched rows with NULL probe columns"
2168 );
2169 }
2170
2171 #[test]
2173 fn test_parallel_full_outer_join() {
2174 let build_rows: Vec<Row> = vec![
2175 Row::from_values(vec![Value::Integer(1), Value::Text("b1".into())]), Row::from_values(vec![Value::Integer(3), Value::Text("b3".into())]), ];
2178
2179 let probe_rows: Vec<Row> = vec![
2180 Row::from_values(vec![Value::Integer(1), Value::Text("p1".into())]), Row::from_values(vec![Value::Integer(2), Value::Text("p2".into())]), ];
2183
2184 let config = ParallelConfig {
2185 min_rows_for_parallel_join: 1,
2186 ..Default::default()
2187 };
2188
2189 let result = parallel_hash_join(
2190 &probe_rows,
2191 &build_rows,
2192 &[0],
2193 &[0],
2194 JoinType::Full,
2195 2,
2196 2,
2197 false,
2198 &config,
2199 );
2200
2201 assert_eq!(result.rows.len(), 3, "FULL OUTER JOIN should have 3 rows");
2206 }
2207
2208 #[test]
2210 fn test_parallel_join_empty_tables() {
2211 let config = ParallelConfig::default();
2212
2213 let result = parallel_hash_join(
2215 &[],
2216 &[Row::from_values(vec![Value::Integer(1)])],
2217 &[0],
2218 &[0],
2219 JoinType::Inner,
2220 1,
2221 1,
2222 false,
2223 &config,
2224 );
2225 assert_eq!(
2226 result.rows.len(),
2227 0,
2228 "Empty probe should give empty result for INNER"
2229 );
2230
2231 let result = parallel_hash_join(
2233 &[Row::from_values(vec![Value::Integer(1)])],
2234 &[],
2235 &[0],
2236 &[0],
2237 JoinType::Inner,
2238 1,
2239 1,
2240 false,
2241 &config,
2242 );
2243 assert_eq!(
2244 result.rows.len(),
2245 0,
2246 "Empty build should give empty result for INNER"
2247 );
2248
2249 let result = parallel_hash_join(
2251 &[
2252 Row::from_values(vec![Value::Integer(1)]),
2253 Row::from_values(vec![Value::Integer(2)]),
2254 ],
2255 &[],
2256 &[0],
2257 &[0],
2258 JoinType::Left,
2259 1,
2260 1,
2261 false,
2262 &config,
2263 );
2264 assert_eq!(
2265 result.rows.len(),
2266 2,
2267 "LEFT JOIN with empty build should have all probe rows"
2268 );
2269 }
2270
2271 #[test]
2278 fn test_parallel_join_swapped() {
2279 let build_rows: Vec<Row> = vec![
2283 Row::from_values(vec![Value::Integer(1), Value::Text("b1".into())]),
2284 Row::from_values(vec![Value::Integer(3), Value::Text("b3".into())]), ];
2286
2287 let probe_rows: Vec<Row> = vec![
2288 Row::from_values(vec![Value::Integer(1), Value::Text("p1".into())]),
2289 Row::from_values(vec![Value::Integer(2), Value::Text("p2".into())]), ];
2291
2292 let config = ParallelConfig {
2293 min_rows_for_parallel_join: 1,
2294 ..Default::default()
2295 };
2296
2297 let result = parallel_hash_join(
2301 &probe_rows,
2302 &build_rows,
2303 &[0],
2304 &[0],
2305 JoinType::Left,
2306 2,
2307 2,
2308 true, &config,
2310 );
2311
2312 assert_eq!(result.rows.len(), 2, "LEFT JOIN swapped should have 2 rows");
2316
2317 let matched_row = result
2320 .rows
2321 .iter()
2322 .find(|r| r.get(0) == Some(&Value::Integer(1)) && r.get(2) == Some(&Value::Integer(1)));
2323 assert!(matched_row.is_some(), "Should have a matched row with id=1");
2324
2325 let unmatched_row = result
2327 .rows
2328 .iter()
2329 .find(|r| r.get(0) == Some(&Value::Integer(3)));
2330 assert!(
2331 unmatched_row.is_some(),
2332 "Should have unmatched build row with id=3"
2333 );
2334 let unmatched = unmatched_row.unwrap();
2335 assert!(
2336 unmatched.get(2).map(|v| v.is_null()).unwrap_or(false),
2337 "Probe col should be NULL"
2338 );
2339 assert!(
2340 unmatched.get(3).map(|v| v.is_null()).unwrap_or(false),
2341 "Probe col should be NULL"
2342 );
2343 }
2344}