Skip to main content

akar_processor/physical/write_ops/
copyfrom.rs

1//! Auto-extracted from physical_operator.rs
2use crate::physical::types::{OperatorResult, PhysicalOperatorExec};
3use akar_common::types::{PhysicalTypeID, Value};
4use akar_common::vector::{DataChunk, ValueVector};
5use akar_storage::table::{ColumnDefinition, TableCatalog};
6use akar_transaction::UndoRecord;
7use std::path::Path;
8use std::sync::{Arc, Mutex};
9
10// ==================== CopyFrom ====================
11
12/// Physical operator for COPY FROM — loads data from CSV/Parquet files into a table.
13///
14/// Detects file type from extension, calls the appropriate reader,
15/// and inserts rows into the target table via the `TableCatalog`.
16pub struct PhysicalCopyFrom {
17    pub table_name: String,
18    pub table_id: u64,
19    pub file_path: String,
20    pub columns: Vec<ColumnDefinition>,
21    pub options: std::collections::HashMap<String, String>,
22    pub table_catalog: Arc<TableCatalog>,
23    pub vfs: Arc<akar_common::file_system::VirtualFileSystemRegistry>,
24    /// Active transaction id (P52.18).
25    pub txn_id: Option<u64>,
26    /// Undo sink for rollback records (P52.18).
27    pub undo_sink: Option<Arc<Mutex<Vec<UndoRecord>>>>,
28}
29
30impl PhysicalOperatorExec for PhysicalCopyFrom {
31    fn operator_type(&self) -> &str {
32        "copy_from"
33    }
34
35    fn execute(&self, _input: Vec<DataChunk>) -> OperatorResult {
36        let path = Path::new(&self.file_path);
37
38        // 1. Detect file type from extension
39        let ext = path
40            .extension()
41            .and_then(|e| e.to_str())
42            .map(|e| e.to_lowercase())
43            .unwrap_or_default();
44
45        // 2. Build config and convert column schema.
46        //    For rel tables the COPY file carries two leading [from, to] columns
47        //    (node PK values) ahead of the user properties. Synthesize those two
48        //    columns (typed as the src/dst node PK types) so readers validate
49        //    `columns.len() + 2` and the insert branch can resolve PKs -> offsets.
50        let mut catalog_cols: Vec<akar_catalog::CatalogColumn> = self
51            .columns
52            .iter()
53            .map(|c| akar_catalog::CatalogColumn {
54                name: c.name.clone(),
55                logical_type: c.logical_type,
56                is_primary_key: c.is_primary_key,
57                compression: akar_common::enums::CompressionType::Uncompressed,
58                default_value: None,
59            })
60            .collect();
61        {
62            let rel_meta = self.table_catalog.get_rel_table_by_name(&self.table_name);
63            if let Some(rel) = &rel_meta {
64                let src_pk_type = self
65                    .table_catalog
66                    .get_node_table(rel.src_table_id)
67                    .and_then(|n| n.columns.get(n.primary_key_column).map(|c| c.logical_type))
68                    .unwrap_or(akar_common::types::LogicalTypeID::Int64);
69                let dst_pk_type = self
70                    .table_catalog
71                    .get_node_table(rel.dst_table_id)
72                    .and_then(|n| n.columns.get(n.primary_key_column).map(|c| c.logical_type))
73                    .unwrap_or(akar_common::types::LogicalTypeID::Int64);
74                let synthetic =
75                    |name: &str, logical_type: akar_common::types::LogicalTypeID| akar_catalog::CatalogColumn {
76                        name: name.to_string(),
77                        logical_type,
78                        is_primary_key: false,
79                        compression: akar_common::enums::CompressionType::Uncompressed,
80                        default_value: None,
81                    };
82                catalog_cols.insert(0, synthetic("from", src_pk_type));
83                catalog_cols.insert(1, synthetic("to", dst_pk_type));
84            }
85        }
86
87        // 3. Read the file
88        let rows = match ext.as_str() {
89            "csv" | "tsv" => {
90                let mut config = akar_storage::csv_reader::CsvReaderConfig::from_options(&self.options);
91                if ext == "tsv" && !self.options.contains_key("DELIM") && !self.options.contains_key("delim") {
92                    config.delimiter = b'\t';
93                }
94
95                akar_storage::csv_reader::read_csv(&self.file_path, &self.vfs, &catalog_cols, &config)
96                    .map_err(|e| format!("CSV read error: {e}"))?
97            }
98            #[cfg(feature = "parquet")]
99            "parquet" => akar_storage::parquet_reader::read_parquet(&self.file_path, &self.vfs, &catalog_cols)
100                .map_err(|e| format!("Parquet read error: {e}"))?,
101            #[cfg(not(feature = "parquet"))]
102            "parquet" => return Err("Parquet support not enabled (feature 'parquet' in akar-storage)".into()),
103            _ => {
104                return Err(format!("Unsupported file type: .{ext} (supported: .csv, .tsv, .parquet)").into());
105            }
106        };
107
108        // 4. Insert rows into the table using batch insert
109        let num_rows = rows.len();
110        if num_rows == 0 {
111            let mut v = ValueVector::new(PhysicalTypeID::Int64, 1);
112            v.resize(1);
113            v.set_i64(0, 0);
114            return Ok(vec![DataChunk::new(
115                vec![akar_common::arrow_vector::ArrowVector::from_legacy(&v).array],
116                vec![akar_common::types::PhysicalTypeID::Int64],
117            )]);
118        }
119
120        if let Some(mut table) = self.table_catalog.get_node_table_by_name_mut(&self.table_name) {
121            let start = table.num_rows;
122            let count = table
123                .insert_rows_batch_with_txn(&rows, self.txn_id)
124                .map_err(|e| format!("Batch insert error: {e}"))?;
125            if let Some(sink) = self.undo_sink.as_ref()
126                && let Ok(mut u) = sink.lock()
127            {
128                for row in start..start + count {
129                    u.push(UndoRecord::insert(self.table_id, row));
130                }
131            }
132            tracing::info!(
133                "COPY FROM: batch-inserted {count} rows into node table '{}'",
134                self.table_name
135            );
136        } else if let Some(mut table) = self.table_catalog.get_rel_table_by_name_mut(&self.table_name) {
137            // Rel COPY files carry [from, to, ...props] where from/to are node PK
138            // values. Resolve them to internal node offsets via the src/dst node
139            // tables' PK index (mirrors C++ IndexLookupInfo).
140            let src_node = self.table_catalog.get_node_table(table.src_table_id);
141            let dst_node = self.table_catalog.get_node_table(table.dst_table_id);
142            let mut rels: Vec<(u64, u64, Vec<Value>)> = Vec::with_capacity(rows.len());
143            for row in &rows {
144                let from = src_node.as_ref().and_then(|n| n.lookup_by_pk(&row[0])).ok_or_else(|| {
145                    format!(
146                        "COPY rel: source node with PK {:?} not found in table '{}'",
147                        row[0], self.table_name
148                    )
149                })?;
150                let to = dst_node.as_ref().and_then(|n| n.lookup_by_pk(&row[1])).ok_or_else(|| {
151                    format!(
152                        "COPY rel: destination node with PK {:?} not found in table '{}'",
153                        row[1], self.table_name
154                    )
155                })?;
156                rels.push((from, to, row[2..].to_vec()));
157            }
158            let start = table.edges.len();
159            let count = table
160                .insert_rels_batch(&rels)
161                .map_err(|e| format!("Batch insert rel error: {e}"))?;
162            if let Some(sink) = self.undo_sink.as_ref()
163                && let Ok(mut u) = sink.lock()
164            {
165                for idx in start..start + count as usize {
166                    u.push(UndoRecord::insert(self.table_id, idx as u64));
167                }
168            }
169            tracing::info!(
170                "COPY FROM: batch-inserted {count} rows into rel table '{}'",
171                self.table_name
172            );
173        } else {
174            return Err(format!("Table '{}' not found in storage catalog", self.table_name).into());
175        }
176
177        // Return success chunk with row count
178        let mut v = ValueVector::new(PhysicalTypeID::Int64, 1);
179        v.resize(1);
180        v.set_i64(0, num_rows as i64);
181        Ok(vec![DataChunk::new(
182            vec![akar_common::arrow_vector::ArrowVector::from_legacy(&v).array],
183            vec![akar_common::types::PhysicalTypeID::Int64],
184        )])
185    }
186}
187
188/// Physical operator for ART index range scans.
189///
190/// Uses the ART index on a node table's PK column to efficiently find rows
191/// within a key range, then fetches the full column data for those rows.
192///
193/// Pattern follows `PhysicalVectorSimilarityScan`.
194#[derive(Debug, Clone)]
195pub struct PhysicalArtIndexRangeScan {
196    pub table_name: String,
197    pub table_id: u64,
198    pub lower_bound: Option<Value>,
199    pub upper_bound: Option<Value>,
200    pub lower_inclusive: bool,
201    pub upper_inclusive: bool,
202    pub table_catalog: Option<Arc<TableCatalog>>,
203}
204
205impl PhysicalOperatorExec for PhysicalArtIndexRangeScan {
206    fn operator_type(&self) -> &str {
207        "art_index_range_scan"
208    }
209
210    fn execute(&self, _input: Vec<DataChunk>) -> OperatorResult {
211        let tc = self
212            .table_catalog
213            .clone()
214            .ok_or_else(|| "No table catalog available for ArtIndexRangeScan".to_string())?;
215
216        let node_table = tc
217            .get_node_table_by_name(&self.table_name)
218            .ok_or_else(|| format!("Node table '{}' not found", self.table_name))?;
219
220        // Verify ART index exists
221        if node_table.art_index.is_none() {
222            return Err(format!("Table '{}' does not have an ART index", self.table_name).into());
223        }
224
225        // Execute range scan on the ART index
226        let row_ids = node_table.lookup_by_pk_range(
227            self.lower_bound.as_ref(),
228            self.lower_inclusive,
229            self.upper_bound.as_ref(),
230            self.upper_inclusive,
231            u64::MAX,
232        );
233        drop(node_table); // Release table ref before cloning data
234
235        if row_ids.is_empty() {
236            return Ok(vec![DataChunk::new(vec![], vec![])]);
237        }
238
239        // Fetch column values for matched row IDs
240        let node_table = tc
241            .get_node_table_by_name(&self.table_name)
242            .ok_or_else(|| format!("Node table '{}' not found", self.table_name))?;
243
244        let num_cols = node_table.columns.len();
245        let num_results = row_ids.len();
246
247        // +1 for the internal node id column (`<var>._id` = row offset)
248        let mut output_columns: Vec<Vec<Value>> = vec![Vec::with_capacity(num_results); num_cols + 1];
249
250        for &row_id in &row_ids {
251            for (col_idx, out_col) in output_columns.iter_mut().enumerate().take(num_cols) {
252                match node_table.get_value(row_id as usize, col_idx) {
253                    Some(val) => out_col.push(val.clone()),
254                    None => out_col.push(Value::Null),
255                }
256            }
257            output_columns[num_cols].push(Value::Int64(row_id as i64));
258        }
259
260        let mut col_types = Vec::with_capacity(num_cols + 1);
261        let mut col_names: Vec<String> = Vec::with_capacity(num_cols + 1);
262        for col in &node_table.columns {
263            col_types.push(col.logical_type);
264            col_names.push(col.name.clone());
265        }
266        col_types.push(akar_common::types::LogicalTypeID::Int64);
267        col_names.push("_id".to_string());
268
269        drop(node_table);
270
271        // Convert column-major Vec<Vec<Value>> to DataChunks
272        use akar_common::types::PhysicalTypeID;
273        use akar_common::vector::{DataChunk, ValueVector};
274
275        let num_rows = output_columns.first().map(|c| c.len()).unwrap_or(0);
276        if num_rows == 0 {
277            return Ok(vec![DataChunk::new(vec![], vec![])]);
278        }
279
280        let mut chunks = Vec::new();
281        let chunk_size = 1024usize;
282        for start in (0..num_rows).step_by(chunk_size) {
283            let end = (start + chunk_size).min(num_rows);
284            let count = end - start;
285            let mut fields = Vec::with_capacity(num_cols + 1);
286
287            for col_idx in 0..num_cols + 1 {
288                let col_data = &output_columns[col_idx];
289                let phys_type = match col_types[col_idx] {
290                    akar_common::types::LogicalTypeID::Bool => PhysicalTypeID::Bool,
291                    akar_common::types::LogicalTypeID::Int64 | akar_common::types::LogicalTypeID::Serial => {
292                        PhysicalTypeID::Int64
293                    }
294                    akar_common::types::LogicalTypeID::Int32 => PhysicalTypeID::Int32,
295                    akar_common::types::LogicalTypeID::Int16 => PhysicalTypeID::Int16,
296                    akar_common::types::LogicalTypeID::Int8 => PhysicalTypeID::Int8,
297                    akar_common::types::LogicalTypeID::UInt64 => PhysicalTypeID::UInt64,
298                    akar_common::types::LogicalTypeID::UInt32 => PhysicalTypeID::UInt32,
299                    akar_common::types::LogicalTypeID::UInt16 => PhysicalTypeID::UInt16,
300                    akar_common::types::LogicalTypeID::UInt8 => PhysicalTypeID::UInt8,
301                    akar_common::types::LogicalTypeID::Double => PhysicalTypeID::Double,
302                    akar_common::types::LogicalTypeID::Float => PhysicalTypeID::Float,
303                    akar_common::types::LogicalTypeID::String => PhysicalTypeID::String,
304                    akar_common::types::LogicalTypeID::Blob => PhysicalTypeID::Blob,
305                    akar_common::types::LogicalTypeID::Date => PhysicalTypeID::Int32,
306                    akar_common::types::LogicalTypeID::Timestamp => PhysicalTypeID::Int64,
307                    akar_common::types::LogicalTypeID::Interval => PhysicalTypeID::Interval,
308                    akar_common::types::LogicalTypeID::List => PhysicalTypeID::List,
309                    akar_common::types::LogicalTypeID::Array => PhysicalTypeID::Array,
310                    akar_common::types::LogicalTypeID::Struct => PhysicalTypeID::Struct,
311                    akar_common::types::LogicalTypeID::Node => PhysicalTypeID::Struct,
312                    akar_common::types::LogicalTypeID::Rel => PhysicalTypeID::Struct,
313                    akar_common::types::LogicalTypeID::InternalID => PhysicalTypeID::Struct, // Internal IDs are Structs
314                    _ => PhysicalTypeID::Any,
315                };
316                let mut vv = ValueVector::new(phys_type, count);
317                vv.resize(count);
318                for row_offset in 0..count {
319                    let val = &col_data[start + row_offset];
320                    match val {
321                        Value::Null => vv.set_null(row_offset, true),
322                        Value::Int64(x) => {
323                            let buf = &mut vv.data_mut()[row_offset * 8..(row_offset + 1) * 8];
324                            buf.copy_from_slice(&x.to_le_bytes());
325                        }
326                        Value::Int32(x) => {
327                            let buf = &mut vv.data_mut()[row_offset * 4..(row_offset + 1) * 4];
328                            buf.copy_from_slice(&x.to_le_bytes());
329                        }
330                        Value::Double(x) => {
331                            let buf = &mut vv.data_mut()[row_offset * 8..(row_offset + 1) * 8];
332                            buf.copy_from_slice(&x.to_le_bytes());
333                        }
334                        Value::String(_) => {
335                            vv.set_value(row_offset, val)?;
336                        }
337                        _ => {}
338                    }
339                }
340                fields.push(vv);
341            }
342
343            let arrow_fields = fields
344                .iter()
345                .map(|v| akar_common::arrow_vector::ArrowVector::from_legacy(v).array)
346                .collect::<Vec<_>>();
347            let arrow_field_types = fields.iter().map(|v| v.physical_type()).collect::<Vec<_>>();
348            chunks.push(DataChunk {
349                fields: arrow_fields,
350                field_types: arrow_field_types,
351                size: count,
352                field_names: col_names.clone(),
353                sel_vector: None,
354            });
355        }
356
357        Ok(chunks)
358    }
359}