use std::collections::HashMap;
use crate::segment::types::{HnswConfig, QuantizationConfig, VectorNameBuf};
use crate::wal::WalOptions;
use crate::edge::config::optimizers::EdgeOptimizersConfig;
use crate::edge::config::shard::EdgeConfig;
use crate::edge::config::vectors::{EdgeSparseVectorParams, EdgeVectorParams};
#[derive(Debug, Default)]
pub struct EdgeConfigBuilder {
on_disk_payload: Option<bool>,
vectors: HashMap<VectorNameBuf, EdgeVectorParams>,
sparse_vectors: HashMap<VectorNameBuf, EdgeSparseVectorParams>,
hnsw_config: Option<HnswConfig>,
quantization_config: Option<QuantizationConfig>,
optimizers: Option<EdgeOptimizersConfig>,
wal_options: Option<WalOptions>,
}
impl EdgeConfigBuilder {
pub fn new() -> Self {
Self::default()
}
pub fn vector(mut self, name: impl Into<VectorNameBuf>, params: EdgeVectorParams) -> Self {
self.vectors.insert(name.into(), params);
self
}
pub fn vectors(mut self, vectors: HashMap<VectorNameBuf, EdgeVectorParams>) -> Self {
self.vectors = vectors;
self
}
pub fn sparse_vector(
mut self,
name: impl Into<VectorNameBuf>,
params: EdgeSparseVectorParams,
) -> Self {
self.sparse_vectors.insert(name.into(), params);
self
}
pub fn sparse_vectors(
mut self,
sparse_vectors: HashMap<VectorNameBuf, EdgeSparseVectorParams>,
) -> Self {
self.sparse_vectors = sparse_vectors;
self
}
pub fn on_disk_payload(mut self, on_disk_payload: bool) -> Self {
self.on_disk_payload = Some(on_disk_payload);
self
}
pub fn hnsw_config(mut self, hnsw_config: HnswConfig) -> Self {
self.hnsw_config = Some(hnsw_config);
self
}
pub fn quantization_config(mut self, quantization_config: QuantizationConfig) -> Self {
self.quantization_config = Some(quantization_config);
self
}
pub fn optimizers(mut self, optimizers: EdgeOptimizersConfig) -> Self {
self.optimizers = Some(optimizers);
self
}
pub fn wal_options(mut self, wal_options: WalOptions) -> Self {
self.wal_options = Some(wal_options);
self
}
pub fn build(self) -> EdgeConfig {
let Self {
on_disk_payload,
vectors,
sparse_vectors,
hnsw_config,
quantization_config,
optimizers,
wal_options,
} = self;
let defaults = EdgeConfig::default();
EdgeConfig {
on_disk_payload: on_disk_payload.unwrap_or(defaults.on_disk_payload),
vectors,
sparse_vectors,
hnsw_config: hnsw_config.unwrap_or(defaults.hnsw_config),
quantization_config: quantization_config.or(defaults.quantization_config),
optimizers: optimizers.unwrap_or(defaults.optimizers),
wal_options: wal_options.or(defaults.wal_options),
}
}
}