#![allow(
clippy::cast_possible_truncation,
clippy::cast_precision_loss,
clippy::if_not_else,
clippy::single_match_else,
clippy::needless_raw_string_hashes
)]
use anyhow::{Context, Result};
use indicatif::{ProgressBar, ProgressStyle};
use serde::Deserialize;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::path::Path;
use velesdb_core::{Database, DistanceMetric, Point, StorageMode};
pub struct ImportConfig {
pub collection: String,
pub dimension: Option<usize>,
pub metric: DistanceMetric,
pub storage_mode: StorageMode,
pub batch_size: usize,
pub id_column: String,
pub vector_column: String,
pub show_progress: bool,
}
impl Default for ImportConfig {
fn default() -> Self {
Self {
collection: String::new(),
dimension: None,
metric: DistanceMetric::Cosine,
storage_mode: StorageMode::Full,
batch_size: 1000,
id_column: "id".to_string(),
vector_column: "vector".to_string(),
show_progress: true,
}
}
}
#[derive(Debug, Deserialize)]
struct JsonRecord {
id: u64,
vector: Vec<f32>,
#[serde(default)]
payload: Option<serde_json::Value>,
}
pub fn import_jsonl(db: &Database, path: &Path, config: &ImportConfig) -> Result<ImportStats> {
let file = File::open(path).context("Failed to open JSONL file")?;
let file_size = file.metadata()?.len();
let total = BufReader::with_capacity(64 * 1024, &file).lines().count();
if total == 0 {
anyhow::bail!("Empty file");
}
let file = File::open(path)?;
let mut reader = BufReader::with_capacity(128 * 1024, file);
let mut first_line = String::new();
reader.read_line(&mut first_line)?;
let first_record: JsonRecord =
serde_json::from_str(&first_line).context("Failed to parse first line")?;
let dimension = config.dimension.unwrap_or(first_record.vector.len());
let collection = get_or_create_collection(
db,
&config.collection,
dimension,
config.metric,
config.storage_mode,
)?;
let progress = create_progress_bar(total, config.show_progress);
if config.show_progress {
progress.set_message(format!(
"Importing {} vectors ({:.1} MB)",
total,
file_size as f64 / (1024.0 * 1024.0)
));
}
let mut stats = ImportStats::default();
let start = std::time::Instant::now();
let mut batch: Vec<Point> = Vec::with_capacity(config.batch_size);
if first_record.vector.len() == dimension {
batch.push(Point::new(
first_record.id,
first_record.vector,
first_record.payload,
));
stats.imported += 1;
} else {
stats.errors += 1;
}
progress.inc(1);
let mut line = String::with_capacity(dimension * 10); while reader.read_line(&mut line)? > 0 {
match serde_json::from_str::<JsonRecord>(&line) {
Ok(record) => {
if record.vector.len() != dimension {
stats.errors += 1;
} else {
batch.push(Point::new(record.id, record.vector, record.payload));
stats.imported += 1;
if batch.len() >= config.batch_size {
collection.upsert_bulk(&batch)?;
batch.clear();
}
}
}
Err(_) => {
stats.errors += 1;
}
}
line.clear(); progress.inc(1);
}
if !batch.is_empty() {
collection.upsert_bulk(&batch)?;
}
progress.finish_with_message("Import complete");
stats.duration_ms = start.elapsed().as_millis() as u64;
stats.total = total;
Ok(stats)
}
#[allow(clippy::too_many_lines)]
pub fn import_csv(db: &Database, path: &Path, config: &ImportConfig) -> Result<ImportStats> {
let file = File::open(path).context("Failed to open CSV file")?;
let file_size = file.metadata()?.len();
let buffered = BufReader::with_capacity(128 * 1024, file);
let mut reader = csv::Reader::from_reader(buffered);
let headers = reader.headers()?.clone();
let id_idx = headers
.iter()
.position(|h| h == config.id_column)
.context(format!("ID column '{}' not found", config.id_column))?;
let vector_idx = headers
.iter()
.position(|h| h == config.vector_column)
.context(format!(
"Vector column '{}' not found",
config.vector_column
))?;
let total = reader.records().count();
if total == 0 {
anyhow::bail!("Empty file");
}
let file = File::open(path)?;
let buffered = BufReader::with_capacity(128 * 1024, file);
let mut reader = csv::Reader::from_reader(buffered);
let first_record = reader.records().next().context("No records in CSV")??;
let vector_str = &first_record[vector_idx];
let first_vector = parse_vector(vector_str)?;
let dimension = config.dimension.unwrap_or(first_vector.len());
let collection = get_or_create_collection(
db,
&config.collection,
dimension,
config.metric,
config.storage_mode,
)?;
let file = File::open(path)?;
let buffered = BufReader::with_capacity(128 * 1024, file);
let mut reader = csv::Reader::from_reader(buffered);
let progress = create_progress_bar(total, config.show_progress);
if config.show_progress {
progress.set_message(format!(
"Importing {} vectors ({:.1} MB)",
total,
file_size as f64 / (1024.0 * 1024.0)
));
}
let mut stats = ImportStats::default();
let start = std::time::Instant::now();
let mut batch: Vec<Point> = Vec::with_capacity(config.batch_size);
for result in reader.records() {
match result {
Ok(record) => {
let id: u64 = match record[id_idx].parse() {
Ok(id) => id,
Err(_) => {
stats.errors += 1;
progress.inc(1);
continue;
}
};
match parse_vector(&record[vector_idx]) {
Ok(vector) => {
if vector.len() != dimension {
stats.errors += 1;
progress.inc(1);
continue;
}
let mut payload = serde_json::Map::new();
for (i, header) in headers.iter().enumerate() {
if i != id_idx && i != vector_idx {
payload.insert(
header.to_string(),
serde_json::Value::String(record[i].to_string()),
);
}
}
let payload_val = if payload.is_empty() {
None
} else {
Some(serde_json::Value::Object(payload))
};
batch.push(Point::new(id, vector, payload_val));
stats.imported += 1;
if batch.len() >= config.batch_size {
collection.upsert_bulk(&batch)?;
batch.clear();
}
}
Err(_) => {
stats.errors += 1;
}
}
}
Err(_) => {
stats.errors += 1;
}
}
progress.inc(1);
}
if !batch.is_empty() {
collection.upsert_bulk(&batch)?;
}
progress.finish_with_message("Import complete");
stats.duration_ms = start.elapsed().as_millis() as u64;
stats.total = total;
Ok(stats)
}
fn parse_vector(s: &str) -> Result<Vec<f32>> {
let s = s.trim();
if s.starts_with('[') {
serde_json::from_str(s).context("Invalid JSON vector")
} else {
s.split(',')
.map(|v| v.trim().parse::<f32>().context("Invalid float value"))
.collect()
}
}
fn get_or_create_collection(
db: &Database,
name: &str,
dimension: usize,
metric: DistanceMetric,
storage_mode: StorageMode,
) -> Result<velesdb_core::VectorCollection> {
if let Some(col) = db.get_vector_collection(name) {
Ok(col)
} else {
db.create_vector_collection_with_options(name, dimension, metric, storage_mode)?;
db.get_vector_collection(name)
.context("Failed to get created collection")
}
}
fn create_progress_bar(total: usize, show: bool) -> ProgressBar {
if show {
let pb = ProgressBar::new(total as u64);
pb.set_style(
ProgressStyle::default_bar()
.template(
"{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} ({eta})",
)
.unwrap()
.progress_chars("#>-"),
);
pb
} else {
ProgressBar::hidden()
}
}
#[derive(Debug, Default)]
pub struct ImportStats {
pub total: usize,
pub imported: usize,
pub errors: usize,
pub duration_ms: u64,
}
impl ImportStats {
#[allow(clippy::cast_precision_loss)]
#[must_use]
pub fn records_per_sec(&self) -> f64 {
if self.duration_ms == 0 {
0.0
} else {
(self.imported as f64) / (self.duration_ms as f64 / 1000.0)
}
}
}
#[cfg(test)]
#[allow(
clippy::cast_precision_loss,
clippy::float_cmp,
clippy::redundant_closure_for_method_calls
)]
#[path = "import_tests.rs"]
mod import_tests;