use std::collections::HashMap;
use std::fs::File;
use std::io::BufReader;
use std::sync::Arc;
use parking_lot::RwLock;
use rayon::prelude::*;
use crate::executor::streaming::operators::spec::CopyTarget;
use crate::executor::streaming::runtime::ExecutionRuntime;
use crate::storage::{QueryStorage, StorageWriter};
use linkrs_core::error::QueryError;
use linkrs_core::types::storage_ids::VertexId;
use linkrs_core::vertex_edge_path::{Edge, Tag, Vertex};
use linkrs_core::Value;
const DEFAULT_BATCH_SIZE: usize = 1000;
#[allow(clippy::too_many_arguments)]
pub fn execute_copy_from(
storage_lock: &Arc<RwLock<dyn QueryStorage>>,
space_name: &str,
target: &CopyTarget,
file_paths: &[String],
by_column: bool,
header: bool,
delimiter: u8,
batch_size: usize,
runtime: Option<Arc<ExecutionRuntime>>,
) -> Result<u64, QueryError> {
if file_paths.is_empty() {
return Err(QueryError::execution(
"COPY FROM requires at least one file".to_string(),
));
}
let batch_sz = if batch_size == 0 {
DEFAULT_BATCH_SIZE
} else {
batch_size
};
match target {
CopyTarget::Vertex(tag) => {
let schema_props = schema_property_names_vertex(storage_lock, space_name, tag, header)?;
run_import(
storage_lock,
space_name,
file_paths,
by_column,
header,
delimiter,
batch_sz,
runtime,
&ImportPlan::Vertices {
tag: tag.clone(),
schema_props,
},
)
}
CopyTarget::Edge(edge_type) => {
let schema_props =
schema_property_names_edge(storage_lock, space_name, edge_type, header)?;
run_import(
storage_lock,
space_name,
file_paths,
by_column,
header,
delimiter,
batch_sz,
runtime,
&ImportPlan::Edges {
edge_type: edge_type.clone(),
schema_props,
},
)
}
}
}
fn schema_property_names_vertex(
storage_lock: &Arc<RwLock<dyn QueryStorage>>,
space_name: &str,
tag: &str,
header: bool,
) -> Result<Option<Vec<String>>, QueryError> {
if header {
return Ok(None);
}
let read = storage_lock.read();
match read.get_tag(space_name, tag) {
Ok(Some(info)) => Ok(Some(
info.properties.iter().map(|p| p.name.clone()).collect(),
)),
Ok(None) => Err(QueryError::execution(format!("Tag '{tag}' not found"))),
Err(e) => Err(QueryError::execution(e.to_string())),
}
}
fn schema_property_names_edge(
storage_lock: &Arc<RwLock<dyn QueryStorage>>,
space_name: &str,
edge_type: &str,
header: bool,
) -> Result<Option<Vec<String>>, QueryError> {
if header {
return Ok(None);
}
let read = storage_lock.read();
match read.get_edge_type(space_name, edge_type) {
Ok(Some(info)) => Ok(Some(
info.properties.iter().map(|p| p.name.clone()).collect(),
)),
Ok(None) => Err(QueryError::execution(format!(
"Edge type '{edge_type}' not found"
))),
Err(e) => Err(QueryError::execution(e.to_string())),
}
}
enum ImportPlan {
Vertices {
tag: String,
schema_props: Option<Vec<String>>,
},
Edges {
edge_type: String,
schema_props: Option<Vec<String>>,
},
}
#[allow(clippy::too_many_arguments)]
fn run_import(
storage_lock: &Arc<RwLock<dyn QueryStorage>>,
space_name: &str,
file_paths: &[String],
by_column: bool,
header: bool,
delimiter: u8,
batch_size: usize,
runtime: Option<Arc<ExecutionRuntime>>,
plan: &ImportPlan,
) -> Result<u64, QueryError> {
let mut merged = MergedCsvSource::open(file_paths, by_column, header, delimiter)?;
let mapping = match plan {
ImportPlan::Vertices { schema_props, .. } => build_column_mapping(
merged.headers(),
header,
schema_props.as_deref(),
KeyLayout::Vertex,
)?,
ImportPlan::Edges { schema_props, .. } => build_column_mapping(
merged.headers(),
header,
schema_props.as_deref(),
KeyLayout::Edge,
)?,
};
let guard = storage_lock.read();
let storage: &dyn QueryStorage = &*guard;
let window = storage.begin_auto_commit_group().map_err(|e| {
QueryError::execution(format!("COPY FROM failed to open write window: {e}"))
})?;
let result = (|| -> Result<u64, QueryError> {
let mut total: u64 = 0;
let mut skipped = 0u64;
while !merged.is_eof() {
if let Some(rt) = &runtime {
rt.ensure_not_cancelled()
.map_err(|e| QueryError::execution(e.to_string()))?;
}
let records = merged.next_batch(batch_size, &mapping, &mut skipped)?;
if records.is_empty() {
continue;
}
let inserted = {
let mut writer = storage.bind_auto_commit_writer(&window).map_err(|e| {
QueryError::execution(format!("COPY FROM failed to bind writer: {e}"))
})?;
match plan {
ImportPlan::Vertices { tag, .. } => {
flush_vertices(&mut *writer, space_name, tag, &mapping, &records)?
}
ImportPlan::Edges { edge_type, .. } => {
flush_edges(&mut *writer, space_name, edge_type, &mapping, &records)?
}
}
};
total += inserted;
}
if skipped > 0 {
log::warn!(
"COPY FROM '{}': skipped {skipped} malformed record(s) \
(wrong field count)",
file_paths.join(", ")
);
}
Ok(total)
})();
match result {
Ok(total) => {
storage
.finalize_auto_commit_group(&window)
.map_err(|e| QueryError::execution(format!("COPY FROM commit failed: {e}")))?;
Ok(total)
}
Err(error) => {
if let Err(rollback_error) = storage.rollback_auto_commit_group(&window) {
log::error!("COPY FROM rollback failed: {rollback_error}");
}
Err(error)
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum KeyLayout {
Vertex,
Edge,
}
fn flush_vertices(
writer: &mut dyn StorageWriter,
space_name: &str,
tag: &str,
mapping: &ColumnMapping,
records: &[Vec<String>],
) -> Result<u64, QueryError> {
let vertices: Vec<Vertex> = records
.par_iter()
.map(|rec| {
let vid_str = rec[mapping.vid_index].trim();
let vid = parse_vid(vid_str)?;
let props = collect_properties(mapping, rec);
Ok(Vertex::new(vid, Tag::new(tag.to_string(), props)))
})
.collect::<Result<Vec<_>, QueryError>>()?;
let count = vertices.len() as u64;
StorageWriter::batch_insert_vertices(writer, space_name, vertices)
.map_err(|e| QueryError::execution(e.to_string()))?;
Ok(count)
}
fn flush_edges(
writer: &mut dyn StorageWriter,
space_name: &str,
edge_type: &str,
mapping: &ColumnMapping,
records: &[Vec<String>],
) -> Result<u64, QueryError> {
let edges: Vec<Edge> = records
.par_iter()
.map(|rec| {
let src = parse_vid(rec[mapping.src_index].trim())?;
let dst = parse_vid(rec[mapping.dst_index].trim())?;
let props = collect_properties(mapping, rec);
Ok(Edge {
src,
dst,
edge_type: edge_type.to_string(),
ranking: 0,
props,
})
})
.collect::<Result<Vec<_>, QueryError>>()?;
let count = edges.len() as u64;
StorageWriter::batch_insert_edges(writer, space_name, edges)
.map_err(|e| QueryError::execution(e.to_string()))?;
Ok(count)
}
fn collect_properties(mapping: &ColumnMapping, rec: &[String]) -> HashMap<Arc<str>, Value> {
let mut props = HashMap::with_capacity(mapping.property_indices.len());
for &(idx, ref name) in &mapping.property_indices {
props.insert(name.as_str().into(), parse_copy_value(rec[idx].trim()));
}
props
}
fn parse_vid(s: &str) -> Result<VertexId, QueryError> {
if s.is_empty() {
return Err(QueryError::execution(
"COPY FROM: empty vertex id cell".to_string(),
));
}
if let Ok(i) = s.parse::<i64>() {
VertexId::try_from_int64(i).map_err(|e| QueryError::execution(e.to_string()))
} else {
VertexId::try_from_string(s).map_err(|e| QueryError::execution(e.to_string()))
}
}
#[derive(Debug, Clone)]
struct ColumnMapping {
vid_index: usize,
src_index: usize,
dst_index: usize,
property_indices: Vec<(usize, String)>,
expected_width: usize,
}
fn find_column(headers: &[String], aliases: &[&str]) -> Option<usize> {
headers
.iter()
.position(|h| aliases.iter().any(|a| h.eq_ignore_ascii_case(a)))
}
fn build_column_mapping(
headers: &[String],
has_header: bool,
schema_props: Option<&[String]>,
layout: KeyLayout,
) -> Result<ColumnMapping, QueryError> {
if !has_header {
let schema_props = schema_props.ok_or_else(|| {
QueryError::execution(
"COPY FROM without HEADER requires the target tag / edge type to \
declare properties"
.to_string(),
)
})?;
let key_count = match layout {
KeyLayout::Vertex => 1usize,
KeyLayout::Edge => 2,
};
let property_indices: Vec<(usize, String)> = schema_props
.iter()
.enumerate()
.map(|(i, name)| (key_count + i, name.clone()))
.collect();
return Ok(ColumnMapping {
vid_index: 0,
src_index: 0,
dst_index: 1,
property_indices,
expected_width: key_count + schema_props.len(),
});
}
match layout {
KeyLayout::Vertex => {
let vid_index =
find_column(headers, &["vid", "id", "_id", "vertex_id"]).ok_or_else(|| {
QueryError::execution(
"COPY FROM: header row has no vertex id column \
(expected one of: vid, id, _id, vertex_id)"
.to_string(),
)
})?;
let property_indices = excluded_pairs(headers, &[vid_index]);
Ok(ColumnMapping {
vid_index,
src_index: vid_index,
dst_index: vid_index,
property_indices,
expected_width: headers.len(),
})
}
KeyLayout::Edge => {
let src_hit = find_column(headers, &["src", "_src", "source"]);
let dst_hit = find_column(headers, &["dst", "_dst", "destination", "dest"]);
let (src_index, dst_index) = match (src_hit, dst_hit) {
(Some(s), Some(d)) => {
if s == d {
return Err(QueryError::execution(
"COPY FROM: src and dst resolve to the same column".to_string(),
));
}
(s, d)
}
(None, None) => (0, 1),
(Some(_), None) => {
return Err(QueryError::execution(
"COPY FROM: header row names a source column but no destination \
column (expected one of: dst, _dst, destination, dest)"
.to_string(),
))
}
(None, Some(_)) => {
return Err(QueryError::execution(
"COPY FROM: header row names a destination column but no source \
column (expected one of: src, _src, source)"
.to_string(),
))
}
};
let property_indices = excluded_pairs(headers, &[src_index, dst_index]);
Ok(ColumnMapping {
vid_index: src_index,
src_index,
dst_index,
property_indices,
expected_width: headers.len(),
})
}
}
}
fn excluded_pairs(headers: &[String], exclude: &[usize]) -> Vec<(usize, String)> {
headers
.iter()
.enumerate()
.filter(|(i, _)| !exclude.contains(i))
.map(|(i, h)| (i, h.clone()))
.collect()
}
struct CsvSource {
headers: Vec<String>,
reader: csv::StringRecordsIntoIter<std::io::BufReader<File>>,
finished: bool,
}
impl CsvSource {
fn open(file_path: &str, header: bool, delimiter: u8) -> Result<Self, QueryError> {
let file = File::open(file_path).map_err(|e| {
QueryError::execution(format!("COPY FROM failed to open '{file_path}': {e}"))
})?;
let mut reader = csv::ReaderBuilder::new()
.has_headers(header)
.delimiter(delimiter)
.trim(csv::Trim::All)
.flexible(true)
.from_reader(BufReader::new(file));
let headers: Vec<String> = if header {
reader
.headers()
.map_err(|e| QueryError::execution(format!("COPY CSV header error: {e}")))?
.iter()
.map(|s| s.to_string())
.collect()
} else {
Vec::new()
};
Ok(Self {
headers,
reader: reader.into_records(),
finished: false,
})
}
fn headers(&self) -> &[String] {
&self.headers
}
fn is_eof(&self) -> bool {
self.finished
}
fn next_batch(
&mut self,
batch_size: usize,
mapping: &ColumnMapping,
skipped: &mut u64,
) -> Result<Vec<Vec<String>>, QueryError> {
let mut out: Vec<Vec<String>> = Vec::with_capacity(batch_size);
while out.len() < batch_size {
let Some(record) = self.next_record()? else {
break;
};
if record.len() != mapping.expected_width {
*skipped += 1;
continue;
}
out.push(record.iter().map(|s| s.to_string()).collect());
}
Ok(out)
}
fn next_record(&mut self) -> Result<Option<csv::StringRecord>, QueryError> {
if self.finished {
return Ok(None);
}
match self.reader.next() {
Some(Ok(rec)) => Ok(Some(rec)),
Some(Err(e)) => Err(QueryError::execution(format!("COPY CSV read error: {e}"))),
None => {
self.finished = true;
Ok(None)
}
}
}
}
enum MergedCsvSource {
Rows {
headers: Vec<String>,
sources: Vec<CsvSource>,
current: usize,
},
Columns {
headers: Vec<String>,
rows: Vec<Vec<String>>,
pos: usize,
},
}
impl MergedCsvSource {
fn open(
file_paths: &[String],
by_column: bool,
header: bool,
delimiter: u8,
) -> Result<Self, QueryError> {
if file_paths.is_empty() {
return Err(QueryError::execution(
"COPY FROM requires at least one file".to_string(),
));
}
if !by_column {
let mut sources = Vec::with_capacity(file_paths.len());
for path in file_paths {
sources.push(CsvSource::open(path, header, delimiter)?);
}
let headers = sources
.first()
.map(|s| s.headers().to_vec())
.unwrap_or_default();
if header {
for (idx, source) in sources.iter().enumerate().skip(1) {
if source.headers() != headers.as_slice() {
return Err(QueryError::execution(format!(
"COPY FROM: header mismatch between '{}' and '{}'",
file_paths[0], file_paths[idx]
)));
}
}
}
return Ok(Self::Rows {
headers,
sources,
current: 0,
});
}
let mut per_file_headers: Vec<Vec<String>> = Vec::with_capacity(file_paths.len());
let mut per_file_rows: Vec<Vec<Vec<String>>> = Vec::with_capacity(file_paths.len());
for path in file_paths {
let mut source = CsvSource::open(path, header, delimiter)?;
let file_headers = source.headers().to_vec();
let file_width = file_headers.len();
let mut rows: Vec<Vec<String>> = Vec::new();
while let Some(record) = source.next_record()? {
rows.push(record.iter().map(|s| s.to_string()).collect());
}
if header {
for (idx, row) in rows.iter().enumerate() {
if row.len() != file_width {
return Err(QueryError::execution(format!(
"COPY FROM BY COLUMN: file '{path}' row {} has {} fields, expected {file_width}",
idx + 1,
row.len()
)));
}
}
}
per_file_headers.push(file_headers);
per_file_rows.push(rows);
}
let row_count = per_file_rows.first().map(|r| r.len()).unwrap_or(0);
for (idx, rows) in per_file_rows.iter().enumerate().skip(1) {
if rows.len() != row_count {
return Err(QueryError::execution(format!(
"COPY FROM BY COLUMN: file '{}' has {} rows, expected {row_count} (from '{}')",
file_paths[idx],
rows.len(),
file_paths[0]
)));
}
}
let headers: Vec<String> = per_file_headers.into_iter().flatten().collect();
let mut rows: Vec<Vec<String>> = Vec::with_capacity(row_count);
for row_idx in 0..row_count {
let mut merged: Vec<String> = Vec::new();
for file_rows in &per_file_rows {
merged.extend(file_rows[row_idx].iter().cloned());
}
rows.push(merged);
}
Ok(Self::Columns {
headers,
rows,
pos: 0,
})
}
fn headers(&self) -> &[String] {
match self {
Self::Rows { headers, .. } => headers,
Self::Columns { headers, .. } => headers,
}
}
fn is_eof(&self) -> bool {
match self {
Self::Rows {
sources, current, ..
} => *current >= sources.len(),
Self::Columns { rows, pos, .. } => *pos >= rows.len(),
}
}
fn next_batch(
&mut self,
batch_size: usize,
mapping: &ColumnMapping,
skipped: &mut u64,
) -> Result<Vec<Vec<String>>, QueryError> {
match self {
Self::Rows {
sources, current, ..
} => {
let mut out: Vec<Vec<String>> = Vec::with_capacity(batch_size);
while out.len() < batch_size && *current < sources.len() {
let batch =
sources[*current].next_batch(batch_size - out.len(), mapping, skipped)?;
if batch.is_empty() {
if sources[*current].is_eof() {
*current += 1;
continue;
}
break;
}
out.extend(batch);
}
Ok(out)
}
Self::Columns { rows, pos, .. } => {
let mut out: Vec<Vec<String>> = Vec::with_capacity(batch_size);
while out.len() < batch_size && *pos < rows.len() {
let row = &rows[*pos];
*pos += 1;
if row.len() != mapping.expected_width {
*skipped += 1;
continue;
}
out.push(row.clone());
}
Ok(out)
}
}
}
}
fn parse_copy_value(s: &str) -> Value {
if s.is_empty() {
return Value::Null(linkrs_core::value::NullType::Null);
}
if s.eq_ignore_ascii_case("true") {
return Value::Bool(true);
}
if s.eq_ignore_ascii_case("false") {
return Value::Bool(false);
}
if s.eq_ignore_ascii_case("null") {
return Value::Null(linkrs_core::value::NullType::Null);
}
if let Ok(i) = s.parse::<i64>() {
return Value::BigInt(i);
}
if let Ok(f) = s.parse::<f64>() {
return Value::Double(f);
}
Value::string(s)
}
pub fn execute_copy_to(
storage_lock: &Arc<RwLock<dyn QueryStorage>>,
space_name: &str,
target: &CopyTarget,
file_path: &str,
header: bool,
delimiter: u8,
) -> Result<u64, QueryError> {
let delim = char::from(delimiter);
let file = File::create(file_path)
.map_err(|e| QueryError::execution(format!("COPY TO: cannot create '{file_path}': {e}")))?;
let mut writer = std::io::BufWriter::new(file);
match target {
CopyTarget::Vertex(tag) => {
let prop_names: Vec<String> = {
let read = storage_lock.read();
match read.get_tag(space_name, tag) {
Ok(Some(info)) => info.properties.iter().map(|p| p.name.clone()).collect(),
Ok(None) => {
return Err(QueryError::execution(format!("Tag '{tag}' not found")))
}
Err(e) => return Err(QueryError::execution(e.to_string())),
}
};
let vertices = {
let read = storage_lock.read();
read.scan_vertices_by_tag(space_name, tag)
.map_err(|e| QueryError::execution(e.to_string()))?
};
if header {
let mut head = vec!["vid".to_string()];
head.extend(prop_names.iter().cloned());
write_csv_line(&mut writer, &head, delim)?;
}
let mut count = 0u64;
for vertex in &vertices {
let props = vertex.properties();
let mut cells = vec![csv_cell(&Value::from(vertex.vid), delim)];
cells.extend(prop_names.iter().map(|name| {
props
.get(name.as_str())
.map_or_else(String::new, |v| csv_cell(v, delim))
}));
write_csv_line(&mut writer, &cells, delim)?;
count += 1;
}
Ok(count)
}
CopyTarget::Edge(edge_type) => {
let prop_names: Vec<String> = {
let read = storage_lock.read();
match read.get_edge_type(space_name, edge_type) {
Ok(Some(info)) => info.properties.iter().map(|p| p.name.clone()).collect(),
Ok(None) => {
return Err(QueryError::execution(format!(
"Edge type '{edge_type}' not found"
)))
}
Err(e) => return Err(QueryError::execution(e.to_string())),
}
};
let edges = {
let read = storage_lock.read();
read.scan_edges_by_type(space_name, edge_type)
.map_err(|e| QueryError::execution(e.to_string()))?
};
if header {
let mut head = vec!["src".to_string(), "dst".to_string()];
head.extend(prop_names.iter().cloned());
write_csv_line(&mut writer, &head, delim)?;
}
let mut count = 0u64;
for edge in &edges {
let mut cells = vec![
csv_cell(&Value::from(*edge.src()), delim),
csv_cell(&Value::from(*edge.dst()), delim),
];
cells.extend(prop_names.iter().map(|name| {
edge.get_property(name.as_str())
.map_or_else(String::new, |v| csv_cell(v, delim))
}));
write_csv_line(&mut writer, &cells, delim)?;
count += 1;
}
Ok(count)
}
}
}
fn csv_cell(value: &Value, delimiter: char) -> String {
let text = match value {
Value::Null(_) => String::new(),
Value::String(s) => s.to_string(),
Value::FixedString(s) => s.to_string(),
other => format!("{other}"),
};
if text.contains(delimiter) || text.contains('"') || text.contains('\n') || text.contains('\r')
{
format!("\"{}\"", text.replace('"', "\"\""))
} else {
text
}
}
fn write_csv_line<W: std::io::Write>(
writer: &mut W,
cells: &[String],
delimiter: char,
) -> Result<(), QueryError> {
let mut line = String::with_capacity(32 * cells.len());
for (i, cell) in cells.iter().enumerate() {
if i > 0 {
line.push(delimiter);
}
line.push_str(cell);
}
line.push('\n');
writer
.write_all(line.as_bytes())
.map_err(|e| QueryError::execution(format!("COPY TO: write failed: {e}")))
}
#[cfg(test)]
mod tests {
use super::*;
fn headers(names: &[&str]) -> Vec<String> {
names.iter().map(|s| s.to_string()).collect()
}
#[test]
fn vertex_header_mapping_locates_vid_and_properties() {
let h = headers(&["name", "vid", "age"]);
let m = build_column_mapping(&h, true, None, KeyLayout::Vertex).expect("mapping");
assert_eq!(m.vid_index, 1);
assert_eq!(
m.property_indices,
vec![(0, "name".to_string()), (2, "age".to_string())]
);
assert_eq!(m.expected_width, 3);
}
#[test]
fn vertex_header_without_id_column_errors() {
let h = headers(&["name", "age"]);
assert!(build_column_mapping(&h, true, None, KeyLayout::Vertex).is_err());
}
#[test]
fn edge_header_named_endpoints() {
let h = headers(&["since", "dst", "src"]);
let m = build_column_mapping(&h, true, None, KeyLayout::Edge).expect("mapping");
assert_eq!(m.src_index, 2);
assert_eq!(m.dst_index, 1);
assert_eq!(m.property_indices, vec![(0, "since".to_string())]);
}
#[test]
fn edge_header_only_src_named_errors() {
let h = headers(&["src", "since"]);
assert!(build_column_mapping(&h, true, None, KeyLayout::Edge).is_err());
}
#[test]
fn edge_header_no_named_endpoints_falls_back_positional() {
let h = headers(&["a", "b", "since"]);
let m = build_column_mapping(&h, true, None, KeyLayout::Edge).expect("mapping");
assert_eq!(m.src_index, 0);
assert_eq!(m.dst_index, 1);
assert_eq!(m.property_indices, vec![(2, "since".to_string())]);
}
#[test]
fn positional_layout_uses_schema_order() {
let schema = vec!["name".to_string(), "age".to_string()];
let m =
build_column_mapping(&[], false, Some(&schema), KeyLayout::Vertex).expect("mapping");
assert_eq!(m.vid_index, 0);
assert_eq!(
m.property_indices,
vec![(1, "name".to_string()), (2, "age".to_string())]
);
assert_eq!(m.expected_width, 3);
}
#[test]
fn parse_value_inference() {
assert_eq!(
parse_copy_value(""),
Value::Null(linkrs_core::value::NullType::Null)
);
assert_eq!(parse_copy_value("true"), Value::Bool(true));
assert_eq!(parse_copy_value("42"), Value::BigInt(42));
assert_eq!(parse_copy_value("4.5"), Value::Double(4.5));
assert_eq!(parse_copy_value("hello"), Value::string("hello"));
}
}