linkrs_query/executor/streaming/operators/
join_operator.rs

1use std::collections::HashSet;
2use std::sync::Arc;
3
4use crate::executor::base::MemoryTracker;
5use crate::executor::streaming::chunk::{use_columnar_path, DataChunk};
6use crate::executor::streaming::executor::FullOuterJoinPhase;
7use crate::executor::streaming::executor::StreamingExecutor;
8use crate::executor::streaming::operators::source_operator::OperatorConfig;
9use crate::executor::streaming::runtime::ExecutionRuntime;
10use crate::executor::streaming::slot::SlotLayout;
11use linkrs_core::error::QueryError;
12use linkrs_core::types::expr::Expression;
13use linkrs_core::Value;
14
15mod cross_join;
16pub mod grace_join;
17mod hash_join;
18mod merge_join;
19mod nested_loop_join;
20mod semi_join;
21
22pub use hash_join::{HashJoinBuildSide, JoinKeyValue};
23
24fn build_combined_names(
25    left_col_names: &[String],
26    right_col_names: &[String],
27    fallback_right_width: usize,
28) -> Vec<String> {
29    let mut names = left_col_names.to_vec();
30    if !right_col_names.is_empty() {
31        names.extend_from_slice(right_col_names);
32    } else {
33        for i in 0..fallback_right_width {
34            names.push(format!("right_{}", i));
35        }
36    }
37    names
38}
39
40/// Defer the typed columnar layout on a freshly materialized join output.
41///
42/// Join outputs are assembled from rows and (unlike the storage scan) start
43/// row-major, which would otherwise drop the typed fast path for every
44/// operator downstream of a join. Instead of building eagerly, the chunk is
45/// marked deferred when the shared [`ColumnarPolicy`] gate allows columnar:
46/// the first typed consumer rebuilds the layout via
47/// [`DataChunk::ensure_typed_columns`], so large join results that are never
48/// read as typed columns pay no build cost. `pub(crate)` so the join modules
49/// can reuse it.
50pub(crate) fn finalize_join_output(
51    mut chunk: DataChunk,
52    runtime: &Option<Arc<ExecutionRuntime>>,
53) -> DataChunk {
54    if use_columnar_path(runtime) {
55        chunk.columnar_build_deferred = true;
56    }
57    chunk
58}
59
60#[derive(Debug)]
61pub enum JoinOperatorKind {
62    HashJoin {
63        join_condition: Option<Expression>,
64        hash_keys: Vec<Expression>,
65        probe_keys: Vec<Expression>,
66        build_side: HashJoinBuildSide,
67        build_done: bool,
68        memory_tracker: MemoryTracker,
69        right_col_names: Vec<String>,
70        build_side_select: super::spec::BuildSide,
71        grace: grace_join::GraceJoinState,
72    },
73    HashLeftJoin {
74        join_condition: Option<Expression>,
75        hash_keys: Vec<Expression>,
76        probe_keys: Vec<Expression>,
77        build_side: HashJoinBuildSide,
78        build_done: bool,
79        memory_tracker: MemoryTracker,
80        right_col_names: Vec<String>,
81        build_side_select: super::spec::BuildSide,
82        grace: grace_join::GraceJoinState,
83    },
84    NestedLoopJoin {
85        join_condition: Option<Expression>,
86        build_side_tuples: Vec<Vec<Value>>,
87        build_done: bool,
88        memory_tracker: MemoryTracker,
89        right_col_names: Vec<String>,
90    },
91    InnerJoin {
92        join_condition: Option<Expression>,
93        build_side_tuples: Vec<Vec<Value>>,
94        build_done: bool,
95        memory_tracker: MemoryTracker,
96        right_col_names: Vec<String>,
97    },
98    LeftJoin {
99        join_condition: Option<Expression>,
100        build_side_tuples: Vec<Vec<Value>>,
101        build_done: bool,
102        memory_tracker: MemoryTracker,
103        right_col_names: Vec<String>,
104    },
105    RightJoin {
106        join_condition: Option<Expression>,
107        build_side_tuples: Vec<Vec<Value>>,
108        right_consumed: bool,
109        memory_tracker: MemoryTracker,
110        right_col_names: Vec<String>,
111    },
112    FullOuterJoin {
113        join_condition: Option<Expression>,
114        left_rows: Vec<Vec<Value>>,
115        right_rows: Vec<Vec<Value>>,
116        matched_right_indices: HashSet<usize>,
117        result_iter: Option<std::vec::IntoIter<Vec<Value>>>,
118        phase: FullOuterJoinPhase,
119        memory_tracker: MemoryTracker,
120        right_col_names: Vec<String>,
121    },
122    CrossJoin {
123        all_left_rows: Vec<Vec<Value>>,
124        all_right_rows: Vec<Vec<Value>>,
125        left_consumed: bool,
126        right_consumed: bool,
127        memory_tracker: MemoryTracker,
128        right_col_names: Vec<String>,
129        output_done: bool,
130    },
131    SemiJoin {
132        join_condition: Option<Expression>,
133        // NOT EXISTS semantics: keep left rows with NO matching right row.
134        anti: bool,
135        right_rows: Vec<Vec<Value>>,
136        right_consumed: bool,
137        memory_tracker: MemoryTracker,
138        right_col_names: Vec<String>,
139    },
140}
141
142/// Join operator.
143///
144/// Wraps [`JoinOperatorKind`] with the runtime context injected at `open()`.
145/// Lifecycle state is owned exclusively by the executor; operators never
146/// write it.
147#[derive(Debug)]
148pub struct JoinOperator {
149    pub kind: JoinOperatorKind,
150    pub runtime: Option<Arc<ExecutionRuntime>>,
151    pub output_layout: Arc<SlotLayout>,
152    pub config: OperatorConfig,
153}
154
155impl JoinOperator {
156    pub fn new(kind: JoinOperatorKind, output_layout: Arc<SlotLayout>) -> Self {
157        Self {
158            kind,
159            runtime: None,
160            output_layout,
161            config: OperatorConfig::default(),
162        }
163    }
164
165    /// Inject the runtime and execution config (called once by the executor
166    /// before this operator produces any data).
167    pub fn inject_context(
168        &mut self,
169        runtime: Option<&Arc<ExecutionRuntime>>,
170        config: OperatorConfig,
171    ) {
172        if let Some(rt) = runtime {
173            self.runtime = Some(rt.clone());
174        }
175        self.config = config;
176    }
177
178    pub fn from_spec(
179        spec: &super::spec::JoinSpec,
180        memory_budget: &crate::executor::base::MemoryBudget,
181        output_layout: Arc<SlotLayout>,
182    ) -> Self {
183        let kind = match spec {
184            super::spec::JoinSpec::HashJoin {
185                join_condition,
186                hash_keys,
187                probe_keys,
188                build_side,
189            } => JoinOperatorKind::HashJoin {
190                join_condition: join_condition.clone(),
191                hash_keys: hash_keys.clone(),
192                probe_keys: probe_keys.clone(),
193                build_side: HashJoinBuildSide::new(),
194                build_done: false,
195                memory_tracker: crate::executor::base::MemoryTracker::new(memory_budget.clone()),
196                right_col_names: Vec::new(),
197                build_side_select: *build_side,
198                grace: grace_join::GraceJoinState::default(),
199            },
200            super::spec::JoinSpec::HashLeftJoin {
201                join_condition,
202                hash_keys,
203                probe_keys,
204                build_side,
205            } => JoinOperatorKind::HashLeftJoin {
206                join_condition: join_condition.clone(),
207                hash_keys: hash_keys.clone(),
208                probe_keys: probe_keys.clone(),
209                build_side: HashJoinBuildSide::new(),
210                build_done: false,
211                memory_tracker: crate::executor::base::MemoryTracker::new(memory_budget.clone()),
212                right_col_names: Vec::new(),
213                build_side_select: *build_side,
214                grace: grace_join::GraceJoinState::default(),
215            },
216            super::spec::JoinSpec::NestedLoopJoin { join_condition } => {
217                JoinOperatorKind::NestedLoopJoin {
218                    join_condition: join_condition.clone(),
219                    build_side_tuples: Vec::new(),
220                    build_done: false,
221                    memory_tracker: crate::executor::base::MemoryTracker::new(
222                        memory_budget.clone(),
223                    ),
224                    right_col_names: Vec::new(),
225                }
226            }
227            super::spec::JoinSpec::InnerJoin { join_condition } => JoinOperatorKind::InnerJoin {
228                join_condition: join_condition.clone(),
229                build_side_tuples: Vec::new(),
230                build_done: false,
231                memory_tracker: crate::executor::base::MemoryTracker::new(memory_budget.clone()),
232                right_col_names: Vec::new(),
233            },
234            super::spec::JoinSpec::LeftJoin { join_condition } => JoinOperatorKind::LeftJoin {
235                join_condition: join_condition.clone(),
236                build_side_tuples: Vec::new(),
237                build_done: false,
238                memory_tracker: crate::executor::base::MemoryTracker::new(memory_budget.clone()),
239                right_col_names: Vec::new(),
240            },
241            super::spec::JoinSpec::RightJoin { join_condition } => JoinOperatorKind::RightJoin {
242                join_condition: join_condition.clone(),
243                build_side_tuples: Vec::new(),
244                right_consumed: false,
245                memory_tracker: crate::executor::base::MemoryTracker::new(memory_budget.clone()),
246                right_col_names: Vec::new(),
247            },
248            super::spec::JoinSpec::FullOuterJoin { join_condition } => {
249                JoinOperatorKind::FullOuterJoin {
250                    join_condition: join_condition.clone(),
251                    left_rows: Vec::new(),
252                    right_rows: Vec::new(),
253                    matched_right_indices: std::collections::HashSet::new(),
254                    result_iter: None,
255                    phase: FullOuterJoinPhase::BuildingRight,
256                    memory_tracker: crate::executor::base::MemoryTracker::new(
257                        memory_budget.clone(),
258                    ),
259                    right_col_names: Vec::new(),
260                }
261            }
262            super::spec::JoinSpec::CrossJoin => JoinOperatorKind::CrossJoin {
263                all_left_rows: Vec::new(),
264                all_right_rows: Vec::new(),
265                left_consumed: false,
266                right_consumed: false,
267                memory_tracker: crate::executor::base::MemoryTracker::new(memory_budget.clone()),
268                right_col_names: Vec::new(),
269                output_done: false,
270            },
271            super::spec::JoinSpec::SemiJoin {
272                join_condition,
273                anti,
274            } => JoinOperatorKind::SemiJoin {
275                join_condition: join_condition.clone(),
276                anti: *anti,
277                right_rows: Vec::new(),
278                right_consumed: false,
279                memory_tracker: crate::executor::base::MemoryTracker::new(memory_budget.clone()),
280                right_col_names: Vec::new(),
281            },
282        };
283        Self::new(kind, output_layout)
284    }
285
286    pub fn memory_tracker(&self) -> &MemoryTracker {
287        match &self.kind {
288            JoinOperatorKind::HashJoin { memory_tracker, .. }
289            | JoinOperatorKind::HashLeftJoin { memory_tracker, .. }
290            | JoinOperatorKind::NestedLoopJoin { memory_tracker, .. }
291            | JoinOperatorKind::InnerJoin { memory_tracker, .. }
292            | JoinOperatorKind::LeftJoin { memory_tracker, .. }
293            | JoinOperatorKind::RightJoin { memory_tracker, .. }
294            | JoinOperatorKind::FullOuterJoin { memory_tracker, .. }
295            | JoinOperatorKind::CrossJoin { memory_tracker, .. }
296            | JoinOperatorKind::SemiJoin { memory_tracker, .. } => memory_tracker,
297        }
298    }
299
300    pub fn open(
301        &mut self,
302        left: &mut StreamingExecutor,
303        right: &mut StreamingExecutor,
304    ) -> Result<(), QueryError> {
305        left.open()?;
306        right.open()?;
307        Ok(())
308    }
309
310    pub fn next(
311        &mut self,
312        left: &mut StreamingExecutor,
313        right: &mut StreamingExecutor,
314    ) -> Result<Option<DataChunk>, QueryError> {
315        let runtime = &self.runtime;
316        let output_layout = &self.output_layout;
317        match &mut self.kind {
318            JoinOperatorKind::HashJoin {
319                join_condition,
320                hash_keys,
321                probe_keys,
322                build_side,
323                build_done,
324                memory_tracker,
325                right_col_names,
326                build_side_select,
327                grace,
328            } => hash_join::next_hash_join(
329                join_condition,
330                hash_keys,
331                probe_keys,
332                build_side,
333                build_done,
334                memory_tracker,
335                right_col_names,
336                *build_side_select,
337                grace,
338                left,
339                right,
340                runtime,
341                output_layout,
342            ),
343            JoinOperatorKind::HashLeftJoin {
344                join_condition,
345                hash_keys,
346                probe_keys,
347                build_side,
348                build_done,
349                memory_tracker,
350                right_col_names,
351                build_side_select,
352                grace,
353            } => hash_join::next_hash_left_join(
354                join_condition,
355                hash_keys,
356                probe_keys,
357                build_side,
358                build_done,
359                memory_tracker,
360                right_col_names,
361                *build_side_select,
362                grace,
363                left,
364                right,
365                runtime,
366                output_layout,
367            ),
368            JoinOperatorKind::NestedLoopJoin {
369                join_condition,
370                build_side_tuples,
371                build_done,
372                memory_tracker,
373                right_col_names,
374            } => nested_loop_join::next_nested_loop_join(
375                join_condition,
376                build_side_tuples,
377                build_done,
378                memory_tracker,
379                right_col_names,
380                left,
381                right,
382                runtime,
383                output_layout,
384            ),
385            JoinOperatorKind::InnerJoin {
386                join_condition,
387                build_side_tuples,
388                build_done,
389                memory_tracker,
390                right_col_names,
391            } => merge_join::next_inner_join(
392                join_condition,
393                build_side_tuples,
394                build_done,
395                memory_tracker,
396                right_col_names,
397                left,
398                right,
399                runtime,
400                output_layout,
401            ),
402            JoinOperatorKind::LeftJoin {
403                join_condition,
404                build_side_tuples,
405                build_done,
406                memory_tracker,
407                right_col_names,
408            } => merge_join::next_left_join(
409                join_condition,
410                build_side_tuples,
411                build_done,
412                memory_tracker,
413                right_col_names,
414                left,
415                right,
416                runtime,
417                output_layout,
418            ),
419            JoinOperatorKind::RightJoin {
420                join_condition,
421                build_side_tuples,
422                right_consumed,
423                memory_tracker,
424                right_col_names,
425            } => merge_join::next_right_join(
426                join_condition,
427                build_side_tuples,
428                right_consumed,
429                memory_tracker,
430                right_col_names,
431                left,
432                right,
433                runtime,
434                output_layout,
435            ),
436            JoinOperatorKind::FullOuterJoin {
437                join_condition,
438                left_rows,
439                right_rows,
440                matched_right_indices,
441                result_iter,
442                phase,
443                memory_tracker,
444                right_col_names,
445            } => merge_join::next_full_outer_join(
446                join_condition,
447                left_rows,
448                right_rows,
449                matched_right_indices,
450                result_iter,
451                phase,
452                memory_tracker,
453                right_col_names,
454                left,
455                right,
456                runtime,
457                output_layout,
458            ),
459            JoinOperatorKind::CrossJoin {
460                all_left_rows,
461                all_right_rows,
462                left_consumed,
463                right_consumed,
464                memory_tracker,
465                right_col_names,
466                output_done,
467            } => cross_join::next_cross_join(
468                all_left_rows,
469                all_right_rows,
470                left_consumed,
471                right_consumed,
472                output_done,
473                memory_tracker,
474                right_col_names,
475                left,
476                right,
477                runtime,
478                output_layout,
479            ),
480            JoinOperatorKind::SemiJoin {
481                join_condition,
482                anti,
483                right_rows,
484                right_consumed,
485                memory_tracker,
486                right_col_names,
487            } => semi_join::next_semi_join(
488                join_condition,
489                *anti,
490                right_rows,
491                right_consumed,
492                memory_tracker,
493                right_col_names,
494                left,
495                right,
496                runtime,
497                output_layout,
498            ),
499        }
500    }
501
502    pub fn stop(&mut self) -> Result<(), QueryError> {
503        Ok(())
504    }
505
506    /// Reset per-run join state (hash tables, buffered sides, phase flags)
507    /// and rewind both inputs so the join re-produces the same result set.
508    pub fn reset(
509        &mut self,
510        left: &mut StreamingExecutor,
511        right: &mut StreamingExecutor,
512    ) -> Result<bool, QueryError> {
513        match &mut self.kind {
514            JoinOperatorKind::HashJoin {
515                build_side,
516                build_done,
517                right_col_names,
518                grace,
519                memory_tracker,
520                ..
521            }
522            | JoinOperatorKind::HashLeftJoin {
523                build_side,
524                build_done,
525                right_col_names,
526                grace,
527                memory_tracker,
528                ..
529            } => {
530                grace.cleanup_files();
531                *grace = grace_join::GraceJoinState::default();
532                *build_side = HashJoinBuildSide::new();
533                *build_done = false;
534                right_col_names.clear();
535                memory_tracker.reset();
536            }
537            JoinOperatorKind::NestedLoopJoin {
538                build_side_tuples,
539                build_done,
540                right_col_names,
541                ..
542            }
543            | JoinOperatorKind::InnerJoin {
544                build_side_tuples,
545                build_done,
546                right_col_names,
547                ..
548            }
549            | JoinOperatorKind::LeftJoin {
550                build_side_tuples,
551                build_done,
552                right_col_names,
553                ..
554            } => {
555                build_side_tuples.clear();
556                *build_done = false;
557                right_col_names.clear();
558            }
559            JoinOperatorKind::RightJoin {
560                build_side_tuples,
561                right_consumed,
562                right_col_names,
563                ..
564            } => {
565                build_side_tuples.clear();
566                *right_consumed = false;
567                right_col_names.clear();
568            }
569            JoinOperatorKind::FullOuterJoin {
570                left_rows,
571                right_rows,
572                matched_right_indices,
573                result_iter,
574                phase,
575                right_col_names,
576                ..
577            } => {
578                left_rows.clear();
579                right_rows.clear();
580                matched_right_indices.clear();
581                *result_iter = None;
582                *phase = FullOuterJoinPhase::BuildingRight;
583                right_col_names.clear();
584            }
585            JoinOperatorKind::CrossJoin {
586                all_left_rows,
587                all_right_rows,
588                left_consumed,
589                right_consumed,
590                right_col_names,
591                output_done,
592                ..
593            } => {
594                all_left_rows.clear();
595                all_right_rows.clear();
596                *left_consumed = false;
597                *right_consumed = false;
598                right_col_names.clear();
599                *output_done = false;
600            }
601            JoinOperatorKind::SemiJoin {
602                right_rows,
603                right_consumed,
604                right_col_names,
605                ..
606            } => {
607                right_rows.clear();
608                *right_consumed = false;
609                right_col_names.clear();
610            }
611        }
612        left.reset()?;
613        right.reset()?;
614        Ok(false)
615    }
616
617    pub fn close(&mut self) -> Result<(), QueryError> {
618        match &mut self.kind {
619            JoinOperatorKind::HashJoin {
620                build_side,
621                memory_tracker,
622                grace,
623                ..
624            } => {
625                grace.cleanup_files();
626                hash_join::close(memory_tracker, build_side)
627            }
628            JoinOperatorKind::HashLeftJoin {
629                build_side,
630                memory_tracker,
631                grace,
632                ..
633            } => {
634                grace.cleanup_files();
635                hash_join::close(memory_tracker, build_side)
636            }
637            JoinOperatorKind::NestedLoopJoin {
638                build_side_tuples,
639                memory_tracker,
640                ..
641            } => nested_loop_join::close(memory_tracker, build_side_tuples),
642            JoinOperatorKind::InnerJoin {
643                build_side_tuples,
644                memory_tracker,
645                ..
646            } => nested_loop_join::close(memory_tracker, build_side_tuples),
647            JoinOperatorKind::LeftJoin {
648                build_side_tuples,
649                memory_tracker,
650                ..
651            } => nested_loop_join::close(memory_tracker, build_side_tuples),
652            JoinOperatorKind::RightJoin {
653                build_side_tuples,
654                memory_tracker,
655                ..
656            } => nested_loop_join::close(memory_tracker, build_side_tuples),
657            JoinOperatorKind::FullOuterJoin {
658                left_rows,
659                right_rows,
660                memory_tracker,
661                ..
662            } => merge_join::close_full_outer(memory_tracker, left_rows, right_rows),
663            JoinOperatorKind::CrossJoin {
664                all_left_rows,
665                all_right_rows,
666                memory_tracker,
667                ..
668            } => cross_join::close_cross(memory_tracker, all_left_rows, all_right_rows),
669            JoinOperatorKind::SemiJoin {
670                right_rows,
671                memory_tracker,
672                ..
673            } => semi_join::close_semi(memory_tracker, right_rows),
674        }
675    }
676
677    /// Spill the hash join build side to partitioned runs.
678    ///
679    /// Only effective during the build phase (before `build_done`): buffered
680    /// build rows are re-partitioned to disk and the remaining build input is
681    /// routed straight to the spill writers by the build loop. Once the probe
682    /// phase has started the in-memory table must stay resident and this is a
683    /// no-op; other join kinds keep their budget-failure semantics.
684    pub fn spill_with_manager(
685        &mut self,
686        sm: &Arc<crate::executor::streaming::spill::SpillManager>,
687    ) -> Result<(), linkrs_core::error::QueryError> {
688        let (build_side, build_done, memory_tracker, right_col_names, grace) = match &mut self.kind
689        {
690            JoinOperatorKind::HashJoin {
691                build_side,
692                build_done,
693                memory_tracker,
694                right_col_names,
695                grace,
696                ..
697            }
698            | JoinOperatorKind::HashLeftJoin {
699                build_side,
700                build_done,
701                memory_tracker,
702                right_col_names,
703                grace,
704                ..
705            } => (
706                build_side,
707                build_done,
708                memory_tracker,
709                right_col_names,
710                grace,
711            ),
712            _ => return Ok(()),
713        };
714        if *build_done || grace.partitioned.is_some() {
715            return Ok(());
716        }
717        if build_side.row_count() == 0 && grace.pending_build.is_none() {
718            return Ok(());
719        }
720        grace_join::spill_build_side(
721            Arc::clone(sm),
722            build_side,
723            memory_tracker,
724            right_col_names,
725            grace,
726        )
727    }
728
729    pub fn spilled_bytes(&self) -> u64 {
730        match &self.kind {
731            JoinOperatorKind::HashJoin { grace, .. }
732            | JoinOperatorKind::HashLeftJoin { grace, .. } => grace.spilled_bytes,
733            _ => 0,
734        }
735    }
736
737    pub fn spill_count(&self) -> u64 {
738        match &self.kind {
739            JoinOperatorKind::HashJoin { grace, .. }
740            | JoinOperatorKind::HashLeftJoin { grace, .. } => grace.spill_runs,
741            _ => 0,
742        }
743    }
744
745    pub fn spilled_rows(&self) -> u64 {
746        match &self.kind {
747            JoinOperatorKind::HashJoin { grace, .. }
748            | JoinOperatorKind::HashLeftJoin { grace, .. } => grace.spilled_rows,
749            _ => 0,
750        }
751    }
752}
linkrs_query/executor/streaming/operators/join_operator.rs

linkrs_query/executor/streaming/operators/
join_operator.rs
