Skip to main content

akar_processor/processor/
join_helpers.rs

1use akar_common::vector::DataChunk;
2use akar_parser::ast::Expression;
3
4pub fn derive_join_column_indices(
5    join_keys: &[Expression],
6    build_chunks: &[DataChunk],
7    probe_chunks: &[DataChunk],
8) -> (Vec<u32>, Vec<u32>) {
9    let build_names: Vec<&str> = build_chunks
10        .first()
11        .map(|c| c.field_names.iter().map(|s| s.as_str()).collect())
12        .unwrap_or_default();
13    let probe_names: Vec<&str> = probe_chunks
14        .first()
15        .map(|c| c.field_names.iter().map(|s| s.as_str()).collect())
16        .unwrap_or_default();
17
18    let mut build_cols: Vec<u32> = Vec::new();
19    let mut probe_cols: Vec<u32> = Vec::new();
20
21    for key in join_keys {
22        if let Expression::BinaryOp(akar_parser::ast::BinaryOp::Equal, left, right) = key {
23            let left_prop = extract_join_prop(left);
24            let right_prop = extract_join_prop(right);
25
26            if let (Some(lp), Some(rp)) = (left_prop, right_prop) {
27                let build_idx = build_names
28                    .iter()
29                    .position(|&n| n == lp)
30                    .or_else(|| build_names.iter().position(|&n| n == rp))
31                    .unwrap_or(0) as u32;
32                let probe_idx = probe_names
33                    .iter()
34                    .position(|&n| n == rp)
35                    .or_else(|| probe_names.iter().position(|&n| n == lp))
36                    .unwrap_or(0) as u32;
37                build_cols.push(build_idx);
38                probe_cols.push(probe_idx);
39            }
40        }
41    }
42
43    if build_cols.is_empty() {
44        (vec![0], vec![0])
45    } else {
46        (build_cols, probe_cols)
47    }
48}
49
50fn extract_join_prop(expr: &Expression) -> Option<String> {
51    match expr {
52        Expression::PropertyAccess(obj, prop) => {
53            if let Expression::Variable(var) = &**obj {
54                Some(format!("{}.{}", var, prop))
55            } else {
56                Some(prop.clone())
57            }
58        }
59        Expression::Variable(name) => Some(name.clone()),
60        _ => None,
61    }
62}