use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use uuid::Uuid;
use super::vector::DenseVector;
use crate::StorageError;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum RecordId {
String(String),
U64(u64),
I64(i64),
Uuid(Uuid),
Opaque(String),
}
impl RecordId {
pub fn as_string_lossy(&self) -> String {
match self {
Self::String(v) => v.clone(),
Self::U64(v) => v.to_string(),
Self::I64(v) => v.to_string(),
Self::Uuid(v) => v.to_string(),
Self::Opaque(v) => v.clone(),
}
}
}
impl From<&str> for RecordId {
fn from(value: &str) -> Self {
Self::String(value.into())
}
}
impl From<String> for RecordId {
fn from(value: String) -> Self {
Self::String(value)
}
}
impl From<u64> for RecordId {
fn from(value: u64) -> Self {
Self::U64(value)
}
}
impl From<i64> for RecordId {
fn from(value: i64) -> Self {
Self::I64(value)
}
}
impl From<Uuid> for RecordId {
fn from(value: Uuid) -> Self {
Self::Uuid(value)
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
pub struct DataRecord {
pub id: Option<RecordId>,
pub vectors: BTreeMap<String, DenseVector>,
pub fields: serde_json::Map<String, Value>,
}
impl DataRecord {
#[must_use]
pub fn new(id: impl Into<RecordId>) -> Self {
Self {
id: Some(id.into()),
..Self::default()
}
}
#[must_use]
pub fn without_id() -> Self {
Self::default()
}
pub fn field<T>(mut self, key: impl Into<String>, value: T) -> Result<Self, StorageError>
where
T: Serialize,
{
let value = serde_json::to_value(value).map_err(|error| {
StorageError::InvalidConfig(format!("failed to serialize record field: {error}"))
})?;
self.fields.insert(key.into(), value);
Ok(self)
}
#[must_use]
pub fn json_field(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
self.fields.insert(key.into(), value);
self
}
#[must_use]
pub fn vector(mut self, key: impl Into<String>, vector: DenseVector) -> Self {
self.vectors.insert(key.into(), vector);
self
}
}