use crate::segment::common::operation_error::{OperationError, OperationResult};
use crate::segment::data_types::fully_qualified_point::{FullyQualifiedPoint, StoredPoint};
use crate::segment::data_types::named_vectors::NamedVectors;
use crate::segment::data_types::segment_record::NamedVectorBytesOwned;
use crate::segment::json_path::JsonPath;
use crate::segment::types::{Payload, PayloadKeyType, PointIdType, SeqNumberType, VectorNameBuf};
pub enum OperationVectors {
Decoded(NamedVectors<'static>),
Raw(NamedVectorBytesOwned),
}
pub enum PointMutation {
Replace {
vectors: OperationVectors,
payload: Payload,
},
Delete,
UpdateVectors(NamedVectors<'static>),
DeleteVectors(Vec<VectorNameBuf>),
SetPayload {
payload: Payload,
key: Option<JsonPath>,
},
OverwritePayload(Payload),
DeletePayload(Vec<PayloadKeyType>),
ClearPayload,
}
impl PointMutation {
fn discards_stored_point(&self) -> bool {
match self {
Self::Replace { .. } | Self::Delete => true,
Self::UpdateVectors(_)
| Self::DeleteVectors(_)
| Self::SetPayload { .. }
| Self::OverwritePayload(_)
| Self::DeletePayload(_)
| Self::ClearPayload => false,
}
}
}
pub struct PointUpdates {
version: SeqNumberType,
mutations: Vec<PointMutation>,
}
impl PointUpdates {
pub(super) fn new(version: SeqNumberType, mutation: PointMutation) -> Self {
Self {
version,
mutations: vec![mutation],
}
}
pub(super) fn push(&mut self, version: SeqNumberType, mutation: PointMutation) {
if mutation.discards_stored_point() {
self.mutations.clear();
}
self.version = self.version.max(version);
self.mutations.push(mutation);
}
pub fn version(&self) -> SeqNumberType {
self.version
}
pub fn needs_stored_point(&self) -> bool {
self.mutations
.first()
.is_none_or(|mutation| !mutation.discards_stored_point())
}
pub fn materialize(
self,
id: PointIdType,
stored: Option<StoredPoint>,
) -> OperationResult<Option<FullyQualifiedPoint>> {
let Self { version, mutations } = self;
let mut exists = stored.is_some();
let (mut stored_vectors, mut payload) = match stored {
Some(stored) => {
let StoredPoint {
internal_id: _,
vectors,
payload,
} = stored;
(vectors, payload)
}
None => (NamedVectorBytesOwned::new(), Payload::default()),
};
let mut updated_vectors = NamedVectors::default();
for mutation in mutations {
match mutation {
PointMutation::Replace {
vectors,
payload: replacement,
} => {
exists = true;
stored_vectors.clear();
updated_vectors = NamedVectors::default();
match vectors {
OperationVectors::Decoded(vectors) => updated_vectors = vectors,
OperationVectors::Raw(vectors) => stored_vectors = vectors,
}
payload = replacement;
}
PointMutation::Delete => {
exists = false;
stored_vectors.clear();
updated_vectors = NamedVectors::default();
payload = Payload::default();
}
PointMutation::UpdateVectors(vectors) => {
if !exists {
return Err(OperationError::PointIdError {
missed_point_id: id,
});
}
updated_vectors.merge(vectors);
}
PointMutation::DeleteVectors(names) => {
for name in &names {
stored_vectors.retain(|(stored_name, _)| stored_name != name);
updated_vectors.remove_ref(name.as_str());
}
}
PointMutation::SetPayload {
payload: values,
key,
} => match key {
Some(key) => payload.merge_by_key(&values, &key),
None => payload.merge(&values),
},
PointMutation::OverwritePayload(values) => payload = values,
PointMutation::DeletePayload(keys) => {
for key in &keys {
payload.remove(key);
}
}
PointMutation::ClearPayload => payload = Payload::default(),
}
}
if !exists {
return Ok(None);
}
Ok(Some(FullyQualifiedPoint {
id,
version,
stored_vectors,
updated_vectors,
payload,
}))
}
}