use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum PointId {
Uuid(String),
Num(u64),
}
impl From<&str> for PointId {
fn from(s: &str) -> Self {
PointId::Uuid(s.to_string())
}
}
impl From<String> for PointId {
fn from(s: String) -> Self {
PointId::Uuid(s)
}
}
impl From<u64> for PointId {
fn from(n: u64) -> Self {
PointId::Num(n)
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum PayloadValue {
String(String),
Integer(i64),
Float(f64),
Bool(bool),
List(Vec<PayloadValue>),
Object(HashMap<String, PayloadValue>),
Null,
}
pub type Payload = HashMap<String, PayloadValue>;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Point {
pub id: PointId,
pub vector: Vec<f32>,
#[serde(default)]
pub payload: Payload,
}
impl Point {
pub fn new(id: impl Into<PointId>, vector: Vec<f32>) -> Self {
Self {
id: id.into(),
vector,
payload: HashMap::new(),
}
}
pub fn new_num(id: u64, vector: Vec<f32>) -> Self {
Self {
id: PointId::Num(id),
vector,
payload: HashMap::new(),
}
}
pub fn with_payload(mut self, key: impl Into<String>, value: impl Into<PayloadValue>) -> Self {
self.payload.insert(key.into(), value.into());
self
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SparseVector {
pub indices: Vec<u32>,
pub values: Vec<f32>,
}
impl SparseVector {
pub fn new(indices: Vec<u32>, values: Vec<f32>) -> Self {
Self { indices, values }
}
pub fn from_dense(dense: &[f32], threshold: f32) -> Self {
let mut indices = Vec::new();
let mut values = Vec::new();
for (i, &v) in dense.iter().enumerate() {
if v.abs() > threshold {
indices.push(i as u32);
values.push(v);
}
}
Self { indices, values }
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum VectorData {
Single(Vec<f32>),
Named(HashMap<String, Vec<f32>>),
Sparse {
indices: Vec<u32>,
values: Vec<f32>,
},
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MultiVectorPoint {
pub id: PointId,
pub vectors: HashMap<String, Vec<f32>>,
#[serde(default)]
pub payload: Payload,
}
impl MultiVectorPoint {
pub fn new(id: impl Into<PointId>) -> Self {
Self {
id: id.into(),
vectors: HashMap::new(),
payload: HashMap::new(),
}
}
pub fn with_vector(mut self, name: impl Into<String>, vector: Vec<f32>) -> Self {
self.vectors.insert(name.into(), vector);
self
}
pub fn with_payload(mut self, key: impl Into<String>, value: impl Into<PayloadValue>) -> Self {
self.payload.insert(key.into(), value.into());
self
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ScoredPoint {
pub id: PointId,
pub score: f32,
#[serde(default)]
pub payload: Payload,
#[serde(default)]
pub vector: Option<Vec<f32>>,
}
impl From<String> for PayloadValue {
fn from(s: String) -> Self {
PayloadValue::String(s)
}
}
impl From<&str> for PayloadValue {
fn from(s: &str) -> Self {
PayloadValue::String(s.to_string())
}
}
impl From<i64> for PayloadValue {
fn from(n: i64) -> Self {
PayloadValue::Integer(n)
}
}
impl From<f64> for PayloadValue {
fn from(n: f64) -> Self {
PayloadValue::Float(n)
}
}
impl From<bool> for PayloadValue {
fn from(b: bool) -> Self {
PayloadValue::Bool(b)
}
}