akar_processor/processor/
union_helpers.rs

1use crate::physical::common::value_hash;
2use crate::processor::chunk_helpers::{extract_all_rows_from_chunks, rows_to_columns};
3use akar_common::error::ProcessorError;
4use akar_common::types::Value;
5use akar_common::vector::DataChunk;
6use akar_planner::logical_operator::LogicalOperator;
7use std::collections::HashMap;
8
9pub fn flatten_union_child(op: &LogicalOperator) -> Vec<LogicalOperator> {
10    match op {
11        LogicalOperator::Projection(p) if p.expressions.is_empty() => p.children.clone(),
12        other => vec![other.clone()],
13    }
14}
15
16pub fn merge_union_chunks(
17    left: Vec<DataChunk>,
18    right: Vec<DataChunk>,
19    all: bool,
20) -> Result<Vec<DataChunk>, ProcessorError> {
21    if left.is_empty() {
22        return Ok(right);
23    }
24    if right.is_empty() {
25        return Ok(left);
26    }
27
28    let num_fields = left[0].num_fields();
29    for chunk in &right {
30        if chunk.num_fields() != num_fields {
31            return Err(format!(
32                "UNION column count mismatch: left has {num_fields} columns, right has {} columns",
33                chunk.num_fields()
34            )
35            .into());
36        }
37    }
38
39    let mut left_rows = extract_all_rows_from_chunks(&left);
40    let right_rows = extract_all_rows_from_chunks(&right);
41    left_rows.extend(right_rows);
42
43    let mut deduped: Vec<Vec<Value>> = Vec::with_capacity(left_rows.len());
44    if !all {
45        for row in &left_rows {
46            if !deduped.contains(row) {
47                deduped.push(row.clone());
48            }
49        }
50    } else {
51        deduped = left_rows;
52    }
53
54    if deduped.is_empty() {
55        return Ok(vec![DataChunk::new(vec![], vec![])]);
56    }
57
58    let (fields, field_types) = rows_to_columns(&deduped);
59    let final_size = deduped.len();
60    let field_names = left.first().map(|c| c.field_names.clone()).unwrap_or_default();
61
62    Ok(vec![DataChunk {
63        fields,
64        field_types,
65        size: final_size,
66        field_names,
67        sel_vector: None,
68    }])
69}
70
71pub fn merge_optional_chunks(left: Vec<DataChunk>, right: Vec<DataChunk>) -> Result<Vec<DataChunk>, ProcessorError> {
72    if left.is_empty() {
73        return Ok(left);
74    }
75    if right.is_empty() {
76        return Ok(left);
77    }
78
79    let left_rows = extract_all_rows_from_chunks(&left);
80    let right_rows = extract_all_rows_from_chunks(&right);
81    if left_rows.is_empty() {
82        return Ok(left);
83    }
84
85    let left_names: Vec<String> = left.first().map(|c| c.field_names.clone()).unwrap_or_default();
86    let right_names: Vec<String> = right.first().map(|c| c.field_names.clone()).unwrap_or_default();
87    let num_right_cols = right_rows.first().map(|r| r.len()).unwrap_or(right_names.len());
88
89    if right_rows.is_empty() {
90        // Left-outer join with no right matches: keep every left row and pad
91        // the right columns with NULLs, preserving the right side's schema so
92        // that RETURN can still resolve the optional columns (e.g. `m.id`).
93        if num_right_cols == 0 {
94            return Ok(left);
95        }
96        let mut combined: Vec<Vec<Value>> = Vec::with_capacity(left_rows.len());
97        for lrow in &left_rows {
98            let mut row = Vec::with_capacity(left_names.len() + num_right_cols);
99            row.extend_from_slice(lrow);
100            row.extend(std::iter::repeat_n(Value::Null, num_right_cols));
101            combined.push(row);
102        }
103        let (fields, field_types) = rows_to_columns(&combined);
104        let mut field_names = left_names;
105        field_names.extend(right_names);
106        return Ok(vec![DataChunk {
107            fields,
108            field_types,
109            size: combined.len(),
110            field_names,
111            sel_vector: None,
112        }]);
113    }
114
115    let num_left_cols = left_rows.first().map(|r| r.len()).unwrap_or(0);
116
117    // OPTIONAL MATCH is a left-outer join on the variables shared between the
118    // two sides. When a shared column exists (identical field name on both
119    // sides) join on it; otherwise the two sides are independent and every
120    // left row pairs with every right row (cross product). Positional merging
121    // would silently drop right rows and mispair the i-th left row with the
122    // i-th right row whenever cardinalities differ.
123    let join_col_left = left_names.iter().position(|n| right_names.contains(n));
124    let join_col_right = join_col_left.and_then(|li| right_names.iter().position(|n| *n == left_names[li]));
125
126    let mut combined: Vec<Vec<Value>> = Vec::new();
127    match (join_col_left, join_col_right) {
128        (Some(li), Some(ri)) => {
129            // Build a hash index over the right side's join column.
130            let mut hash: HashMap<u64, Vec<usize>> = HashMap::new();
131            for (i, row) in right_rows.iter().enumerate() {
132                if let Some(val) = row.get(ri) {
133                    if !matches!(val, Value::Null) {
134                        hash.entry(value_hash(val)).or_default().push(i);
135                    }
136                }
137            }
138            for lrow in &left_rows {
139                let key = lrow.get(li);
140                let mut emitted = false;
141                if let Some(key) = key {
142                    if !matches!(key, Value::Null) {
143                        if let Some(bucket) = hash.get(&value_hash(key)) {
144                            for &i in bucket {
145                                if &right_rows[i][ri] == key {
146                                    let mut row = Vec::with_capacity(num_left_cols + num_right_cols);
147                                    row.extend_from_slice(lrow);
148                                    row.extend_from_slice(&right_rows[i]);
149                                    combined.push(row);
150                                    emitted = true;
151                                }
152                            }
153                        }
154                    }
155                }
156                if !emitted {
157                    let mut row = Vec::with_capacity(num_left_cols + num_right_cols);
158                    row.extend_from_slice(lrow);
159                    row.extend(std::iter::repeat_n(Value::Null, num_right_cols));
160                    combined.push(row);
161                }
162            }
163        }
164        _ => {
165            // Cross product left-outer join: each left row pairs with all
166            // right rows.
167            for lrow in &left_rows {
168                for rrow in &right_rows {
169                    let mut row = Vec::with_capacity(num_left_cols + num_right_cols);
170                    row.extend_from_slice(lrow);
171                    row.extend_from_slice(rrow);
172                    combined.push(row);
173                }
174            }
175        }
176    }
177
178    if combined.is_empty() {
179        return Ok(vec![]);
180    }
181
182    let (fields, field_types) = rows_to_columns(&combined);
183    let size = combined.len();
184    let mut field_names = left_names;
185    field_names.extend(right_names);
186    Ok(vec![DataChunk {
187        fields,
188        field_types,
189        size,
190        field_names,
191        sel_vector: None,
192    }])
193}
194
195#[cfg(test)]
196mod tests {
197    use super::*;
198    use akar_common::types::PhysicalTypeID;
199    use akar_common::vector::ValueVector;
200
201    fn make_chunk(cols: &[Vec<Value>], names: &[&str]) -> DataChunk {
202        let mut fields = Vec::with_capacity(cols.len());
203        let mut types = Vec::with_capacity(cols.len());
204        for col in cols {
205            let first = col.iter().find(|v| !matches!(v, Value::Null)).unwrap_or(&Value::Int64(0));
206            let ptype = match first {
207                Value::Int64(_) => PhysicalTypeID::Int64,
208                Value::String(_) => PhysicalTypeID::String,
209                _ => PhysicalTypeID::Int64,
210            };
211            let mut v = ValueVector::new(ptype, col.len().max(1));
212            for (i, val) in col.iter().enumerate() {
213                let _ = v.set_value(i, val);
214            }
215            v.resize(col.len());
216            fields.push(akar_common::arrow_vector::ArrowVector::from_legacy(&v).array);
217            types.push(ptype);
218        }
219        let mut chunk = DataChunk::new(fields, types);
220        chunk.field_names = names.iter().map(|s| s.to_string()).collect();
221        chunk
222    }
223
224    #[test]
225    fn test_merge_optional_empty_right_returns_left() {
226        let left = make_chunk(&[vec![Value::Int64(1), Value::Int64(2)]], &["a.id"]);
227        let merged = merge_optional_chunks(vec![left.clone()], vec![]).unwrap();
228        assert_eq!(merged.len(), 1);
229        assert_eq!(merged[0].size, 2);
230    }
231
232    #[test]
233    fn test_merge_optional_zero_row_right_pads_nulls() {
234        // Right side yields a 0-row chunk but still carries its schema. The
235        // merge must keep both left rows and pad `m.id` with NULLs so RETURN
236        // resolves the optional column instead of falling back to a wrong one.
237        let left = make_chunk(&[vec![Value::Int64(1), Value::Int64(2)]], &["n.id"]);
238        let right = make_chunk(&[Vec::<Value>::new()], &["m.id"]);
239        let merged = merge_optional_chunks(vec![left], vec![right]).unwrap();
240        assert_eq!(merged[0].size, 2, "every left row survives");
241        assert_eq!(merged[0].field_names, vec!["n.id", "m.id"]);
242        assert!(merged[0].is_null(1, 0), "m.id must be null");
243        assert!(merged[0].is_null(1, 1), "m.id must be null");
244    }
245
246    #[test]
247    fn test_merge_optional_cross_product_no_dropped_rows() {
248        // Different cardinalities (2 left, 2 right) — no shared column, so the
249        // merge is a cross product. Every left row is kept and no right row is
250        // dropped.
251        let left = make_chunk(
252            &[vec![Value::Int64(1), Value::Int64(2)], vec![Value::Int64(10), Value::Int64(20)]],
253            &["a.id", "a.x"],
254        );
255        let right = make_chunk(&[vec![Value::Int64(5), Value::Int64(6)]], &["b.id"]);
256        let merged = merge_optional_chunks(vec![left], vec![right]).unwrap();
257        assert_eq!(merged[0].size, 4, "expected 2 left rows x 2 right rows");
258        assert_eq!(merged[0].field_names, vec!["a.id", "a.x", "b.id"]);
259    }
260
261    #[test]
262    fn test_merge_optional_joins_on_shared_column() {
263        // Shared column `a.id` on both sides → proper left-outer join.
264        let left = make_chunk(
265            &[vec![Value::Int64(1), Value::Int64(2)], vec![Value::String("Alice".into()), Value::String("Bob".into())]],
266            &["a.id", "a.name"],
267        );
268        let right = make_chunk(
269            &[vec![Value::Int64(1), Value::Int64(3)], vec![Value::String("Rex".into()), Value::String("Tom".into())]],
270            &["a.id", "a.pet"],
271        );
272        let merged = merge_optional_chunks(vec![left], vec![right]).unwrap();
273        assert_eq!(merged[0].size, 2, "Alice joins Rex; Bob has no match -> nulls");
274        let pet_col = merged[0].get_string(3, 0).map(str::to_string);
275        assert_eq!(pet_col.as_deref(), Some("Rex"));
276        assert!(merged[0].is_null(3, 1), "Bob's pet must be null");
277    }
278}
akar_processor/processor/union_helpers.rs

akar_processor/processor/
union_helpers.rs
