use std::fmt;
use crate::common::counter::hardware_counter::HardwareCounterCell;
use crate::segment::common::operation_error::{OperationError, OperationResult};
use crate::shard::operations::vector_name_ops::VectorNameConfig;
use crate::shard::operations::{CollectionUpdateOperations, VectorNameOperations};
use crate::shard::update::*;
use crate::shard::wal::WalRawRecord;
use crate::edge::EdgeShard;
use crate::edge::config::vectors::{EdgeSparseVectorParams, EdgeVectorParams};
impl EdgeShard {
pub fn update(&self, operation: CollectionUpdateOperations) -> OperationResult<()> {
if let CollectionUpdateOperations::VectorNameOperation(
VectorNameOperations::CreateVectorName(create),
) = &operation
{
self.check_vector_name_create_compatible(&create.vector_name, &create.config)?;
}
let record = WalRawRecord::new(&operation).map_err(service_error)?;
let mut wal = self.wal.lock();
let operation_id = wal.write(&record).map_err(service_error)?;
let hw_counter = HardwareCounterCell::disposable();
let _update_guard = self.segments.acquire_updates_lock();
let segments_guard = self.segments.read();
let result = match operation {
CollectionUpdateOperations::PointOperation(point_operation) => {
process_point_operation(&segments_guard, operation_id, point_operation, &hw_counter)
}
CollectionUpdateOperations::VectorOperation(vector_operation) => {
process_vector_operation(
&segments_guard,
operation_id,
vector_operation,
&hw_counter,
)
}
CollectionUpdateOperations::PayloadOperation(payload_operation) => {
process_payload_operation(
&segments_guard,
operation_id,
payload_operation,
&hw_counter,
)
}
CollectionUpdateOperations::FieldIndexOperation(index_operation) => {
process_field_index_operation(
&segments_guard,
operation_id,
&index_operation,
&hw_counter,
)
}
CollectionUpdateOperations::VectorNameOperation(ref vector_name_operation) => {
let result = process_vector_name_operation(
&segments_guard,
operation_id,
vector_name_operation,
);
if let Ok(changed) = result {
self.apply_vector_name_to_config(vector_name_operation, changed)?;
}
result
}
#[cfg(feature = "staging")]
CollectionUpdateOperations::StagingOperation(staging_operation) => {
crate::shard::update::process_staging_operation(
&segments_guard,
operation_id,
staging_operation,
)
}
};
result.map(|_| ())
}
fn check_vector_name_create_compatible(
&self,
name: &str,
config: &VectorNameConfig,
) -> OperationResult<()> {
let conflict = match requested_vector_params(config) {
RequestedVectorParams::Dense(params) => self
.config
.read()
.vectors
.get(name)
.is_some_and(|existing| !dense_identity_matches(existing, ¶ms)),
RequestedVectorParams::Sparse(params) => self
.config
.read()
.sparse_vectors
.get(name)
.is_some_and(|existing| !sparse_identity_matches(existing, ¶ms)),
};
if conflict {
return Err(OperationError::validation_error(format!(
"vector '{name}' already exists with different parameters; delete \
it before recreating with a different configuration"
)));
}
Ok(())
}
fn apply_vector_name_to_config(
&self,
operation: &VectorNameOperations,
changed: usize,
) -> OperationResult<()> {
match operation {
VectorNameOperations::CreateVectorName(create) => {
if changed == 0 {
return Ok(());
}
self.config
.write(|config| match requested_vector_params(&create.config) {
RequestedVectorParams::Dense(params) => {
config.vectors.insert(create.vector_name.clone(), params);
}
RequestedVectorParams::Sparse(params) => {
config
.sparse_vectors
.insert(create.vector_name.clone(), params);
}
})
.map_err(service_error)?;
}
VectorNameOperations::DeleteVectorName(delete) => {
self.config
.write(|config| {
config.vectors.remove(&delete.vector_name);
config.sparse_vectors.remove(&delete.vector_name);
})
.map_err(service_error)?;
}
}
Ok(())
}
}
enum RequestedVectorParams {
Dense(EdgeVectorParams),
Sparse(EdgeSparseVectorParams),
}
fn requested_vector_params(config: &VectorNameConfig) -> RequestedVectorParams {
match config {
VectorNameConfig::Dense(wrapper) => RequestedVectorParams::Dense(EdgeVectorParams {
size: wrapper.dense.size,
distance: wrapper.dense.distance,
on_disk: None,
multivector_config: wrapper.dense.multivector_config,
datatype: wrapper.dense.datatype,
quantization_config: None,
hnsw_config: None,
}),
VectorNameConfig::Sparse(wrapper) => {
RequestedVectorParams::Sparse(EdgeSparseVectorParams {
full_scan_threshold: None,
on_disk: None,
modifier: wrapper.sparse.modifier,
datatype: wrapper.sparse.datatype,
})
}
}
}
fn dense_identity_matches(existing: &EdgeVectorParams, requested: &EdgeVectorParams) -> bool {
existing.size == requested.size
&& existing.distance == requested.distance
&& existing.multivector_config == requested.multivector_config
&& existing.datatype == requested.datatype
}
fn sparse_identity_matches(
existing: &EdgeSparseVectorParams,
requested: &EdgeSparseVectorParams,
) -> bool {
existing.modifier == requested.modifier && existing.datatype == requested.datatype
}
fn service_error(err: impl fmt::Display) -> OperationError {
OperationError::service_error(err.to_string())
}