use crate::segment::data_types::modifier::Modifier;
use crate::segment::index::sparse_index::sparse_index_config::{SparseIndexConfig, SparseIndexType};
use crate::segment::types::{
Distance, HnswConfig, Indexes, MultiVectorConfig, QuantizationConfig, SparseVectorDataConfig,
SparseVectorStorageType, VectorDataConfig, VectorStorageDatatype, VectorStorageType,
};
use serde::{Deserialize, Serialize};
use crate::shard::optimizers::config::DenseVectorOptimizerConfig;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct EdgeVectorParams {
pub size: usize,
pub distance: Distance,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub on_disk: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub multivector_config: Option<MultiVectorConfig>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub datatype: Option<VectorStorageDatatype>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub quantization_config: Option<QuantizationConfig>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub hnsw_config: Option<HnswConfig>,
}
impl EdgeVectorParams {
pub fn builder(size: usize, distance: Distance) -> crate::edge::builders::EdgeVectorParamsBuilder {
crate::edge::builders::EdgeVectorParamsBuilder::new(size, distance)
}
pub fn to_plain_vector_data_config(
&self,
global_quantization: Option<&QuantizationConfig>,
) -> VectorDataConfig {
let EdgeVectorParams {
size,
distance,
on_disk,
multivector_config,
datatype,
quantization_config,
hnsw_config: _, } = self;
let resolved_quantization_config = quantization_config.as_ref().or(global_quantization);
let quantization_config =
QuantizationConfig::for_appendable_segment(resolved_quantization_config);
VectorDataConfig {
size: *size,
distance: *distance,
storage_type: VectorStorageType::from_on_disk(on_disk.unwrap_or_default()),
index: Indexes::Plain {},
quantization_config,
multivector_config: *multivector_config,
datatype: *datatype,
}
}
pub fn to_dense_vector_optimizer_config(
&self,
global_hnsw_config: &HnswConfig,
global_quantization_config: Option<&QuantizationConfig>,
) -> DenseVectorOptimizerConfig {
let EdgeVectorParams {
size: _,
distance: _,
on_disk,
multivector_config: _,
datatype: _,
quantization_config,
hnsw_config,
} = self;
DenseVectorOptimizerConfig {
on_disk: *on_disk,
memory: None,
hnsw_config: hnsw_config.unwrap_or(*global_hnsw_config),
quantization_config: quantization_config
.clone()
.or_else(|| global_quantization_config.cloned()),
}
}
pub fn from_vector_data_config(v: &VectorDataConfig) -> Self {
let VectorDataConfig {
size,
distance,
storage_type,
index,
quantization_config, multivector_config,
datatype,
} = v;
Self {
size: *size,
distance: *distance,
on_disk: Some(storage_type.is_on_disk()),
multivector_config: *multivector_config,
datatype: *datatype,
quantization_config: quantization_config.clone(),
hnsw_config: match index {
Indexes::Plain {} => None,
Indexes::Hnsw(hnsw_config) => Some(*hnsw_config),
},
}
}
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct EdgeSparseVectorParams {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub full_scan_threshold: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub on_disk: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub modifier: Option<Modifier>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub datatype: Option<VectorStorageDatatype>,
}
impl EdgeSparseVectorParams {
pub fn builder() -> crate::edge::builders::EdgeSparseVectorParamsBuilder {
crate::edge::builders::EdgeSparseVectorParamsBuilder::new()
}
pub fn to_plain_sparse_vector_data_config(&self) -> SparseVectorDataConfig {
let EdgeSparseVectorParams {
full_scan_threshold,
on_disk: _,
modifier,
datatype,
} = self;
SparseVectorDataConfig {
index: SparseIndexConfig {
full_scan_threshold: *full_scan_threshold,
index_type: SparseIndexType::default(),
datatype: *datatype,
memory: None,
},
storage_type: SparseVectorStorageType::Mmap,
modifier: *modifier,
}
}
pub fn to_sparse_vector_optimizer_config(
&self,
) -> crate::shard::optimizers::config::SparseVectorOptimizerConfig {
let EdgeSparseVectorParams {
full_scan_threshold: _,
on_disk,
modifier: _,
datatype: _,
} = self;
crate::shard::optimizers::config::SparseVectorOptimizerConfig {
on_disk: *on_disk,
memory: None,
}
}
pub fn from_sparse_vector_data_config(s: &SparseVectorDataConfig) -> Self {
let SparseVectorDataConfig {
index,
storage_type: _, modifier,
} = s;
let SparseIndexConfig {
full_scan_threshold,
index_type,
datatype,
memory: _,
} = index;
Self {
full_scan_threshold: *full_scan_threshold,
on_disk: Some(index_type.is_on_disk()),
modifier: *modifier,
datatype: *datatype,
}
}
}