akar_processor/physical/write_ops/
packedextend.rs

1//! Auto-extracted from physical_operator.rs
2use crate::physical::types::{OperatorResult, PhysicalOperatorExec};
3use akar_common::types::Value;
4use akar_common::vector::{DataChunk, ValueVector};
5use akar_storage::table::TableCatalog;
6use std::sync::Arc;
7
8// ==================== PackedExtend ====================
9
10/// Physical operator for multi-rel extend, producing flattened output rows.
11///
12/// Extends from a source node using a `CsrIndex` or adjacency list and
13/// produces one output row per relationship, duplicating the source node
14/// properties for each destination neighbor. This normalized format is
15/// what downstream operators (HashJoin, Filter, etc.) expect.
16pub struct PhysicalPackedExtend {
17    pub rel_table_name: String,
18    pub rel_table_id: u64,
19    pub bound_node_var: String,
20    pub direction: akar_parser::ast::EdgeDirection,
21    pub dst_node_var: String,
22    pub table_catalog: Arc<TableCatalog>,
23}
24
25impl PhysicalOperatorExec for PhysicalPackedExtend {
26    fn operator_type(&self) -> &str {
27        "packed_extend"
28    }
29
30    fn execute(&self, input: Vec<DataChunk>) -> OperatorResult {
31        if input.is_empty() || input.iter().all(|c| c.size == 0) {
32            return Ok(input);
33        }
34
35        let rel_table = self
36            .table_catalog
37            .get_rel_table_by_name(&self.rel_table_name)
38            .ok_or_else(|| format!("Rel table {} not found", self.rel_table_name))?;
39
40        let mut output_chunks = Vec::new();
41
42        for chunk in input {
43            if chunk.size == 0 {
44                continue;
45            }
46
47            // Find bound node column index
48            let bound_idx = chunk
49                .field_names
50                .iter()
51                .position(|name| name == &self.bound_node_var)
52                .unwrap_or(0);
53
54            // --- Pass 1: collect neighbor lists and estimate total output size ---
55            let mut per_src_neighbors: Vec<Vec<u64>> = Vec::with_capacity(chunk.size);
56            let mut total_output_rows: usize = 0;
57
58            for i in 0..chunk.size {
59                if chunk.fields[bound_idx].is_null(i) {
60                    per_src_neighbors.push(Vec::new());
61                    continue;
62                }
63
64                let src_id = match chunk.get_value(bound_idx, i) {
65                    Some(Value::Int64(id)) => id as u64,
66                    Some(Value::UInt64(id)) => id,
67                    _ => {
68                        per_src_neighbors.push(Vec::new());
69                        continue;
70                    }
71                };
72
73                let neighbors = self.fetch_neighbors(&rel_table, src_id);
74                total_output_rows += neighbors.len();
75                per_src_neighbors.push(neighbors);
76            }
77
78            if total_output_rows == 0 {
79                continue;
80            }
81
82            // --- Pass 2: build flat output vectors with pre-allocated capacity ---
83            let num_input_cols = chunk.fields.len();
84            let mut out_fields: Vec<arrow::array::ArrayRef> = Vec::with_capacity(num_input_cols + 1);
85            let mut out_field_types: Vec<akar_common::types::PhysicalTypeID> = Vec::with_capacity(num_input_cols + 1);
86
87            // Duplicate each input column for every output row
88            for col_idx in 0..num_input_cols {
89                let phys_type = chunk.field_types[col_idx];
90                // Arrow builder path for strings too (255-byte ValueVector cap, P63).
91                if matches!(
92                    phys_type,
93                    akar_common::types::PhysicalTypeID::String
94                        | akar_common::types::PhysicalTypeID::List
95                        | akar_common::types::PhysicalTypeID::Array
96                        | akar_common::types::PhysicalTypeID::Struct
97                ) {
98                    let mut flat: Vec<Value> = Vec::with_capacity(total_output_rows);
99                    let mut out_pos = 0;
100                    for (src_row, neighbors) in per_src_neighbors.iter().enumerate() {
101                        if neighbors.is_empty() {
102                            continue;
103                        }
104                        let is_null = chunk.fields[col_idx].is_null(src_row);
105                        let val = chunk.get_value(col_idx, src_row);
106                        for _ in neighbors {
107                            if is_null {
108                                flat.push(Value::Null);
109                            } else {
110                                flat.push(val.clone().unwrap_or(Value::Null));
111                            }
112                            out_pos += 1;
113                        }
114                    }
115                    let arr = crate::expression_evaluator::build_arrow_from_values(&flat, phys_type, out_pos)
116                        .map_err(|e| e.to_string())?;
117                    out_fields.push(arr.array);
118                    out_field_types.push(phys_type);
119                } else {
120                    let mut v = ValueVector::new(phys_type, total_output_rows);
121                    v.resize(total_output_rows);
122
123                    let mut out_pos = 0;
124                    for (src_row, neighbors) in per_src_neighbors.iter().enumerate() {
125                        if neighbors.is_empty() {
126                            continue;
127                        }
128                        // Copy source row value once, then duplicate for each neighbor
129                        let is_null = chunk.fields[col_idx].is_null(src_row);
130                        let val = chunk.get_value(col_idx, src_row);
131                        for _ in neighbors {
132                            if is_null {
133                                v.set_null(out_pos, true);
134                            } else if let Some(ref val) = val {
135                                crate::physical::common::store_value_in_vector(&mut v, out_pos, val)?;
136                            }
137                            out_pos += 1;
138                        }
139                    }
140                    out_fields.push(akar_common::arrow_vector::ArrowVector::from_legacy(&v).array);
141                    out_field_types.push(phys_type);
142                }
143            }
144
145            // Destination column: flat Int64 of neighbor node IDs
146            let mut dst_field = ValueVector::new(akar_common::types::PhysicalTypeID::Int64, total_output_rows);
147            dst_field.resize(total_output_rows);
148
149            let mut out_pos = 0;
150            for neighbors in &per_src_neighbors {
151                for &dst_id in neighbors {
152                    crate::physical::common::store_value_in_vector(
153                        &mut dst_field,
154                        out_pos,
155                        &Value::Int64(dst_id as i64),
156                    )?;
157                    out_pos += 1;
158                }
159            }
160            out_fields.push(akar_common::arrow_vector::ArrowVector::from_legacy(&dst_field).array);
161            out_field_types.push(akar_common::types::PhysicalTypeID::Int64);
162
163            let mut new_names = chunk.field_names.clone();
164            new_names.push(self.dst_node_var.clone());
165
166            output_chunks.push(DataChunk {
167                fields: out_fields,
168                field_types: out_field_types,
169                size: total_output_rows,
170                field_names: new_names,
171                sel_vector: None,
172            });
173        }
174
175        if output_chunks.is_empty() {
176            Ok(vec![DataChunk {
177                fields: vec![],
178                field_types: vec![],
179                size: 0,
180                field_names: vec![],
181                sel_vector: None,
182            }])
183        } else {
184            Ok(output_chunks)
185        }
186    }
187}
188
189impl PhysicalPackedExtend {
190    fn fetch_neighbors(&self, rel_table: &akar_storage::table::RelTable, src_id: u64) -> Vec<u64> {
191        if let Some(csr) = &rel_table.csr_index {
192            let is_fwd = matches!(self.direction, akar_parser::ast::EdgeDirection::LeftToRight);
193            csr.get_neighbors(src_id, is_fwd).unwrap_or_default()
194        } else {
195            match self.direction {
196                akar_parser::ast::EdgeDirection::LeftToRight => rel_table
197                    .get_outgoing_edges(src_id)
198                    .into_iter()
199                    .map(|(dst, _)| dst)
200                    .collect(),
201                akar_parser::ast::EdgeDirection::RightToLeft => rel_table
202                    .get_incoming_edges(src_id)
203                    .into_iter()
204                    .map(|(dst, _)| dst)
205                    .collect(),
206                akar_parser::ast::EdgeDirection::Both => {
207                    let mut n: Vec<u64> = rel_table
208                        .get_outgoing_edges(src_id)
209                        .into_iter()
210                        .map(|(dst, _)| dst)
211                        .collect();
212                    n.extend(rel_table.get_incoming_edges(src_id).into_iter().map(|(dst, _)| dst));
213                    n
214                }
215            }
216        }
217    }
218}
akar_processor/physical/write_ops/packedextend.rs

akar_processor/physical/write_ops/
packedextend.rs
