akar_processor/physical/write_ops/
optionalextend.rs1use 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
11pub struct PhysicalOptionalExtend {
25 pub rel_table_name: String,
27 pub rel_table_id: u64,
29 pub rel_var: String,
32 pub src_node_var: String,
34 pub dst_node_var: String,
36 pub direction: EdgeDirection,
38 pub table_catalog: Arc<TableCatalog>,
40}
41
42fn 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
62fn 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 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 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 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 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 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}