akar_processor/processor/
union_helpers.rs

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

akar_processor/processor/
union_helpers.rs
