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

akar_processor/physical/write_ops/
copyfrom.rs
