use radixdb_core::{time_compat::Instant, CompactArc, SmartString};
use radixdb_core::{DataType, Error, Result, Row, Schema, Value};
use radixdb_sql::ast::{CopyFormat, CopyStatement};
use radixdb_storage::traits::{Engine, QueryResult, Table};
use rustc_hash::{FxHashMap, FxHashSet};
use super::dml_support::evaluate_default_expr;
use crate::context::{
invalidate_in_subquery_cache_for_table, invalidate_scalar_subquery_cache_for_table,
invalidate_semi_join_cache_for_table, ExecutionContext,
};
use crate::mutation::host::MutationHost;
use crate::mutation::validation::{
compile_table_check_constraints, prepare_insert_row_constraints,
};
use crate::result::ExecResult;
const COPY_TRANSACTION_MEMORY_AMPLIFICATION: usize = 8;
const COPY_INSERT_BATCH_ROWS: usize = 4096;
#[inline]
fn flush_copy_insert_batch(table: &mut Box<dyn Table>, rows: &mut Vec<Row>) -> Result<()> {
if rows.is_empty() {
return Ok(());
}
let next = Vec::with_capacity(COPY_INSERT_BATCH_ROWS);
table.insert_batch(std::mem::replace(rows, next))
}
fn account_copy_transaction_row(
used_bytes: &mut usize,
limit_bytes: usize,
row: &Row,
row_number: i64,
) -> Result<()> {
let row_bytes = radixdb_storage::mvcc::version_store::estimate_row_hot_bytes(row)
.saturating_mul(COPY_TRANSACTION_MEMORY_AMPLIFICATION);
let attempted_bytes = used_bytes.saturating_add(row_bytes);
if attempted_bytes > limit_bytes {
return Err(Error::CopyTransactionMemoryLimit {
row: row_number,
limit_bytes,
attempted_bytes,
});
}
*used_bytes = attempted_bytes;
Ok(())
}
#[inline]
fn parse_field(field: &str, target_type: DataType, col_name: &str) -> Result<Value> {
match target_type {
DataType::Integer => field.parse::<i64>().map(Value::Integer).map_err(|_| {
Error::Type(format!(
"cannot convert value '{}' to INTEGER for column '{}'",
field, col_name
))
}),
DataType::Float => field.parse::<f64>().map(Value::Float).map_err(|_| {
Error::Type(format!(
"cannot convert value '{}' to FLOAT for column '{}'",
field, col_name
))
}),
DataType::Boolean => {
if field.eq_ignore_ascii_case("true")
|| field.eq_ignore_ascii_case("t")
|| field.eq_ignore_ascii_case("yes")
|| field.eq_ignore_ascii_case("y")
|| field == "1"
{
Ok(Value::Boolean(true))
} else if field.eq_ignore_ascii_case("false")
|| field.eq_ignore_ascii_case("f")
|| field.eq_ignore_ascii_case("no")
|| field.eq_ignore_ascii_case("n")
|| field == "0"
{
Ok(Value::Boolean(false))
} else {
Err(Error::Type(format!(
"cannot convert value '{}' to BOOLEAN for column '{}'",
field, col_name
)))
}
}
DataType::Timestamp => radixdb_core::parse_timestamp(field)
.map(Value::Timestamp)
.map_err(|_| {
Error::Type(format!(
"cannot convert value '{}' to TIMESTAMP for column '{}'",
field, col_name
))
}),
DataType::Decimal => radixdb_core::value::parse_decimal_str(field)
.and_then(|(unscaled, precision, scale)| {
Value::try_decimal(unscaled, precision, scale).ok()
})
.ok_or_else(|| {
Error::Type(format!(
"cannot convert value '{}' to DECIMAL for column '{}'",
field, col_name
))
}),
DataType::Date => radixdb_core::value::parse_date_days_since_unix_epoch(field)
.map(Value::date)
.ok_or_else(|| {
Error::Type(format!(
"cannot convert value '{}' to DATE for column '{}'",
field, col_name
))
}),
DataType::Bytes => Ok(Value::bytes(field.as_bytes().to_vec())),
DataType::Text => {
Ok(Value::Text(SmartString::new(field)))
}
DataType::Json => Value::try_json(field).map_err(|_| {
Error::Type(format!(
"cannot convert value '{}' to JSON for column '{}'",
field, col_name
))
}),
DataType::Uuid => radixdb_core::value::parse_uuid_str(field)
.map(Value::uuid)
.ok_or_else(|| {
Error::Type(format!(
"cannot convert value '{}' to UUID for column '{}'",
field, col_name
))
}),
_ => {
let text_val = Value::text(field);
let coerced = text_val.coerce_to_type(target_type);
if !text_val.is_null() && coerced.is_null() {
return Err(Error::Type(format!(
"cannot convert value '{}' to {:?} for column '{}'",
field, target_type, col_name
)));
}
Ok(coerced)
}
}
}
#[doc(hidden)]
pub trait CopyExecutorExt: MutationHost {
fn execute_copy(
&self,
stmt: &CopyStatement,
_ctx: &ExecutionContext,
) -> Result<Box<dyn QueryResult>> {
let copy_started = Instant::now();
let table_name = &stmt.table_name.value_lower;
{
let active_tx = self.mutation_active_transaction().lock().unwrap();
if active_tx.is_some() {
return Err(Error::InvalidArgument(
"COPY FROM cannot be used inside an explicit transaction".to_string(),
));
}
}
let mut transaction = self.mutation_engine().begin_transaction()?;
let mut table = transaction.get_table(table_name)?;
let schema = table.schema().clone();
let schema_column_count = schema.columns.len();
let copy_max_transaction_bytes = self
.mutation_engine()
.config()
.persistence
.copy_max_transaction_bytes;
let mut copy_transaction_bytes = 0usize;
let all_column_types: Vec<DataType> = schema.columns.iter().map(|c| c.data_type).collect();
let all_vector_dims: Vec<u16> =
schema.columns.iter().map(|c| c.vector_dimensions).collect();
let default_exprs: Vec<Option<String>> = schema
.columns
.iter()
.map(|c| c.default_expr.clone())
.collect();
let column_indices: Vec<usize> = if stmt.columns.is_empty() {
(0..schema_column_count).collect()
} else {
let mut seen = FxHashSet::default();
for identifier in &stmt.columns {
if !seen.insert(identifier.value_lower.clone()) {
return Err(Error::InvalidArgument(format!(
"duplicate COPY target column '{}'",
identifier.value
)));
}
}
let col_map = schema.column_index_map();
stmt.columns
.iter()
.map(|id| {
col_map
.get(id.value_lower.as_str())
.copied()
.ok_or_else(|| Error::ColumnNotFound(id.value.to_string()))
})
.collect::<Result<Vec<_>>>()?
};
let fk_schema: Option<CompactArc<Schema>> = if !schema.foreign_keys.is_empty() {
Some(CompactArc::new(schema.clone()))
} else {
None
};
let parse_started = Instant::now();
let outcome = match stmt.format {
CopyFormat::Csv => self.copy_from_csv(
stmt,
&mut table,
&schema,
&column_indices,
&all_column_types,
&all_vector_dims,
&default_exprs,
&fk_schema,
schema_column_count,
&mut copy_transaction_bytes,
copy_max_transaction_bytes,
),
CopyFormat::Json => self.copy_from_json(
stmt,
&mut table,
&schema,
&column_indices,
&all_column_types,
&all_vector_dims,
&default_exprs,
&fk_schema,
schema_column_count,
&mut copy_transaction_bytes,
copy_max_transaction_bytes,
),
};
let parse_elapsed = parse_started.elapsed();
drop(table);
let rows_affected = match outcome {
Ok(rows_affected) => {
let commit_started = Instant::now();
let commit_result = transaction.commit();
let commit_elapsed = commit_started.elapsed();
if let Err(error) = commit_result {
radixdb_storage::instrumentation::record_copy(
rows_affected.max(0) as u64,
parse_elapsed,
commit_elapsed,
copy_started.elapsed(),
);
return Err(error);
}
if rows_affected > 0 {
self.invalidate_copy_caches(table_name);
}
radixdb_storage::instrumentation::record_copy(
rows_affected.max(0) as u64,
parse_elapsed,
commit_elapsed,
copy_started.elapsed(),
);
rows_affected
}
Err(error) => {
transaction.rollback()?;
radixdb_storage::instrumentation::record_copy(
0,
parse_elapsed,
std::time::Duration::ZERO,
copy_started.elapsed(),
);
return Err(error);
}
};
Ok(Box::new(ExecResult::with_rows_affected(rows_affected)))
}
fn invalidate_copy_caches(&self, table_name: &str) {
self.mutation_invalidate_semantic_cache(table_name);
invalidate_semi_join_cache_for_table(table_name);
invalidate_scalar_subquery_cache_for_table(table_name);
invalidate_in_subquery_cache_for_table(table_name);
}
#[allow(clippy::too_many_arguments)]
fn copy_from_csv(
&self,
stmt: &CopyStatement,
table: &mut Box<dyn Table>,
schema: &Schema,
column_indices: &[usize],
all_column_types: &[DataType],
all_vector_dims: &[u16],
default_exprs: &[Option<String>],
fk_schema: &Option<CompactArc<Schema>>,
schema_column_count: usize,
copy_transaction_bytes: &mut usize,
copy_max_transaction_bytes: usize,
) -> Result<i64> {
let file = std::fs::File::open(&stmt.file_path).map_err(|e| {
Error::InvalidArgument(format!("cannot open file '{}': {}", stmt.file_path, e))
})?;
let mut reader = csv::ReaderBuilder::new()
.has_headers(stmt.header)
.delimiter(stmt.delimiter)
.from_reader(std::io::BufReader::new(file));
let field_to_col: Option<Vec<usize>> = if stmt.header && stmt.columns.is_empty() {
let headers = reader
.headers()
.map_err(|e| Error::InvalidArgument(format!("cannot read CSV headers: {}", e)))?;
let col_map = {
let m = schema.column_index_map();
m.clone()
};
let mut mapping = Vec::with_capacity(headers.len());
let mut seen_headers = FxHashSet::default();
for h in headers.iter() {
let lower = h.to_lowercase();
if !seen_headers.insert(lower.clone()) {
return Err(Error::InvalidArgument(format!(
"duplicate CSV header '{}'",
h
)));
}
if let Some(&idx) = col_map.get(lower.as_str()) {
mapping.push(idx);
} else {
return Err(Error::ColumnNotFound(h.to_string()));
}
}
Some(mapping)
} else {
None
};
let null_str = stmt.null_string.as_deref().unwrap_or("");
let mut rows_affected = 0i64;
let compiled_table_checks = compile_table_check_constraints(schema)?;
let mut table_check_vm = crate::expression::ExprVM::new();
let mut insert_batch = Vec::with_capacity(COPY_INSERT_BATCH_ROWS);
for result in reader.records() {
let record = result.map_err(|e| {
Error::InvalidArgument(format!(
"CSV parse error at row {}: {}",
rows_affected + 1,
e
))
})?;
let effective_indices = field_to_col.as_deref().unwrap_or(column_indices);
if record.len() != effective_indices.len() {
return Err(Error::InvalidArgument(format!(
"CSV row {} has {} fields but expected {}",
rows_affected + 1,
record.len(),
effective_indices.len()
)));
}
let mut row_values =
build_default_row(default_exprs, all_column_types, schema_column_count)?;
for (i, field) in record.iter().enumerate() {
let col_idx = effective_indices[i];
if field == null_str {
row_values[col_idx] = Value::null_unknown();
continue;
}
let target_type = all_column_types[col_idx];
let col_name = schema.columns[col_idx].name.as_str();
let value = parse_field(field, target_type, col_name)?;
validate_vector_dims(&value, target_type, all_vector_dims[col_idx])?;
row_values[col_idx] = value;
}
let mut row = Row::from_values(row_values);
prepare_insert_row_constraints(
table.as_mut(),
schema,
&compiled_table_checks,
&mut row,
&mut table_check_vm,
)?;
if let Some(ref fks) = fk_schema {
crate::mutation::foreign_key::check_parent_exists(
self.mutation_engine(),
table.txn_id(),
fks,
&row,
)?;
}
account_copy_transaction_row(
copy_transaction_bytes,
copy_max_transaction_bytes,
&row,
rows_affected + 1,
)?;
insert_batch.push(row);
if fk_schema.is_some() || insert_batch.len() == COPY_INSERT_BATCH_ROWS {
flush_copy_insert_batch(table, &mut insert_batch)?;
}
rows_affected += 1;
}
flush_copy_insert_batch(table, &mut insert_batch)?;
Ok(rows_affected)
}
#[allow(clippy::too_many_arguments)]
fn copy_from_json(
&self,
stmt: &CopyStatement,
table: &mut Box<dyn Table>,
schema: &Schema,
column_indices: &[usize],
all_column_types: &[DataType],
all_vector_dims: &[u16],
default_exprs: &[Option<String>],
fk_schema: &Option<CompactArc<Schema>>,
schema_column_count: usize,
copy_transaction_bytes: &mut usize,
copy_max_transaction_bytes: usize,
) -> Result<i64> {
let null_str = stmt.null_string.as_deref();
let use_columns = !stmt.columns.is_empty();
let compiled_table_checks = compile_table_check_constraints(schema)?;
let mut table_check_vm = crate::expression::ExprVM::new();
let mut insert_batch = Vec::with_capacity(COPY_INSERT_BATCH_ROWS);
let col_name_lower_map: FxHashMap<String, usize> = if use_columns {
stmt.columns
.iter()
.enumerate()
.map(|(i, column)| (column.value_lower.to_string(), column_indices[i]))
.collect()
} else {
schema
.columns
.iter()
.enumerate()
.map(|(idx, c)| (c.name.to_lowercase(), idx))
.collect()
};
let file = std::fs::File::open(&stmt.file_path).map_err(|e| {
Error::InvalidArgument(format!("cannot open file '{}': {}", stmt.file_path, e))
})?;
let reader = JsonArrayStripper::new(std::io::BufReader::new(file));
let stream = serde_json::Deserializer::from_reader(reader).into_iter::<serde_json::Value>();
let mut rows_affected = 0i64;
for (idx, result) in stream.enumerate() {
let item = result.map_err(|e| {
Error::InvalidArgument(format!("JSON parse error at object {}: {}", idx + 1, e))
})?;
let obj = item.as_object().ok_or_else(|| {
Error::InvalidArgument(format!("JSON item {} is not an object", idx + 1))
})?;
let row = self.prepare_json_row(
obj,
table,
schema,
default_exprs,
schema_column_count,
&col_name_lower_map,
use_columns,
all_column_types,
all_vector_dims,
null_str,
&compiled_table_checks,
&mut table_check_vm,
fk_schema,
copy_transaction_bytes,
copy_max_transaction_bytes,
rows_affected + 1,
)?;
insert_batch.push(row);
if fk_schema.is_some() || insert_batch.len() == COPY_INSERT_BATCH_ROWS {
flush_copy_insert_batch(table, &mut insert_batch)?;
}
rows_affected += 1;
}
flush_copy_insert_batch(table, &mut insert_batch)?;
Ok(rows_affected)
}
#[allow(clippy::too_many_arguments)]
fn prepare_json_row(
&self,
obj: &serde_json::Map<String, serde_json::Value>,
table: &mut Box<dyn Table>,
schema: &Schema,
default_exprs: &[Option<String>],
schema_column_count: usize,
col_name_lower_map: &FxHashMap<String, usize>,
use_columns: bool,
all_column_types: &[DataType],
all_vector_dims: &[u16],
null_str: Option<&str>,
compiled_table_checks: &[(String, crate::expression::SharedProgram)],
table_check_vm: &mut crate::expression::ExprVM,
fk_schema: &Option<CompactArc<Schema>>,
copy_transaction_bytes: &mut usize,
copy_max_transaction_bytes: usize,
row_number: i64,
) -> Result<Row> {
let mut normalized_obj = FxHashMap::default();
for (key, value) in obj {
let normalized = key.to_lowercase();
if normalized_obj.insert(normalized, value).is_some() {
return Err(Error::InvalidArgument(format!(
"duplicate JSON key '{}' after case normalization",
key
)));
}
}
for lower_key in normalized_obj.keys() {
if !col_name_lower_map.contains_key(lower_key) {
return Err(Error::ColumnNotFound(lower_key.clone()));
}
}
let mut row_values =
build_default_row(default_exprs, all_column_types, schema_column_count)?;
if use_columns {
for (lower_name, col_idx) in col_name_lower_map {
let target_type = all_column_types[*col_idx];
if let Some(v) = normalized_obj.get(lower_name) {
let value = json_value_to_radixdb(v, target_type, lower_name, null_str)?;
validate_vector_dims(&value, target_type, all_vector_dims[*col_idx])?;
row_values[*col_idx] = value;
}
}
} else {
for (lower_key, json_val) in normalized_obj {
let col_idx = col_name_lower_map[&lower_key];
let target_type = all_column_types[col_idx];
let value = json_value_to_radixdb(json_val, target_type, &lower_key, null_str)?;
validate_vector_dims(&value, target_type, all_vector_dims[col_idx])?;
row_values[col_idx] = value;
}
}
let mut row = Row::from_values(row_values);
prepare_insert_row_constraints(
&mut **table,
schema,
compiled_table_checks,
&mut row,
table_check_vm,
)?;
if let Some(ref fks) = fk_schema {
crate::mutation::foreign_key::check_parent_exists(
self.mutation_engine(),
table.txn_id(),
fks,
&row,
)?;
}
account_copy_transaction_row(
copy_transaction_bytes,
copy_max_transaction_bytes,
&row,
row_number,
)?;
Ok(row)
}
}
impl<T: MutationHost + ?Sized> CopyExecutorExt for T {}
fn build_default_row(
default_exprs: &[Option<String>],
all_column_types: &[DataType],
schema_column_count: usize,
) -> Result<Vec<Value>> {
let mut row = Vec::with_capacity(schema_column_count);
for i in 0..schema_column_count {
if let Some(ref default_expr) = default_exprs[i] {
let default_type = all_column_types[i];
row.push(evaluate_default_expr(default_expr, default_type)?);
} else {
row.push(Value::null_unknown());
}
}
Ok(row)
}
#[inline]
fn validate_vector_dims(value: &Value, target_type: DataType, expected_dims: u16) -> Result<()> {
if target_type == DataType::Vector && expected_dims > 0 {
if let Value::Extension(data) = value {
if data.first() == Some(&(DataType::Vector as u8)) {
let got_dim = u16::try_from((data.len() - 1) / 4).unwrap_or(u16::MAX);
if got_dim != expected_dims {
return Err(Error::VectorDimensionMismatch {
expected: expected_dims,
got: got_dim,
});
}
}
}
}
Ok(())
}
fn json_value_to_radixdb(
v: &serde_json::Value,
target_type: DataType,
col_name: &str,
null_str: Option<&str>,
) -> Result<Value> {
let val = match v {
serde_json::Value::Null => return Ok(Value::null_unknown()),
serde_json::Value::Bool(b) => Value::Boolean(*b),
serde_json::Value::Number(n) => {
Value::text(n.to_string())
}
serde_json::Value::String(s) => {
if let Some(ns) = null_str {
if s == ns {
return Ok(Value::null_unknown());
}
}
Value::text(s)
}
serde_json::Value::Object(_) | serde_json::Value::Array(_) => Value::text(v.to_string()),
};
val.try_coerce_to_type(target_type).map_err(|error| {
Error::Type(format!(
"cannot convert value '{}' to {:?} for column '{}': {}",
val, target_type, col_name, error
))
})
}
struct JsonArrayStripper<R> {
inner: R,
is_array: bool,
depth: u32,
in_string: bool,
escape: bool,
pending: Option<u8>,
expect_value: bool,
seen_value: bool,
array_closed: bool,
terminal_checked: bool,
}
impl<R: std::io::Read> JsonArrayStripper<R> {
fn new(mut inner: R) -> Self {
let mut first = [0u8; 1];
let (is_array, pending) = loop {
match inner.read(&mut first) {
Ok(1) if first[0].is_ascii_whitespace() => continue,
Ok(1) if first[0] == b'[' => break (true, None), Ok(1) => break (false, Some(first[0])), _ => break (false, None), }
};
JsonArrayStripper {
inner,
is_array,
depth: 0,
in_string: false,
escape: false,
pending,
expect_value: true,
seen_value: false,
array_closed: false,
terminal_checked: false,
}
}
}
impl<R: std::io::Read> std::io::Read for JsonArrayStripper<R> {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
if let Some(b) = self.pending.take() {
buf[0] = b;
if buf.len() == 1 {
return Ok(1);
}
let n = self.inner.read(&mut buf[1..])?;
return Ok(1 + n);
}
let n = self.inner.read(buf)?;
if self.is_array {
if n == 0 {
if !self.terminal_checked {
self.terminal_checked = true;
if !self.array_closed || self.in_string || self.depth != 0 {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"unterminated JSON array",
));
}
}
} else {
self.strip_array_syntax(buf, n)?;
}
}
Ok(n)
}
}
impl<R> JsonArrayStripper<R> {
fn strip_array_syntax(&mut self, buf: &mut [u8], len: usize) -> std::io::Result<()> {
for b in &mut buf[..len] {
if self.array_closed {
if !b.is_ascii_whitespace() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"trailing data after JSON array",
));
}
continue;
}
if self.in_string {
if self.escape {
self.escape = false;
} else if *b == b'\\' {
self.escape = true;
} else if *b == b'"' {
self.in_string = false;
}
continue;
}
match *b {
b'"' => {
if self.depth == 0 {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"JSON array COPY items must be objects",
));
}
self.in_string = true;
}
b'{' if self.depth == 0 => {
if !self.expect_value {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"missing comma between JSON array items",
));
}
self.expect_value = false;
self.seen_value = true;
self.depth = 1;
}
b'{' | b'[' => {
self.depth += 1;
}
b'}' | b']' => {
if self.depth > 0 {
self.depth -= 1;
} else if *b == b']' && (!self.expect_value || !self.seen_value) {
*b = b' ';
self.array_closed = true;
} else {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"invalid JSON array closing delimiter",
));
}
}
b',' if self.depth == 0 => {
if self.expect_value || !self.seen_value {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"unexpected comma in JSON array",
));
}
*b = b' ';
self.expect_value = true;
}
_ if self.depth == 0 && !b.is_ascii_whitespace() => {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"JSON array COPY items must be objects",
));
}
_ => {}
}
}
Ok(())
}
}