use bincode::{Decode, Encode};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Snapshot {
pub id: String,
pub collection_name: String,
pub created_at: DateTime<Utc>,
pub vectors_count: usize,
pub checksum: String,
pub size_bytes: u64,
}
#[derive(Debug, Serialize, Deserialize, Encode, Decode)]
pub struct SnapshotData {
pub metadata: SnapshotMetadata,
pub config: CollectionConfig,
pub vectors: Vec<VectorRecord>,
}
impl SnapshotData {
pub fn new(
collection_name: String,
config: CollectionConfig,
vectors: Vec<VectorRecord>,
) -> Self {
Self {
metadata: SnapshotMetadata {
id: uuid::Uuid::new_v4().to_string(),
collection_name,
created_at: Utc::now().to_rfc3339(),
version: env!("CARGO_PKG_VERSION").to_string(),
},
config,
vectors,
}
}
pub fn vectors_count(&self) -> usize {
self.vectors.len()
}
pub fn id(&self) -> &str {
&self.metadata.id
}
pub fn collection_name(&self) -> &str {
&self.metadata.collection_name
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Encode, Decode)]
pub struct SnapshotMetadata {
pub id: String,
pub collection_name: String,
pub created_at: String,
pub version: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, Encode, Decode)]
pub struct CollectionConfig {
pub dimension: usize,
pub metric: DistanceMetric,
pub hnsw_config: Option<HnswConfig>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Encode, Decode)]
pub enum DistanceMetric {
Cosine,
Euclidean,
DotProduct,
}
#[derive(Debug, Clone, Serialize, Deserialize, Encode, Decode)]
pub struct HnswConfig {
pub m: usize,
pub ef_construction: usize,
pub ef_search: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize, Encode, Decode)]
pub struct VectorRecord {
pub id: String,
pub vector: Vec<f32>,
#[serde(skip)]
#[bincode(with_serde)]
payload_json: Option<String>,
}
impl VectorRecord {
pub fn new(id: String, vector: Vec<f32>, payload: Option<Value>) -> Self {
let payload_json = payload.and_then(|v| serde_json::to_string(&v).ok());
Self {
id,
vector,
payload_json,
}
}
pub fn payload(&self) -> Option<Value> {
self.payload_json
.as_ref()
.and_then(|s| serde_json::from_str(s).ok())
}
pub fn set_payload(&mut self, payload: Option<Value>) {
self.payload_json = payload.and_then(|v| serde_json::to_string(&v).ok());
}
pub fn dimension(&self) -> usize {
self.vector.len()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_vector_record_creation() {
let record = VectorRecord::new("test-1".to_string(), vec![1.0, 2.0, 3.0], None);
assert_eq!(record.id, "test-1");
assert_eq!(record.dimension(), 3);
}
#[test]
fn test_snapshot_data_creation() {
let config = CollectionConfig {
dimension: 3,
metric: DistanceMetric::Cosine,
hnsw_config: None,
};
let vectors = vec![
VectorRecord::new("v1".to_string(), vec![1.0, 0.0, 0.0], None),
VectorRecord::new("v2".to_string(), vec![0.0, 1.0, 0.0], None),
];
let data = SnapshotData::new("test-collection".to_string(), config, vectors);
assert_eq!(data.vectors_count(), 2);
assert_eq!(data.collection_name(), "test-collection");
assert!(!data.id().is_empty());
}
}