Skip to main content

alopex_sql/executor/query/
join.rs

1use crate::executor::evaluator::EvalContext;
2use crate::executor::{ExecutorError, Result, Row};
3use crate::planner::logical_plan::JoinType;
4use crate::planner::typed_expr::{TypedExpr, TypedExprKind};
5use crate::storage::SqlValue;
6use alopex_core::sql::join::{
7    RowLike, combine_rows as core_combine_rows, hash_join as core_hash_join,
8    nested_loop_join as core_nested_loop_join, pad_left as core_pad_left,
9    pad_right as core_pad_right,
10};
11
12impl RowLike for Row {
13    type Value = SqlValue;
14
15    fn values(&self) -> &[Self::Value] {
16        &self.values
17    }
18}
19
20/// Execute JOIN operation.
21pub fn execute_join(
22    left_rows: Vec<Row>,
23    right_rows: Vec<Row>,
24    join_type: JoinType,
25    condition: Option<&TypedExpr>,
26) -> Result<Vec<Row>> {
27    let left_width = left_rows.first().map_or(0, Row::len);
28    let right_width = right_rows.first().map_or(0, Row::len);
29    execute_join_with_widths(
30        left_rows,
31        right_rows,
32        join_type,
33        condition,
34        left_width,
35        right_width,
36    )
37}
38
39pub(crate) fn execute_join_with_widths(
40    left_rows: Vec<Row>,
41    right_rows: Vec<Row>,
42    join_type: JoinType,
43    condition: Option<&TypedExpr>,
44    left_width: usize,
45    right_width: usize,
46) -> Result<Vec<Row>> {
47    if matches!(join_type, JoinType::Cross) || condition.is_none() {
48        return nested_loop_join_with_widths(
49            &left_rows,
50            &right_rows,
51            condition,
52            join_type,
53            left_width,
54            right_width,
55        );
56    }
57
58    if let Some((left_key, right_key)) = condition.and_then(|expr| equi_join_keys(expr, left_width))
59    {
60        return hash_join_with_widths(
61            &left_rows,
62            &right_rows,
63            left_key,
64            right_key,
65            join_type,
66            left_width,
67            right_width,
68        );
69    }
70
71    nested_loop_join_with_widths(
72        &left_rows,
73        &right_rows,
74        condition,
75        join_type,
76        left_width,
77        right_width,
78    )
79}
80
81/// Nested loop join implementation.
82pub fn nested_loop_join(
83    left: &[Row],
84    right: &[Row],
85    condition: &TypedExpr,
86    join_type: JoinType,
87) -> Result<Vec<Row>> {
88    let left_width = left.first().map_or(0, Row::len);
89    let right_width = right.first().map_or(0, Row::len);
90    nested_loop_join_with_widths(
91        left,
92        right,
93        Some(condition),
94        join_type,
95        left_width,
96        right_width,
97    )
98}
99
100fn nested_loop_join_with_widths(
101    left: &[Row],
102    right: &[Row],
103    condition: Option<&TypedExpr>,
104    join_type: JoinType,
105    left_width: usize,
106    right_width: usize,
107) -> Result<Vec<Row>> {
108    let pairs = core_nested_loop_join(
109        left,
110        right,
111        matches!(join_type, JoinType::Left | JoinType::Full),
112        matches!(join_type, JoinType::Right | JoinType::Full),
113        |left_row, right_row| {
114            let joined = core_combine_rows(left_row, right_row, 0, Row::new);
115            condition_matches(condition, &joined)
116        },
117    )?;
118    Ok(materialize_join_pairs(pairs, left_width, right_width))
119}
120
121/// Hash join implementation for equi-joins.
122pub fn hash_join(
123    left: &[Row],
124    right: &[Row],
125    left_key: usize,
126    right_key: usize,
127    join_type: JoinType,
128) -> Result<Vec<Row>> {
129    let left_width = left.first().map_or(0, Row::len);
130    let right_width = right.first().map_or(0, Row::len);
131    hash_join_with_widths(
132        left,
133        right,
134        left_key,
135        right_key,
136        join_type,
137        left_width,
138        right_width,
139    )
140}
141
142fn hash_join_with_widths(
143    left: &[Row],
144    right: &[Row],
145    left_key: usize,
146    right_key: usize,
147    join_type: JoinType,
148    left_width: usize,
149    right_width: usize,
150) -> Result<Vec<Row>> {
151    let pairs = core_hash_join(
152        left,
153        right,
154        matches!(join_type, JoinType::Left | JoinType::Full),
155        matches!(join_type, JoinType::Right | JoinType::Full),
156        |row| {
157            row.get(left_key).map(hash_key).ok_or({
158                ExecutorError::Evaluation(crate::executor::EvaluationError::InvalidColumnRef {
159                    index: left_key,
160                })
161            })
162        },
163        |row| {
164            row.get(right_key).map(hash_key).ok_or({
165                ExecutorError::Evaluation(crate::executor::EvaluationError::InvalidColumnRef {
166                    index: right_key,
167                })
168            })
169        },
170    )?;
171    Ok(materialize_join_pairs(pairs, left_width, right_width))
172}
173
174fn materialize_join_pairs(
175    pairs: Vec<(Option<Row>, Option<Row>)>,
176    left_width: usize,
177    right_width: usize,
178) -> Vec<Row> {
179    pairs
180        .into_iter()
181        .enumerate()
182        .map(|(row_id, (left, right))| match (left, right) {
183            (Some(left), Some(right)) => core_combine_rows(&left, &right, row_id as u64, Row::new),
184            (Some(left), None) => core_pad_right(
185                &left,
186                right_width,
187                row_id as u64,
188                || SqlValue::Null,
189                Row::new,
190            ),
191            (None, Some(right)) => core_pad_left(
192                &right,
193                left_width,
194                row_id as u64,
195                || SqlValue::Null,
196                Row::new,
197            ),
198            (None, None) => unreachable!("core join never emits an empty join pair"),
199        })
200        .collect()
201}
202
203fn condition_matches(condition: Option<&TypedExpr>, row: &Row) -> Result<bool> {
204    let Some(condition) = condition else {
205        return Ok(true);
206    };
207    let ctx = EvalContext::new(&row.values);
208    match crate::executor::evaluator::evaluate(condition, &ctx)? {
209        SqlValue::Boolean(true) => Ok(true),
210        _ => Ok(false),
211    }
212}
213
214fn equi_join_keys(condition: &TypedExpr, left_width: usize) -> Option<(usize, usize)> {
215    let TypedExprKind::BinaryOp { left, op, right } = &condition.kind else {
216        return None;
217    };
218    if !matches!(op, crate::ast::expr::BinaryOp::Eq) {
219        return None;
220    }
221    let (
222        TypedExprKind::ColumnRef {
223            column_index: l, ..
224        },
225        TypedExprKind::ColumnRef {
226            column_index: r, ..
227        },
228    ) = (&left.kind, &right.kind)
229    else {
230        return None;
231    };
232    match (*l < left_width, *r < left_width) {
233        (true, false) => Some((*l, *r - left_width)),
234        (false, true) => Some((*r, *l - left_width)),
235        _ => None,
236    }
237}
238
239fn hash_key(value: &SqlValue) -> String {
240    format!("{value:?}")
241}