akar_processor/processor/mapper/
map_join.rs

1use super::ExecutionContext;
2use crate::physical_operator::*;
3use akar_common::error::ProcessorError;
4use akar_common::types::{PhysicalTypeID, Value};
5use akar_common::vector::DataChunk;
6use akar_parser::ast::Expression;
7use akar_planner::logical_operator::LogicalOperator;
8
9use crate::processor::join_helpers::{JoinKeyBinding, derive_join_bindings};
10use crate::processor::union_helpers::{flatten_union_child, merge_optional_chunks};
11
12pub fn map_and_execute_join(
13    op: &LogicalOperator,
14    current_input: Vec<DataChunk>,
15    ctx: &mut ExecutionContext,
16) -> Result<Vec<DataChunk>, ProcessorError> {
17    match op {
18        LogicalOperator::HashJoin(h) => {
19            let left_ops = flatten_union_child(&h.build_side);
20            let right_ops = flatten_union_child(&h.probe_side);
21
22            let build_chunks = ctx.execute_children(&left_ops)?;
23            let probe_chunks = ctx.execute_children(&right_ops)?;
24
25            let (build_orig, probe_orig) = (field_count(&build_chunks), field_count(&probe_chunks));
26            let (build_chunks, build_cols, build_appended, probe_chunks, probe_cols, probe_appended) =
27                prepare_join_sides(&h.join_keys, build_chunks, probe_chunks)?;
28
29            let join = PhysicalHashJoin::new(build_cols, probe_cols);
30            let result = join.execute_binary(&build_chunks, &probe_chunks)?;
31
32            Ok(strip_join_synthetic_columns(
33                result,
34                build_orig,
35                build_appended,
36                probe_orig,
37                probe_appended,
38                true,
39            ))
40        }
41        LogicalOperator::SemiJoin(s) => {
42            let left_ops = flatten_union_child(&s.left);
43            let right_ops = flatten_union_child(&s.right);
44
45            let build_chunks = ctx.execute_children(&left_ops)?;
46            let probe_chunks = ctx.execute_children(&right_ops)?;
47
48            let (_build_orig, probe_orig) = (field_count(&build_chunks), field_count(&probe_chunks));
49            let (build_chunks, build_cols, _build_appended, probe_chunks, probe_cols, probe_appended) =
50                prepare_join_sides(&s.join_keys, build_chunks, probe_chunks)?;
51
52            let semi = PhysicalSemiJoin {
53                build_columns: build_cols,
54                probe_columns: probe_cols,
55            };
56            let result = semi.execute_binary(&build_chunks, &probe_chunks)?;
57            Ok(strip_join_synthetic_columns(
58                result,
59                0,
60                0,
61                probe_orig,
62                probe_appended,
63                false,
64            ))
65        }
66        LogicalOperator::AntiJoin(a) => {
67            let left_ops = flatten_union_child(&a.left);
68            let right_ops = flatten_union_child(&a.right);
69
70            let build_chunks = ctx.execute_children(&left_ops)?;
71            let probe_chunks = ctx.execute_children(&right_ops)?;
72
73            let (_build_orig, probe_orig) = (field_count(&build_chunks), field_count(&probe_chunks));
74            let (build_chunks, build_cols, _build_appended, probe_chunks, probe_cols, probe_appended) =
75                prepare_join_sides(&a.join_keys, build_chunks, probe_chunks)?;
76
77            let anti = PhysicalAntiJoin {
78                build_columns: build_cols,
79                probe_columns: probe_cols,
80            };
81            let result = anti.execute_binary(&build_chunks, &probe_chunks)?;
82            Ok(strip_join_synthetic_columns(
83                result,
84                0,
85                0,
86                probe_orig,
87                probe_appended,
88                false,
89            ))
90        }
91        LogicalOperator::Intersect(ic) => {
92            // Build side is a (possibly nested) Union of per-pattern pipelines,
93            // each executed independently so the intersect has one hash table per
94            // pattern. Probe side is the shared-node scan.
95            let build_sides = collect_union_sides(&ic.left);
96            let probe_ops = flatten_union_child(&ic.right);
97
98            let build_chunk_sides: Vec<Vec<DataChunk>> = build_sides
99                .iter()
100                .map(|ops| ctx.execute_children(ops))
101                .collect::<Result<_, _>>()?;
102            let probe_chunks = ctx.execute_children(&probe_ops)?;
103
104            let (probe_key_col, build_key_col) =
105                resolve_intersect_key_cols(&ic.build_key_exprs, &probe_chunks, &build_chunk_sides);
106
107            let intersect = PhysicalIntersect {
108                num_build_sides: build_chunk_sides.len() as u32,
109                probe_key_col,
110                build_key_col,
111            };
112            let result = intersect.execute_sides(&build_chunk_sides, &probe_chunks)?;
113            Ok(result)
114        }
115        LogicalOperator::CrossProduct(cp) => {
116            let left_ops = flatten_union_child(&cp.left);
117            let right_ops = flatten_union_child(&cp.right);
118
119            let left_chunks = ctx.execute_children(&left_ops)?;
120            let right_chunks = ctx.execute_children(&right_ops)?;
121
122            let cross = PhysicalCrossProduct;
123            let result = cross.execute_binary(&left_chunks, &right_chunks)?;
124
125            Ok(result)
126        }
127        LogicalOperator::OptionalMatch(om) => {
128            // Execute left (required) subtree
129            let left_ops = flatten_union_child(&om.left);
130            let left_result = ctx.execute_children(&left_ops)?;
131
132            // Execute right (optional) subtree
133            let right_ops = flatten_union_child(&om.right);
134            let right_result = ctx.execute_children(&right_ops)?;
135
136            // Combine: use flattened row-level merge
137            let merged = merge_optional_chunks(left_result, right_result)?;
138            Ok(merged)
139        }
140        LogicalOperator::RecursiveExtend(re) => {
141            let scan = PhysicalRecursiveExtend {
142                source_table_id: re.source_table_id,
143                rel_table_ids: re.rel_table_ids.clone(),
144                lower_bound: re.lower_bound,
145                upper_bound: re.upper_bound,
146                direction: re.direction,
147                semantic: re.semantic,
148                table_catalog: ctx.table_catalog.clone(),
149                weight_property: re.weight_property.clone(),
150                cost_output_name: re.cost_output_name.clone(),
151            };
152            let result = scan.execute(current_input)?;
153            Ok(result)
154        }
155        _ => Err(format!("Not a join operator: {:?}", op).into()),
156    }
157}
158
159fn field_count(chunks: &[DataChunk]) -> usize {
160    chunks.first().map(|c| c.fields.len()).unwrap_or(0)
161}
162
163/// Split join bindings into per-side (column, map_key) lists.
164fn split_bindings(bindings: &[JoinKeyBinding]) -> (Vec<u32>, Vec<Option<String>>, Vec<u32>, Vec<Option<String>>) {
165    let mut build_cols = Vec::new();
166    let mut build_keys = Vec::new();
167    let mut probe_cols = Vec::new();
168    let mut probe_keys = Vec::new();
169    for b in bindings {
170        build_cols.push(b.build_col);
171        build_keys.push(b.build_map_key.clone());
172        probe_cols.push(b.probe_col);
173        probe_keys.push(b.probe_map_key.clone());
174    }
175    (build_cols, build_keys, probe_cols, probe_keys)
176}
177
178/// Resolve join key columns and materialize map/struct key extraction (P53.26).
179/// Returns the (possibly extended) build/probe chunks with the resolved column
180/// indices, plus how many synthetic columns were appended to each side.
181fn prepare_join_sides(
182    join_keys: &[Expression],
183    build_chunks: Vec<DataChunk>,
184    probe_chunks: Vec<DataChunk>,
185) -> Result<(Vec<DataChunk>, Vec<u32>, usize, Vec<DataChunk>, Vec<u32>, usize), ProcessorError> {
186    let bindings = derive_join_bindings(join_keys, &build_chunks, &probe_chunks);
187    let (build_cols, build_keys, probe_cols, probe_keys) = split_bindings(&bindings);
188    let (build_chunks, build_cols, build_appended) = materialize_map_keys(&build_chunks, &build_cols, &build_keys)?;
189    let (probe_chunks, probe_cols, probe_appended) = materialize_map_keys(&probe_chunks, &probe_cols, &probe_keys)?;
190    Ok((
191        build_chunks,
192        build_cols,
193        build_appended,
194        probe_chunks,
195        probe_cols,
196        probe_appended,
197    ))
198}
199
200/// Append synthetic columns holding map/struct key values for every binding
201/// that needs extraction, and point the join column at them. Complex values are
202/// preserved because `build_arrow_from_values` emits Arrow arrays directly
203/// (unlike `store_value_in_vector`, which drops them to NULL).
204fn materialize_map_keys(
205    chunks: &[DataChunk],
206    cols: &[u32],
207    keys: &[Option<String>],
208) -> Result<(Vec<DataChunk>, Vec<u32>, usize), ProcessorError> {
209    let mut out: Vec<DataChunk> = chunks.to_vec();
210    let mut new_cols: Vec<u32> = cols.to_vec();
211    let base_count = out.first().map(|c| c.fields.len()).unwrap_or(0);
212    let mut appended = 0usize;
213    // Dedupe repeated (column, key) extractions so each synthetic column is
214    // added only once and shared across join keys.
215    let mut done: std::collections::HashMap<(u32, String), u32> = std::collections::HashMap::new();
216
217    for (i, key) in keys.iter().enumerate() {
218        let Some(key_name) = key else { continue };
219        let col = cols[i];
220        if let Some(&idx) = done.get(&(col, key_name.clone())) {
221            new_cols[i] = idx;
222            continue;
223        }
224        let new_idx = (base_count + appended) as u32;
225        let mut ok = false;
226        for chunk in out.iter_mut() {
227            if col as usize >= chunk.fields.len() {
228                continue;
229            }
230            let extracted: Vec<Value> = (0..chunk.size)
231                .map(|row| {
232                    crate::expression_evaluator::map_property_value(
233                        &chunk.get_value(col as usize, row).unwrap_or(Value::Null),
234                        key_name,
235                    )
236                })
237                .collect();
238            let t = extracted
239                .iter()
240                .find(|v| !matches!(v, Value::Null))
241                .map(|v| v.physical_type())
242                .unwrap_or(PhysicalTypeID::Int64);
243            let arr = crate::expression_evaluator::build_arrow_from_values(&extracted, t, chunk.size)
244                .map_err(|e| e.to_string())?;
245            chunk.fields.push(arr.array);
246            chunk.field_types.push(arr.physical_type);
247            chunk.field_names.push(format!("__join_extract_{}_{}", col, key_name));
248            ok = true;
249        }
250        if ok {
251            appended += 1;
252        }
253        done.insert((col, key_name.clone()), new_idx);
254        new_cols[i] = new_idx;
255    }
256    Ok((out, new_cols, appended))
257}
258
259/// Remove synthetic extraction columns from join output chunks.
260///
261/// For hash joins the output is `[build columns..., probe columns...]`; for
262/// semi/anti joins it is `[probe columns...]`. Synthetic columns are the last
263/// `build_appended` / `probe_appended` fields of their side.
264fn strip_join_synthetic_columns(
265    result: Vec<DataChunk>,
266    build_orig: usize,
267    build_appended: usize,
268    probe_orig: usize,
269    probe_appended: usize,
270    output_has_build: bool,
271) -> Vec<DataChunk> {
272    let mut to_remove: Vec<usize> = Vec::new();
273    if output_has_build {
274        to_remove.extend(build_orig..build_orig + build_appended);
275        let probe_base = build_orig + build_appended + probe_orig;
276        to_remove.extend(probe_base..probe_base + probe_appended);
277    } else {
278        to_remove.extend(probe_orig..probe_orig + probe_appended);
279    }
280    if to_remove.is_empty() {
281        return result;
282    }
283    result
284        .into_iter()
285        .map(|mut chunk| {
286            for &idx in to_remove.iter().rev() {
287                if idx < chunk.fields.len() {
288                    chunk.fields.remove(idx);
289                    if idx < chunk.field_types.len() {
290                        chunk.field_types.remove(idx);
291                    }
292                    if idx < chunk.field_names.len() {
293                        chunk.field_names.remove(idx);
294                    }
295                }
296            }
297            chunk
298        })
299        .collect()
300}
301
302/// Flatten a (possibly nested) `Union` subtree into a list of independent
303/// operator pipelines — one per WCOJ build side.
304fn collect_union_sides(op: &LogicalOperator) -> Vec<Vec<LogicalOperator>> {
305    match op {
306        LogicalOperator::Union(u) => {
307            let mut sides = collect_union_sides(&u.left);
308            sides.extend(collect_union_sides(&u.right));
309            sides
310        }
311        other => vec![flatten_union_child(other)],
312    }
313}
314
315/// Resolve the shared-node key column index on the probe and build sides.
316///
317/// The key is derived from the first build key expression (a reference to the
318/// shared variable, e.g. `a`), resolved against `field_names` as `a.id`.
319fn resolve_intersect_key_cols(
320    build_key_exprs: &[Expression],
321    probe_chunks: &[DataChunk],
322    build_sides: &[Vec<DataChunk>],
323) -> (u32, u32) {
324    let var = build_key_exprs.first().and_then(|e| match e {
325        Expression::Variable(v) => Some(v.clone()),
326        Expression::PropertyAccess(obj, _) => {
327            if let Expression::Variable(v) = &**obj {
328                Some(v.clone())
329            } else {
330                None
331            }
332        }
333        _ => None,
334    });
335
336    let probe_names: Vec<&str> = probe_chunks
337        .first()
338        .map(|c| c.field_names.iter().map(|s| s.as_str()).collect())
339        .unwrap_or_default();
340    let build_names: Vec<&str> = build_sides
341        .first()
342        .and_then(|s| s.first())
343        .map(|c| c.field_names.iter().map(|s| s.as_str()).collect())
344        .unwrap_or_default();
345
346    let mut probe_col = 0u32;
347    let mut build_col = 0u32;
348    if let Some(var) = var {
349        let candidates = [format!("{var}._id"), format!("{var}.id"), var];
350        for c in &candidates {
351            if let Some(idx) = probe_names.iter().position(|n| n == c) {
352                probe_col = idx as u32;
353                break;
354            }
355        }
356        for c in &candidates {
357            if let Some(idx) = build_names.iter().position(|n| n == c) {
358                build_col = idx as u32;
359                break;
360            }
361        }
362    }
363    (probe_col, build_col)
364}
akar_processor/processor/mapper/map_join.rs

akar_processor/processor/mapper/
map_join.rs
