Skip to main content

akar_processor/physical/write_ops/
optionalextend.rs

1use crate::physical::common::store_value_in_vector;
2use crate::physical::scan_filter::PhysicalScan;
3use akar_common::error::ProcessorError;
4use akar_common::types::{PhysicalTypeID, Value};
5use akar_common::vector::{DataChunk, ValueVector};
6use akar_parser::ast::EdgeDirection;
7use akar_storage::table::TableCatalog;
8use std::collections::HashMap;
9use std::sync::Arc;
10
11// ==================== OptionalExtend ====================
12
13/// Physical operator for `OPTIONAL MATCH` over an already-bound pair of node
14/// variables, e.g. `OPTIONAL MATCH (a)-[existing:Connected]-(b)` where both
15/// `a` and `b` come from the mandatory side (P53.25).
16///
17/// For each input row the relationship-table adjacency is probed for an edge
18/// between the source and destination node ids. When an edge exists, its
19/// property columns (plus an internal `{rel_var}._id` holding the edge index)
20/// are emitted; when no edge exists, those columns are NULL-padded and exactly
21/// one row is still produced (outer-join semantics).
22///
23/// Output layout: `[input_fields | rel_properties | {rel_var}._id]`.
24pub struct PhysicalOptionalExtend {
25    /// Name of the relationship table to probe.
26    pub rel_table_name: String,
27    /// ID of the relationship table.
28    pub rel_table_id: u64,
29    /// Variable name of the relationship (e.g., "existing"); prefix for the
30    /// emitted edge property columns.
31    pub rel_var: String,
32    /// Variable name of the bound source node (e.g., "a").
33    pub src_node_var: String,
34    /// Variable name of the bound destination node (e.g., "b").
35    pub dst_node_var: String,
36    /// Direction of the probe (forward, backward, or both).
37    pub direction: EdgeDirection,
38    /// Table catalog for data access.
39    pub table_catalog: Arc<TableCatalog>,
40}
41
42/// Resolve the internal node-id column (`{var}._id`, falling back to the bare
43/// variable or the primary key column) in the input chunk.
44fn find_node_id_col(chunk: &DataChunk, var: &str) -> Result<usize, ProcessorError> {
45    let name_id = format!("{}.{}", var, "_id");
46    let name_pk = format!("{}.{}", var, "id");
47    let idx = chunk
48        .field_names
49        .iter()
50        .position(|n| n == &name_id)
51        .or_else(|| chunk.field_names.iter().position(|n| n == var))
52        .or_else(|| chunk.field_names.iter().position(|n| n == &name_pk));
53    idx.ok_or_else(|| {
54        format!(
55            "Node variable {} not found in OptionalExtend input. Available fields: {:?}",
56            var, chunk.field_names
57        )
58        .into()
59    })
60}
61
62/// Find an edge index between `src` and `dst` in the forward/reverse adjacency
63/// maps, honoring the probe direction. Returns the first matching edge index.
64fn probe_edge(
65    fwd_adj: &HashMap<u64, Vec<(u64, usize)>>,
66    rev_adj: &HashMap<u64, Vec<(u64, usize)>>,
67    src: u64,
68    dst: u64,
69    direction: &EdgeDirection,
70) -> Option<usize> {
71    match direction {
72        EdgeDirection::LeftToRight => fwd_adj
73            .get(&src)
74            .and_then(|e| e.iter().find(|(o, _)| *o == dst).map(|(_, i)| *i)),
75        EdgeDirection::RightToLeft => rev_adj
76            .get(&src)
77            .and_then(|e| e.iter().find(|(o, _)| *o == dst).map(|(_, i)| *i)),
78        EdgeDirection::Both => fwd_adj
79            .get(&src)
80            .and_then(|e| e.iter().find(|(o, _)| *o == dst).map(|(_, i)| *i))
81            .or_else(|| {
82                rev_adj
83                    .get(&src)
84                    .and_then(|e| e.iter().find(|(o, _)| *o == dst).map(|(_, i)| *i))
85            }),
86    }
87}
88
89impl PhysicalOptionalExtend {
90    pub fn execute(&self, input: Vec<DataChunk>) -> Result<Vec<DataChunk>, ProcessorError> {
91        if input.is_empty() {
92            return Ok(input);
93        }
94
95        // Collect rel table data upfront (owned)
96        let (fwd_adj, rev_adj, rel_props, rel_cols) = {
97            let rel_table = self
98                .table_catalog
99                .get_rel_table_by_name(&self.rel_table_name)
100                .ok_or_else(|| format!("Rel table {} not found", self.rel_table_name))?;
101            (
102                rel_table.fwd_adj.clone(),
103                rel_table.rev_adj.clone(),
104                rel_table.properties.clone(),
105                rel_table.columns.clone(),
106            )
107        };
108
109        let num_rel_cols = rel_cols.len();
110        let rel_prefix = if self.rel_var.is_empty() {
111            self.rel_table_name.clone()
112        } else {
113            self.rel_var.clone()
114        };
115        let rel_field_names: Vec<String> = rel_cols.iter().map(|c| format!("{}.{}", rel_prefix, c.name)).collect();
116
117        let mut output = Vec::with_capacity(input.len());
118
119        for chunk in input {
120            if chunk.size == 0 {
121                output.push(chunk);
122                continue;
123            }
124
125            let src_idx = find_node_id_col(&chunk, &self.src_node_var)?;
126            let dst_idx = find_node_id_col(&chunk, &self.dst_node_var)?;
127
128            let num_input_fields = chunk.fields.len();
129            let num_out_cols = num_input_fields + num_rel_cols + 1;
130            let mut out_data: Vec<Vec<Value>> = vec![Vec::with_capacity(chunk.size); num_out_cols];
131
132            for i in 0..chunk.size {
133                for col in 0..num_input_fields {
134                    let val = chunk.get_value(col, i).unwrap_or(Value::Null);
135                    out_data[col].push(val);
136                }
137
138                // Probe the adjacency for an edge between src and dst.
139                let edge_idx = match (chunk.get_value(src_idx, i), chunk.get_value(dst_idx, i)) {
140                    (Some(Value::Int64(s)), Some(Value::Int64(d))) => {
141                        probe_edge(&fwd_adj, &rev_adj, s as u64, d as u64, &self.direction)
142                    }
143                    _ => None,
144                };
145
146                for col in 0..num_rel_cols {
147                    let val = match edge_idx {
148                        Some(ei) => rel_props
149                            .get(col)
150                            .and_then(|c| c.get(ei))
151                            .cloned()
152                            .unwrap_or(Value::Null),
153                        None => Value::Null,
154                    };
155                    out_data[num_input_fields + col].push(val);
156                }
157                // Internal edge index — non-NULL iff an edge exists, so a
158                // zero-property rel table still answers `{rel_var} IS NULL`.
159                out_data[num_input_fields + num_rel_cols].push(match edge_idx {
160                    Some(ei) => Value::Int64(ei as i64),
161                    None => Value::Null,
162                });
163            }
164
165            // Build output columns.
166            let mut fields = Vec::with_capacity(num_out_cols);
167            let mut field_types = Vec::with_capacity(num_out_cols);
168            let mut field_names = Vec::with_capacity(num_out_cols);
169
170            for col in 0..num_input_fields {
171                let phys_type = chunk.field_types[col];
172                let mut v = ValueVector::new(phys_type, chunk.size);
173                v.resize(chunk.size);
174                for row in 0..chunk.size {
175                    store_value_in_vector(&mut v, row, &out_data[col][row])?;
176                }
177                fields.push(v);
178                field_types.push(phys_type);
179                field_names.push(if col < chunk.field_names.len() {
180                    chunk.field_names[col].clone()
181                } else {
182                    format!("field_{}", col)
183                });
184            }
185            for col in 0..num_rel_cols {
186                let phys_type = if col < rel_cols.len() {
187                    PhysicalScan::logical_to_physical(&rel_cols[col].logical_type)
188                } else {
189                    PhysicalTypeID::Int64
190                };
191                let mut v = ValueVector::new(phys_type, chunk.size);
192                v.resize(chunk.size);
193                for row in 0..chunk.size {
194                    store_value_in_vector(&mut v, row, &out_data[num_input_fields + col][row])?;
195                }
196                fields.push(v);
197                field_types.push(phys_type);
198                field_names.push(rel_field_names[col].clone());
199            }
200            // Internal edge index column (`{rel_var}._id`).
201            let mut id_v = ValueVector::new(PhysicalTypeID::Int64, chunk.size);
202            id_v.resize(chunk.size);
203            for row in 0..chunk.size {
204                store_value_in_vector(&mut id_v, row, &out_data[num_input_fields + num_rel_cols][row])?;
205            }
206            fields.push(id_v);
207            field_types.push(PhysicalTypeID::Int64);
208            field_names.push(format!("{}.{}", rel_prefix, "_id"));
209
210            let arrow_fields = fields
211                .iter()
212                .map(|v| akar_common::arrow_vector::ArrowVector::from_legacy(v).array)
213                .collect::<Vec<_>>();
214            let arrow_field_types = fields.iter().map(|v| v.physical_type()).collect::<Vec<_>>();
215            output.push(DataChunk {
216                fields: arrow_fields,
217                field_types: arrow_field_types,
218                size: chunk.size,
219                field_names,
220                sel_vector: None,
221            });
222        }
223
224        Ok(output)
225    }
226}