Skip to main content

grafeo_core/execution/operators/
join.rs

1//! Join operators for combining data from two sources.
2//!
3//! This module provides:
4//! - `HashJoinOperator`: Efficient hash-based join for equality conditions
5//! - `NestedLoopJoinOperator`: General-purpose join for any condition
6
7use std::collections::HashMap;
8
9use grafeo_common::types::{LogicalType, Value};
10
11use super::{Operator, OperatorError, OperatorResult};
12use crate::execution::chunk::DataChunkBuilder;
13use crate::execution::{DataChunk, ValueVector};
14
15/// The type of join to perform.
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum JoinType {
18    /// Inner join: only matching rows from both sides.
19    Inner,
20    /// Left outer join: all rows from left, matching from right (nulls if no match).
21    Left,
22    /// Right outer join: all rows from right, matching from left (nulls if no match).
23    Right,
24    /// Full outer join: all rows from both sides.
25    Full,
26    /// Cross join: cartesian product of both sides.
27    Cross,
28    /// Semi join: rows from left that have a match in right.
29    Semi,
30    /// Anti join: rows from left that have no match in right.
31    Anti,
32}
33
34/// A hash key that can be hashed and compared for join operations.
35#[derive(Debug, Clone, PartialEq, Eq, Hash)]
36pub enum HashKey {
37    /// Null key.
38    Null,
39    /// Boolean key.
40    Bool(bool),
41    /// Integer key.
42    Int64(i64),
43    /// String key (using the string content for hashing).
44    String(String),
45    /// Composite key for multi-column joins.
46    Composite(Vec<HashKey>),
47}
48
49impl HashKey {
50    /// Creates a hash key from a Value.
51    pub fn from_value(value: &Value) -> Self {
52        match value {
53            Value::Null => HashKey::Null,
54            Value::Bool(b) => HashKey::Bool(*b),
55            Value::Int64(i) => HashKey::Int64(*i),
56            Value::Float64(f) => {
57                // Convert float to bits for consistent hashing
58                HashKey::Int64(f.to_bits() as i64)
59            }
60            Value::String(s) => HashKey::String(s.to_string()),
61            Value::Bytes(b) => {
62                // Use byte content for hashing
63                HashKey::String(format!("{b:?}"))
64            }
65            Value::Timestamp(t) => HashKey::Int64(t.as_micros()),
66            Value::List(items) => {
67                HashKey::Composite(items.iter().map(HashKey::from_value).collect())
68            }
69            Value::Map(map) => {
70                let mut keys: Vec<_> = map
71                    .iter()
72                    .map(|(k, v)| {
73                        HashKey::Composite(vec![
74                            HashKey::String(k.to_string()),
75                            HashKey::from_value(v),
76                        ])
77                    })
78                    .collect();
79                keys.sort_by(|a, b| format!("{a:?}").cmp(&format!("{b:?}")));
80                HashKey::Composite(keys)
81            }
82            Value::Vector(v) => {
83                // Hash vectors by converting each f32 to its bit representation
84                HashKey::Composite(
85                    v.iter()
86                        .map(|f| HashKey::Int64(f.to_bits() as i64))
87                        .collect(),
88                )
89            }
90        }
91    }
92
93    /// Creates a hash key from a column value at a given row.
94    pub fn from_column(column: &ValueVector, row: usize) -> Option<Self> {
95        column.get_value(row).map(|v| Self::from_value(&v))
96    }
97}
98
99/// Hash join operator.
100///
101/// Builds a hash table from the build side (right) and probes with the probe side (left).
102/// Efficient for equality joins on one or more columns.
103pub struct HashJoinOperator {
104    /// Left (probe) side operator.
105    probe_side: Box<dyn Operator>,
106    /// Right (build) side operator.
107    build_side: Box<dyn Operator>,
108    /// Column indices on the probe side for join keys.
109    probe_keys: Vec<usize>,
110    /// Column indices on the build side for join keys.
111    build_keys: Vec<usize>,
112    /// Join type.
113    join_type: JoinType,
114    /// Output schema (combined from both sides).
115    output_schema: Vec<LogicalType>,
116    /// Hash table: key -> list of (chunk_index, row_index).
117    hash_table: HashMap<HashKey, Vec<(usize, usize)>>,
118    /// Materialized build side chunks.
119    build_chunks: Vec<DataChunk>,
120    /// Whether the build phase is complete.
121    build_complete: bool,
122    /// Current probe chunk being processed.
123    current_probe_chunk: Option<DataChunk>,
124    /// Current row in the probe chunk.
125    current_probe_row: usize,
126    /// Current position in the hash table matches for the current probe row.
127    current_match_position: usize,
128    /// Current matches for the current probe row.
129    current_matches: Vec<(usize, usize)>,
130    /// For left/full outer joins: track which probe rows had matches.
131    probe_matched: Vec<bool>,
132    /// For right/full outer joins: track which build rows were matched.
133    build_matched: Vec<Vec<bool>>,
134    /// Whether we're in the emit unmatched phase (for outer joins).
135    emitting_unmatched: bool,
136    /// Current chunk index when emitting unmatched rows.
137    unmatched_chunk_idx: usize,
138    /// Current row index when emitting unmatched rows.
139    unmatched_row_idx: usize,
140}
141
142impl HashJoinOperator {
143    /// Creates a new hash join operator.
144    ///
145    /// # Arguments
146    /// * `probe_side` - Left side operator (will be probed).
147    /// * `build_side` - Right side operator (will build hash table).
148    /// * `probe_keys` - Column indices on probe side for join keys.
149    /// * `build_keys` - Column indices on build side for join keys.
150    /// * `join_type` - Type of join to perform.
151    /// * `output_schema` - Schema of the output (probe columns + build columns).
152    pub fn new(
153        probe_side: Box<dyn Operator>,
154        build_side: Box<dyn Operator>,
155        probe_keys: Vec<usize>,
156        build_keys: Vec<usize>,
157        join_type: JoinType,
158        output_schema: Vec<LogicalType>,
159    ) -> Self {
160        Self {
161            probe_side,
162            build_side,
163            probe_keys,
164            build_keys,
165            join_type,
166            output_schema,
167            hash_table: HashMap::new(),
168            build_chunks: Vec::new(),
169            build_complete: false,
170            current_probe_chunk: None,
171            current_probe_row: 0,
172            current_match_position: 0,
173            current_matches: Vec::new(),
174            probe_matched: Vec::new(),
175            build_matched: Vec::new(),
176            emitting_unmatched: false,
177            unmatched_chunk_idx: 0,
178            unmatched_row_idx: 0,
179        }
180    }
181
182    /// Builds the hash table from the build side.
183    fn build_hash_table(&mut self) -> Result<(), OperatorError> {
184        while let Some(chunk) = self.build_side.next()? {
185            let chunk_idx = self.build_chunks.len();
186
187            // Initialize match tracking for outer joins
188            if matches!(self.join_type, JoinType::Right | JoinType::Full) {
189                self.build_matched.push(vec![false; chunk.row_count()]);
190            }
191
192            // Add each row to the hash table
193            for row in chunk.selected_indices() {
194                let key = self.extract_key(&chunk, row, &self.build_keys)?;
195
196                // Skip null keys for inner/semi/anti joins
197                if matches!(key, HashKey::Null)
198                    && !matches!(
199                        self.join_type,
200                        JoinType::Left | JoinType::Right | JoinType::Full
201                    )
202                {
203                    continue;
204                }
205
206                self.hash_table
207                    .entry(key)
208                    .or_default()
209                    .push((chunk_idx, row));
210            }
211
212            self.build_chunks.push(chunk);
213        }
214
215        self.build_complete = true;
216        Ok(())
217    }
218
219    /// Extracts a hash key from a chunk row.
220    fn extract_key(
221        &self,
222        chunk: &DataChunk,
223        row: usize,
224        key_columns: &[usize],
225    ) -> Result<HashKey, OperatorError> {
226        if key_columns.len() == 1 {
227            let col = chunk.column(key_columns[0]).ok_or_else(|| {
228                OperatorError::ColumnNotFound(format!("column {}", key_columns[0]))
229            })?;
230            Ok(HashKey::from_column(col, row).unwrap_or(HashKey::Null))
231        } else {
232            let keys: Vec<HashKey> = key_columns
233                .iter()
234                .map(|&col_idx| {
235                    chunk
236                        .column(col_idx)
237                        .and_then(|col| HashKey::from_column(col, row))
238                        .unwrap_or(HashKey::Null)
239                })
240                .collect();
241            Ok(HashKey::Composite(keys))
242        }
243    }
244
245    /// Produces an output row from a probe row and build row.
246    fn produce_output_row(
247        &self,
248        builder: &mut DataChunkBuilder,
249        probe_chunk: &DataChunk,
250        probe_row: usize,
251        build_chunk: Option<&DataChunk>,
252        build_row: Option<usize>,
253    ) -> Result<(), OperatorError> {
254        let probe_col_count = probe_chunk.column_count();
255
256        // Copy probe side columns
257        for col_idx in 0..probe_col_count {
258            let src_col = probe_chunk
259                .column(col_idx)
260                .ok_or_else(|| OperatorError::ColumnNotFound(format!("probe column {col_idx}")))?;
261            let dst_col = builder
262                .column_mut(col_idx)
263                .ok_or_else(|| OperatorError::ColumnNotFound(format!("output column {col_idx}")))?;
264
265            if let Some(value) = src_col.get_value(probe_row) {
266                dst_col.push_value(value);
267            } else {
268                dst_col.push_value(Value::Null);
269            }
270        }
271
272        // Copy build side columns
273        match (build_chunk, build_row) {
274            (Some(chunk), Some(row)) => {
275                for col_idx in 0..chunk.column_count() {
276                    let src_col = chunk.column(col_idx).ok_or_else(|| {
277                        OperatorError::ColumnNotFound(format!("build column {col_idx}"))
278                    })?;
279                    let dst_col =
280                        builder
281                            .column_mut(probe_col_count + col_idx)
282                            .ok_or_else(|| {
283                                OperatorError::ColumnNotFound(format!(
284                                    "output column {}",
285                                    probe_col_count + col_idx
286                                ))
287                            })?;
288
289                    if let Some(value) = src_col.get_value(row) {
290                        dst_col.push_value(value);
291                    } else {
292                        dst_col.push_value(Value::Null);
293                    }
294                }
295            }
296            _ => {
297                // Emit nulls for build side (left outer join case)
298                if !self.build_chunks.is_empty() {
299                    let build_col_count = self.build_chunks[0].column_count();
300                    for col_idx in 0..build_col_count {
301                        let dst_col =
302                            builder
303                                .column_mut(probe_col_count + col_idx)
304                                .ok_or_else(|| {
305                                    OperatorError::ColumnNotFound(format!(
306                                        "output column {}",
307                                        probe_col_count + col_idx
308                                    ))
309                                })?;
310                        dst_col.push_value(Value::Null);
311                    }
312                }
313            }
314        }
315
316        builder.advance_row();
317        Ok(())
318    }
319
320    /// Gets the next probe chunk.
321    fn get_next_probe_chunk(&mut self) -> Result<bool, OperatorError> {
322        let chunk = self.probe_side.next()?;
323        if let Some(ref c) = chunk {
324            // Initialize match tracking for outer joins
325            if matches!(self.join_type, JoinType::Left | JoinType::Full) {
326                self.probe_matched = vec![false; c.row_count()];
327            }
328        }
329        let has_chunk = chunk.is_some();
330        self.current_probe_chunk = chunk;
331        self.current_probe_row = 0;
332        Ok(has_chunk)
333    }
334
335    /// Emits unmatched build rows for right/full outer joins.
336    fn emit_unmatched_build(&mut self) -> OperatorResult {
337        if self.build_matched.is_empty() {
338            return Ok(None);
339        }
340
341        let mut builder = DataChunkBuilder::with_capacity(&self.output_schema, 2048);
342
343        // Determine probe column count from schema or first probe chunk
344        let probe_col_count = if !self.build_chunks.is_empty() {
345            self.output_schema.len() - self.build_chunks[0].column_count()
346        } else {
347            0
348        };
349
350        while self.unmatched_chunk_idx < self.build_chunks.len() {
351            let chunk = &self.build_chunks[self.unmatched_chunk_idx];
352            let matched = &self.build_matched[self.unmatched_chunk_idx];
353
354            while self.unmatched_row_idx < matched.len() {
355                if !matched[self.unmatched_row_idx] {
356                    // This row was not matched - emit with nulls on probe side
357
358                    // Emit nulls for probe side
359                    for col_idx in 0..probe_col_count {
360                        if let Some(dst_col) = builder.column_mut(col_idx) {
361                            dst_col.push_value(Value::Null);
362                        }
363                    }
364
365                    // Copy build side values
366                    for col_idx in 0..chunk.column_count() {
367                        if let (Some(src_col), Some(dst_col)) = (
368                            chunk.column(col_idx),
369                            builder.column_mut(probe_col_count + col_idx),
370                        ) {
371                            if let Some(value) = src_col.get_value(self.unmatched_row_idx) {
372                                dst_col.push_value(value);
373                            } else {
374                                dst_col.push_value(Value::Null);
375                            }
376                        }
377                    }
378
379                    builder.advance_row();
380
381                    if builder.is_full() {
382                        self.unmatched_row_idx += 1;
383                        return Ok(Some(builder.finish()));
384                    }
385                }
386
387                self.unmatched_row_idx += 1;
388            }
389
390            self.unmatched_chunk_idx += 1;
391            self.unmatched_row_idx = 0;
392        }
393
394        if builder.row_count() > 0 {
395            Ok(Some(builder.finish()))
396        } else {
397            Ok(None)
398        }
399    }
400}
401
402impl Operator for HashJoinOperator {
403    fn next(&mut self) -> OperatorResult {
404        // Phase 1: Build hash table
405        if !self.build_complete {
406            self.build_hash_table()?;
407        }
408
409        // Phase 3: Emit unmatched build rows (right/full outer join)
410        if self.emitting_unmatched {
411            return self.emit_unmatched_build();
412        }
413
414        // Phase 2: Probe
415        let mut builder = DataChunkBuilder::with_capacity(&self.output_schema, 2048);
416
417        loop {
418            // Get current probe chunk or fetch new one
419            if self.current_probe_chunk.is_none() && !self.get_next_probe_chunk()? {
420                // No more probe data
421                if matches!(self.join_type, JoinType::Right | JoinType::Full) {
422                    self.emitting_unmatched = true;
423                    return self.emit_unmatched_build();
424                }
425                return if builder.row_count() > 0 {
426                    Ok(Some(builder.finish()))
427                } else {
428                    Ok(None)
429                };
430            }
431
432            // Invariant: current_probe_chunk is Some here - the guard at line 396 either
433            // populates it via get_next_probe_chunk() or returns from the function
434            let probe_chunk = self
435                .current_probe_chunk
436                .as_ref()
437                .expect("probe chunk is Some: guard at line 396 ensures this");
438            let probe_rows: Vec<usize> = probe_chunk.selected_indices().collect();
439
440            while self.current_probe_row < probe_rows.len() {
441                let probe_row = probe_rows[self.current_probe_row];
442
443                // If we don't have current matches, look them up
444                if self.current_matches.is_empty() && self.current_match_position == 0 {
445                    let key = self.extract_key(probe_chunk, probe_row, &self.probe_keys)?;
446
447                    // Handle semi/anti joins differently
448                    match self.join_type {
449                        JoinType::Semi => {
450                            if self.hash_table.contains_key(&key) {
451                                // Emit probe row only
452                                for col_idx in 0..probe_chunk.column_count() {
453                                    if let (Some(src_col), Some(dst_col)) =
454                                        (probe_chunk.column(col_idx), builder.column_mut(col_idx))
455                                        && let Some(value) = src_col.get_value(probe_row)
456                                    {
457                                        dst_col.push_value(value);
458                                    }
459                                }
460                                builder.advance_row();
461                            }
462                            self.current_probe_row += 1;
463                            continue;
464                        }
465                        JoinType::Anti => {
466                            if !self.hash_table.contains_key(&key) {
467                                // Emit probe row only
468                                for col_idx in 0..probe_chunk.column_count() {
469                                    if let (Some(src_col), Some(dst_col)) =
470                                        (probe_chunk.column(col_idx), builder.column_mut(col_idx))
471                                        && let Some(value) = src_col.get_value(probe_row)
472                                    {
473                                        dst_col.push_value(value);
474                                    }
475                                }
476                                builder.advance_row();
477                            }
478                            self.current_probe_row += 1;
479                            continue;
480                        }
481                        _ => {
482                            self.current_matches =
483                                self.hash_table.get(&key).cloned().unwrap_or_default();
484                        }
485                    }
486                }
487
488                // Process matches
489                if self.current_matches.is_empty() {
490                    // No matches - for left/full outer join, emit with nulls
491                    if matches!(self.join_type, JoinType::Left | JoinType::Full) {
492                        self.produce_output_row(&mut builder, probe_chunk, probe_row, None, None)?;
493                    }
494                    self.current_probe_row += 1;
495                    self.current_match_position = 0;
496                } else {
497                    // Process each match
498                    while self.current_match_position < self.current_matches.len() {
499                        let (build_chunk_idx, build_row) =
500                            self.current_matches[self.current_match_position];
501                        let build_chunk = &self.build_chunks[build_chunk_idx];
502
503                        // Mark as matched for outer joins
504                        if matches!(self.join_type, JoinType::Left | JoinType::Full)
505                            && probe_row < self.probe_matched.len()
506                        {
507                            self.probe_matched[probe_row] = true;
508                        }
509                        if matches!(self.join_type, JoinType::Right | JoinType::Full)
510                            && build_chunk_idx < self.build_matched.len()
511                            && build_row < self.build_matched[build_chunk_idx].len()
512                        {
513                            self.build_matched[build_chunk_idx][build_row] = true;
514                        }
515
516                        self.produce_output_row(
517                            &mut builder,
518                            probe_chunk,
519                            probe_row,
520                            Some(build_chunk),
521                            Some(build_row),
522                        )?;
523
524                        self.current_match_position += 1;
525
526                        if builder.is_full() {
527                            return Ok(Some(builder.finish()));
528                        }
529                    }
530
531                    // Done with this probe row
532                    self.current_probe_row += 1;
533                    self.current_matches.clear();
534                    self.current_match_position = 0;
535                }
536
537                if builder.is_full() {
538                    return Ok(Some(builder.finish()));
539                }
540            }
541
542            // Done with current probe chunk
543            self.current_probe_chunk = None;
544            self.current_probe_row = 0;
545
546            if builder.row_count() > 0 {
547                return Ok(Some(builder.finish()));
548            }
549        }
550    }
551
552    fn reset(&mut self) {
553        self.probe_side.reset();
554        self.build_side.reset();
555        self.hash_table.clear();
556        self.build_chunks.clear();
557        self.build_complete = false;
558        self.current_probe_chunk = None;
559        self.current_probe_row = 0;
560        self.current_match_position = 0;
561        self.current_matches.clear();
562        self.probe_matched.clear();
563        self.build_matched.clear();
564        self.emitting_unmatched = false;
565        self.unmatched_chunk_idx = 0;
566        self.unmatched_row_idx = 0;
567    }
568
569    fn name(&self) -> &'static str {
570        "HashJoin"
571    }
572}
573
574/// Nested loop join operator.
575///
576/// Performs a cartesian product of both sides, filtering by the join condition.
577/// Less efficient than hash join but supports any join condition.
578pub struct NestedLoopJoinOperator {
579    /// Left side operator.
580    left: Box<dyn Operator>,
581    /// Right side operator.
582    right: Box<dyn Operator>,
583    /// Join condition predicate (if any).
584    condition: Option<Box<dyn JoinCondition>>,
585    /// Join type.
586    join_type: JoinType,
587    /// Output schema.
588    output_schema: Vec<LogicalType>,
589    /// Materialized right side chunks.
590    right_chunks: Vec<DataChunk>,
591    /// Whether the right side is materialized.
592    right_materialized: bool,
593    /// Current left chunk.
594    current_left_chunk: Option<DataChunk>,
595    /// Current row in the left chunk.
596    current_left_row: usize,
597    /// Current chunk index in the right side.
598    current_right_chunk: usize,
599    /// Whether the current left row has been matched (for Left Join).
600    current_left_matched: bool,
601    /// Current row in the current right chunk.
602    current_right_row: usize,
603}
604
605/// Trait for join conditions.
606pub trait JoinCondition: Send + Sync {
607    /// Evaluates the condition for a pair of rows.
608    fn evaluate(
609        &self,
610        left_chunk: &DataChunk,
611        left_row: usize,
612        right_chunk: &DataChunk,
613        right_row: usize,
614    ) -> bool;
615}
616
617/// A simple equality condition for nested loop joins.
618pub struct EqualityCondition {
619    /// Column index on the left side.
620    left_column: usize,
621    /// Column index on the right side.
622    right_column: usize,
623}
624
625impl EqualityCondition {
626    /// Creates a new equality condition.
627    pub fn new(left_column: usize, right_column: usize) -> Self {
628        Self {
629            left_column,
630            right_column,
631        }
632    }
633}
634
635impl JoinCondition for EqualityCondition {
636    fn evaluate(
637        &self,
638        left_chunk: &DataChunk,
639        left_row: usize,
640        right_chunk: &DataChunk,
641        right_row: usize,
642    ) -> bool {
643        let left_val = left_chunk
644            .column(self.left_column)
645            .and_then(|c| c.get_value(left_row));
646        let right_val = right_chunk
647            .column(self.right_column)
648            .and_then(|c| c.get_value(right_row));
649
650        match (left_val, right_val) {
651            (Some(l), Some(r)) => l == r,
652            _ => false,
653        }
654    }
655}
656
657impl NestedLoopJoinOperator {
658    /// Creates a new nested loop join operator.
659    pub fn new(
660        left: Box<dyn Operator>,
661        right: Box<dyn Operator>,
662        condition: Option<Box<dyn JoinCondition>>,
663        join_type: JoinType,
664        output_schema: Vec<LogicalType>,
665    ) -> Self {
666        Self {
667            left,
668            right,
669            condition,
670            join_type,
671            output_schema,
672            right_chunks: Vec::new(),
673            right_materialized: false,
674            current_left_chunk: None,
675            current_left_row: 0,
676            current_right_chunk: 0,
677            current_right_row: 0,
678            current_left_matched: false,
679        }
680    }
681
682    /// Materializes the right side.
683    fn materialize_right(&mut self) -> Result<(), OperatorError> {
684        while let Some(chunk) = self.right.next()? {
685            self.right_chunks.push(chunk);
686        }
687        self.right_materialized = true;
688        Ok(())
689    }
690
691    /// Produces an output row.
692    fn produce_row(
693        &self,
694        builder: &mut DataChunkBuilder,
695        left_chunk: &DataChunk,
696        left_row: usize,
697        right_chunk: &DataChunk,
698        right_row: usize,
699    ) {
700        // Copy left columns
701        for col_idx in 0..left_chunk.column_count() {
702            if let (Some(src), Some(dst)) =
703                (left_chunk.column(col_idx), builder.column_mut(col_idx))
704            {
705                if let Some(val) = src.get_value(left_row) {
706                    dst.push_value(val);
707                } else {
708                    dst.push_value(Value::Null);
709                }
710            }
711        }
712
713        // Copy right columns
714        let left_col_count = left_chunk.column_count();
715        for col_idx in 0..right_chunk.column_count() {
716            if let (Some(src), Some(dst)) = (
717                right_chunk.column(col_idx),
718                builder.column_mut(left_col_count + col_idx),
719            ) {
720                if let Some(val) = src.get_value(right_row) {
721                    dst.push_value(val);
722                } else {
723                    dst.push_value(Value::Null);
724                }
725            }
726        }
727
728        builder.advance_row();
729    }
730
731    /// Produces an output row with NULLs for the right side (for unmatched left rows in Left Join).
732    fn produce_left_unmatched_row(
733        &self,
734        builder: &mut DataChunkBuilder,
735        left_chunk: &DataChunk,
736        left_row: usize,
737        right_col_count: usize,
738    ) {
739        // Copy left columns
740        for col_idx in 0..left_chunk.column_count() {
741            if let (Some(src), Some(dst)) =
742                (left_chunk.column(col_idx), builder.column_mut(col_idx))
743            {
744                if let Some(val) = src.get_value(left_row) {
745                    dst.push_value(val);
746                } else {
747                    dst.push_value(Value::Null);
748                }
749            }
750        }
751
752        // Fill right columns with NULLs
753        let left_col_count = left_chunk.column_count();
754        for col_idx in 0..right_col_count {
755            if let Some(dst) = builder.column_mut(left_col_count + col_idx) {
756                dst.push_value(Value::Null);
757            }
758        }
759
760        builder.advance_row();
761    }
762}
763
764impl Operator for NestedLoopJoinOperator {
765    fn next(&mut self) -> OperatorResult {
766        // Materialize right side
767        if !self.right_materialized {
768            self.materialize_right()?;
769        }
770
771        // If right side is empty and not a left outer join, return nothing
772        if self.right_chunks.is_empty() && !matches!(self.join_type, JoinType::Left) {
773            return Ok(None);
774        }
775
776        let mut builder = DataChunkBuilder::with_capacity(&self.output_schema, 2048);
777
778        loop {
779            // Get current left chunk
780            if self.current_left_chunk.is_none() {
781                self.current_left_chunk = self.left.next()?;
782                self.current_left_row = 0;
783                self.current_right_chunk = 0;
784                self.current_right_row = 0;
785
786                if self.current_left_chunk.is_none() {
787                    // No more left data
788                    return if builder.row_count() > 0 {
789                        Ok(Some(builder.finish()))
790                    } else {
791                        Ok(None)
792                    };
793                }
794            }
795
796            let left_chunk = self.current_left_chunk.as_ref().unwrap();
797            let left_rows: Vec<usize> = left_chunk.selected_indices().collect();
798
799            // Calculate right column count for potential unmatched rows
800            let right_col_count = if !self.right_chunks.is_empty() {
801                self.right_chunks[0].column_count()
802            } else {
803                // Infer from output schema
804                self.output_schema
805                    .len()
806                    .saturating_sub(left_chunk.column_count())
807            };
808
809            // Process current left row against all right rows
810            while self.current_left_row < left_rows.len() {
811                let left_row = left_rows[self.current_left_row];
812
813                // Reset match tracking for this left row
814                if self.current_right_chunk == 0 && self.current_right_row == 0 {
815                    self.current_left_matched = false;
816                }
817
818                // Cross join or inner/other join
819                while self.current_right_chunk < self.right_chunks.len() {
820                    let right_chunk = &self.right_chunks[self.current_right_chunk];
821                    let right_rows: Vec<usize> = right_chunk.selected_indices().collect();
822
823                    while self.current_right_row < right_rows.len() {
824                        let right_row = right_rows[self.current_right_row];
825
826                        // Check condition
827                        let matches = match &self.condition {
828                            Some(cond) => {
829                                cond.evaluate(left_chunk, left_row, right_chunk, right_row)
830                            }
831                            None => true, // Cross join
832                        };
833
834                        if matches {
835                            self.current_left_matched = true;
836                            self.produce_row(
837                                &mut builder,
838                                left_chunk,
839                                left_row,
840                                right_chunk,
841                                right_row,
842                            );
843
844                            if builder.is_full() {
845                                self.current_right_row += 1;
846                                return Ok(Some(builder.finish()));
847                            }
848                        }
849
850                        self.current_right_row += 1;
851                    }
852
853                    self.current_right_chunk += 1;
854                    self.current_right_row = 0;
855                }
856
857                // Done processing all right rows for this left row
858                // For Left Join, emit unmatched left row with NULLs
859                if matches!(self.join_type, JoinType::Left) && !self.current_left_matched {
860                    self.produce_left_unmatched_row(
861                        &mut builder,
862                        left_chunk,
863                        left_row,
864                        right_col_count,
865                    );
866
867                    if builder.is_full() {
868                        self.current_left_row += 1;
869                        self.current_right_chunk = 0;
870                        self.current_right_row = 0;
871                        return Ok(Some(builder.finish()));
872                    }
873                }
874
875                // Move to next left row
876                self.current_left_row += 1;
877                self.current_right_chunk = 0;
878                self.current_right_row = 0;
879            }
880
881            // Done with current left chunk
882            self.current_left_chunk = None;
883
884            if builder.row_count() > 0 {
885                return Ok(Some(builder.finish()));
886            }
887        }
888    }
889
890    fn reset(&mut self) {
891        self.left.reset();
892        self.right.reset();
893        self.right_chunks.clear();
894        self.right_materialized = false;
895        self.current_left_chunk = None;
896        self.current_left_row = 0;
897        self.current_right_chunk = 0;
898        self.current_right_row = 0;
899        self.current_left_matched = false;
900    }
901
902    fn name(&self) -> &'static str {
903        "NestedLoopJoin"
904    }
905}
906
907#[cfg(test)]
908mod tests {
909    use super::*;
910    use crate::execution::chunk::DataChunkBuilder;
911
912    /// Mock operator for testing.
913    struct MockOperator {
914        chunks: Vec<DataChunk>,
915        position: usize,
916    }
917
918    impl MockOperator {
919        fn new(chunks: Vec<DataChunk>) -> Self {
920            Self {
921                chunks,
922                position: 0,
923            }
924        }
925    }
926
927    impl Operator for MockOperator {
928        fn next(&mut self) -> OperatorResult {
929            if self.position < self.chunks.len() {
930                let chunk = std::mem::replace(&mut self.chunks[self.position], DataChunk::empty());
931                self.position += 1;
932                Ok(Some(chunk))
933            } else {
934                Ok(None)
935            }
936        }
937
938        fn reset(&mut self) {
939            self.position = 0;
940        }
941
942        fn name(&self) -> &'static str {
943            "Mock"
944        }
945    }
946
947    fn create_int_chunk(values: &[i64]) -> DataChunk {
948        let mut builder = DataChunkBuilder::new(&[LogicalType::Int64]);
949        for &v in values {
950            builder.column_mut(0).unwrap().push_int64(v);
951            builder.advance_row();
952        }
953        builder.finish()
954    }
955
956    #[test]
957    fn test_hash_join_inner() {
958        // Left: [1, 2, 3, 4]
959        // Right: [2, 3, 4, 5]
960        // Inner join on column 0 should produce: [2, 3, 4]
961
962        let left = MockOperator::new(vec![create_int_chunk(&[1, 2, 3, 4])]);
963        let right = MockOperator::new(vec![create_int_chunk(&[2, 3, 4, 5])]);
964
965        let output_schema = vec![LogicalType::Int64, LogicalType::Int64];
966        let mut join = HashJoinOperator::new(
967            Box::new(left),
968            Box::new(right),
969            vec![0],
970            vec![0],
971            JoinType::Inner,
972            output_schema,
973        );
974
975        let mut results = Vec::new();
976        while let Some(chunk) = join.next().unwrap() {
977            for row in chunk.selected_indices() {
978                let left_val = chunk.column(0).unwrap().get_int64(row).unwrap();
979                let right_val = chunk.column(1).unwrap().get_int64(row).unwrap();
980                results.push((left_val, right_val));
981            }
982        }
983
984        results.sort();
985        assert_eq!(results, vec![(2, 2), (3, 3), (4, 4)]);
986    }
987
988    #[test]
989    fn test_hash_join_left_outer() {
990        // Left: [1, 2, 3]
991        // Right: [2, 3]
992        // Left outer join should produce: [(1, null), (2, 2), (3, 3)]
993
994        let left = MockOperator::new(vec![create_int_chunk(&[1, 2, 3])]);
995        let right = MockOperator::new(vec![create_int_chunk(&[2, 3])]);
996
997        let output_schema = vec![LogicalType::Int64, LogicalType::Int64];
998        let mut join = HashJoinOperator::new(
999            Box::new(left),
1000            Box::new(right),
1001            vec![0],
1002            vec![0],
1003            JoinType::Left,
1004            output_schema,
1005        );
1006
1007        let mut results = Vec::new();
1008        while let Some(chunk) = join.next().unwrap() {
1009            for row in chunk.selected_indices() {
1010                let left_val = chunk.column(0).unwrap().get_int64(row).unwrap();
1011                let right_val = chunk.column(1).unwrap().get_int64(row);
1012                results.push((left_val, right_val));
1013            }
1014        }
1015
1016        results.sort_by_key(|(l, _)| *l);
1017        assert_eq!(results.len(), 3);
1018        assert_eq!(results[0], (1, None)); // No match
1019        assert_eq!(results[1], (2, Some(2)));
1020        assert_eq!(results[2], (3, Some(3)));
1021    }
1022
1023    #[test]
1024    fn test_nested_loop_cross_join() {
1025        // Left: [1, 2]
1026        // Right: [10, 20]
1027        // Cross join should produce: [(1,10), (1,20), (2,10), (2,20)]
1028
1029        let left = MockOperator::new(vec![create_int_chunk(&[1, 2])]);
1030        let right = MockOperator::new(vec![create_int_chunk(&[10, 20])]);
1031
1032        let output_schema = vec![LogicalType::Int64, LogicalType::Int64];
1033        let mut join = NestedLoopJoinOperator::new(
1034            Box::new(left),
1035            Box::new(right),
1036            None,
1037            JoinType::Cross,
1038            output_schema,
1039        );
1040
1041        let mut results = Vec::new();
1042        while let Some(chunk) = join.next().unwrap() {
1043            for row in chunk.selected_indices() {
1044                let left_val = chunk.column(0).unwrap().get_int64(row).unwrap();
1045                let right_val = chunk.column(1).unwrap().get_int64(row).unwrap();
1046                results.push((left_val, right_val));
1047            }
1048        }
1049
1050        results.sort();
1051        assert_eq!(results, vec![(1, 10), (1, 20), (2, 10), (2, 20)]);
1052    }
1053
1054    #[test]
1055    fn test_hash_join_semi() {
1056        // Left: [1, 2, 3, 4]
1057        // Right: [2, 4]
1058        // Semi join should produce: [2, 4] (only left rows that have matches)
1059
1060        let left = MockOperator::new(vec![create_int_chunk(&[1, 2, 3, 4])]);
1061        let right = MockOperator::new(vec![create_int_chunk(&[2, 4])]);
1062
1063        // Semi join only outputs probe (left) columns
1064        let output_schema = vec![LogicalType::Int64];
1065        let mut join = HashJoinOperator::new(
1066            Box::new(left),
1067            Box::new(right),
1068            vec![0],
1069            vec![0],
1070            JoinType::Semi,
1071            output_schema,
1072        );
1073
1074        let mut results = Vec::new();
1075        while let Some(chunk) = join.next().unwrap() {
1076            for row in chunk.selected_indices() {
1077                let val = chunk.column(0).unwrap().get_int64(row).unwrap();
1078                results.push(val);
1079            }
1080        }
1081
1082        results.sort();
1083        assert_eq!(results, vec![2, 4]);
1084    }
1085
1086    #[test]
1087    fn test_hash_join_anti() {
1088        // Left: [1, 2, 3, 4]
1089        // Right: [2, 4]
1090        // Anti join should produce: [1, 3] (left rows with no matches)
1091
1092        let left = MockOperator::new(vec![create_int_chunk(&[1, 2, 3, 4])]);
1093        let right = MockOperator::new(vec![create_int_chunk(&[2, 4])]);
1094
1095        let output_schema = vec![LogicalType::Int64];
1096        let mut join = HashJoinOperator::new(
1097            Box::new(left),
1098            Box::new(right),
1099            vec![0],
1100            vec![0],
1101            JoinType::Anti,
1102            output_schema,
1103        );
1104
1105        let mut results = Vec::new();
1106        while let Some(chunk) = join.next().unwrap() {
1107            for row in chunk.selected_indices() {
1108                let val = chunk.column(0).unwrap().get_int64(row).unwrap();
1109                results.push(val);
1110            }
1111        }
1112
1113        results.sort();
1114        assert_eq!(results, vec![1, 3]);
1115    }
1116}