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