From 75434ca565f96494deaf1858318fe8aa71edab9f Mon Sep 17 00:00:00 2001 From: Andrey Vasnetsov Date: Mon, 20 Oct 2025 12:39:01 +0200 Subject: [PATCH] implement all point update operation constructors (#7424) --- lib/edge/python/src/lib.rs | 2 +- lib/edge/python/src/types/mod.rs | 2 + lib/edge/python/src/types/point_vectors.rs | 23 ++ lib/edge/python/src/update.rs | 287 ++++++++++++++++++++- lib/edge/python/src/update/upsert.rs | 42 --- 5 files changed, 311 insertions(+), 45 deletions(-) create mode 100644 lib/edge/python/src/types/point_vectors.rs delete mode 100644 lib/edge/python/src/update/upsert.rs diff --git a/lib/edge/python/src/lib.rs b/lib/edge/python/src/lib.rs index cca1ea78b0..362ef53d71 100644 --- a/lib/edge/python/src/lib.rs +++ b/lib/edge/python/src/lib.rs @@ -48,7 +48,7 @@ mod qdrant_edge { #[pymodule_export] use super::types::filter::{PyFilter, field_condition::PyFieldCondition}; #[pymodule_export] - use super::types::{PyPoint, PyRecord, PySparseVector}; + use super::types::{PyPoint, PyPointVectors, PyRecord, PySparseVector}; #[pymodule_export] use super::update::PyUpdateOperation; } diff --git a/lib/edge/python/src/types/mod.rs b/lib/edge/python/src/types/mod.rs index 45bff82708..93cf4eecde 100644 --- a/lib/edge/python/src/types/mod.rs +++ b/lib/edge/python/src/types/mod.rs @@ -2,6 +2,7 @@ pub mod filter; pub mod payload; pub mod point; pub mod point_id; +pub mod point_vectors; pub mod record; pub mod vector; @@ -9,5 +10,6 @@ pub use filter::*; pub use payload::*; pub use point::*; pub use point_id::*; +pub use point_vectors::*; pub use record::*; pub use vector::*; diff --git a/lib/edge/python/src/types/point_vectors.rs b/lib/edge/python/src/types/point_vectors.rs new file mode 100644 index 0000000000..a4ad39048d --- /dev/null +++ b/lib/edge/python/src/types/point_vectors.rs @@ -0,0 +1,23 @@ +use derive_more::Into; +use pyo3::{pyclass, pymethods}; +use segment::types::PointIdType; +use shard::operations::point_ops::VectorStructPersisted; +use shard::operations::vector_ops::PointVectorsPersisted; + +use crate::types::{PyPointId, PyVector}; + +#[pyclass(name = "PointVectors")] +#[derive(Clone, Debug, Into)] +#[repr(transparent)] +pub struct PyPointVectors(pub PointVectorsPersisted); + +#[pymethods] +impl PyPointVectors { + #[new] + pub fn new(id: PyPointId, vector: PyVector) -> Self { + Self(PointVectorsPersisted { + id: PointIdType::from(id), + vector: VectorStructPersisted::from(vector), + }) + } +} diff --git a/lib/edge/python/src/update.rs b/lib/edge/python/src/update.rs index 44a3e329b2..d1cff4111e 100644 --- a/lib/edge/python/src/update.rs +++ b/lib/edge/python/src/update.rs @@ -1,9 +1,292 @@ -pub mod upsert; +use std::str::FromStr; use derive_more::Into; +use pyo3::exceptions::PyValueError; use pyo3::prelude::*; -use shard::operations::CollectionUpdateOperations; +use segment::json_path::JsonPath; +use segment::types::{Filter, Payload, VectorNameBuf}; +use shard::operations::point_ops::{ + PointIdsList, PointInsertOperationsInternal, PointStructPersisted, +}; +use shard::operations::vector_ops::PointVectorsPersisted; +use shard::operations::{CollectionUpdateOperations, payload_ops, point_ops, vector_ops}; + +use crate::types::{PyFilter, PyPayload, PyPoint, PyPointId, PyPointVectors}; #[pyclass(name = "UpdateOperation")] #[derive(Clone, Debug, Into)] pub struct PyUpdateOperation(CollectionUpdateOperations); + +#[pymethods] // Can't split impl block due to pyo3 limitations, so all constructors go here +impl PyUpdateOperation { + #[staticmethod] + pub fn upsert_points(points: Vec) -> Self { + let points = points.into_iter().map(PointStructPersisted::from).collect(); + + let operation = + CollectionUpdateOperations::PointOperation(point_ops::PointOperations::UpsertPoints( + PointInsertOperationsInternal::PointsList(points), + )); + + Self(operation) + } + + #[staticmethod] + pub fn upsert_points_conditional(points: Vec, condition: PyFilter) -> Self { + let points = points.into_iter().map(PointStructPersisted::from).collect(); + let points_op = PointInsertOperationsInternal::PointsList(points); + + let condition = Filter::from(condition); + + let operation = CollectionUpdateOperations::PointOperation( + point_ops::PointOperations::UpsertPointsConditional( + point_ops::ConditionalInsertOperationInternal { + points_op, + condition, + }, + ), + ); + + Self(operation) + } + + #[staticmethod] + pub fn delete_points(ids: Vec) -> Self { + let point_ids = PyPointId::into_rust_vec(ids); + + let operation = + CollectionUpdateOperations::PointOperation(point_ops::PointOperations::DeletePoints { + ids: point_ids, + }); + + Self(operation) + } + + #[staticmethod] + pub fn delete_points_by_filter(filter: PyFilter) -> Self { + let filter = Filter::from(filter); + + let operation = CollectionUpdateOperations::PointOperation( + point_ops::PointOperations::DeletePointsByFilter(filter), + ); + + Self(operation) + } + + #[staticmethod] + pub fn update_vectors(point_vectors: Vec) -> Self { + let points = point_vectors + .into_iter() + .map(PointVectorsPersisted::from) + .collect(); + + let operation = CollectionUpdateOperations::VectorOperation( + vector_ops::VectorOperations::UpdateVectors(vector_ops::UpdateVectorsOp { + points, + update_filter: None, + }), + ); + + Self(operation) + } + + #[staticmethod] + pub fn update_vectors_conditional( + point_vectors: Vec, + filter: PyFilter, + ) -> Self { + let points = point_vectors + .into_iter() + .map(PointVectorsPersisted::from) + .collect(); + let filter = Filter::from(filter); + let operation = CollectionUpdateOperations::VectorOperation( + vector_ops::VectorOperations::UpdateVectors(vector_ops::UpdateVectorsOp { + points, + update_filter: Some(filter), + }), + ); + Self(operation) + } + + #[staticmethod] + pub fn delete_vectors(ids: Vec, vector_names: Vec) -> Self { + let point_ids = PyPointId::into_rust_vec(ids); + let operation = CollectionUpdateOperations::VectorOperation( + vector_ops::VectorOperations::DeleteVectors( + PointIdsList::from(point_ids), + vector_names, + ), + ); + Self(operation) + } + + #[staticmethod] + pub fn delete_vectors_by_filter(filter: PyFilter, vector_names: Vec) -> Self { + let filter = Filter::from(filter); + let operation = CollectionUpdateOperations::VectorOperation( + vector_ops::VectorOperations::DeleteVectorsByFilter(filter, vector_names), + ); + Self(operation) + } + + #[staticmethod] + #[pyo3(signature = (ids, payload, key=None))] + pub fn set_payload( + ids: Vec, + payload: PyPayload, + key: Option, + ) -> Result { + let point_ids = PyPointId::into_rust_vec(ids); + let payload = Payload::from(payload); + + let key = key + .map(|k| JsonPath::from_str(&k).map_err(|_| PyErr::new::(k))) + .transpose()?; + + let operation = CollectionUpdateOperations::PayloadOperation( + payload_ops::PayloadOps::SetPayload(payload_ops::SetPayloadOp { + payload, + points: Some(point_ids), + filter: None, + key, + }), + ); + + Ok(Self(operation)) + } + + #[staticmethod] + #[pyo3(signature = (filter, payload, key=None))] + pub fn set_payload_by_filter( + filter: PyFilter, + payload: PyPayload, + key: Option, + ) -> Result { + let filter = Filter::from(filter); + let payload = Payload::from(payload); + + let key = key + .map(|k| JsonPath::from_str(&k).map_err(|_| PyErr::new::(k))) + .transpose()?; + + let operation = CollectionUpdateOperations::PayloadOperation( + payload_ops::PayloadOps::SetPayload(payload_ops::SetPayloadOp { + payload, + points: None, + filter: Some(filter), + key, + }), + ); + Ok(Self(operation)) + } + + #[staticmethod] + #[pyo3(signature = (ids, keys))] + pub fn delete_payload(ids: Vec, keys: Vec) -> Result { + let point_ids = PyPointId::into_rust_vec(ids); + + let keys: Vec<_> = keys + .into_iter() + .map(|k| JsonPath::from_str(&k).map_err(|_| PyErr::new::(k))) + .collect::>()?; + + let operation = CollectionUpdateOperations::PayloadOperation( + payload_ops::PayloadOps::DeletePayload(payload_ops::DeletePayloadOp { + keys, + points: Some(point_ids), + filter: None, + }), + ); + Ok(Self(operation)) + } + + #[staticmethod] + #[pyo3(signature = (filter, keys))] + pub fn delete_payload_by_filter(filter: PyFilter, keys: Vec) -> Result { + let filter = Filter::from(filter); + let keys: Vec<_> = keys + .into_iter() + .map(|k| JsonPath::from_str(&k).map_err(|_| PyErr::new::(k))) + .collect::>()?; + + let operation = CollectionUpdateOperations::PayloadOperation( + payload_ops::PayloadOps::DeletePayload(payload_ops::DeletePayloadOp { + keys, + points: None, + filter: Some(filter), + }), + ); + Ok(Self(operation)) + } + + #[staticmethod] + pub fn clear_payload(ids: Vec) -> Self { + let point_ids = PyPointId::into_rust_vec(ids); + let operation = + CollectionUpdateOperations::PayloadOperation(payload_ops::PayloadOps::ClearPayload { + points: point_ids, + }); + Self(operation) + } + + #[staticmethod] + pub fn clear_payload_by_filter(filter: PyFilter) -> Self { + let filter = Filter::from(filter); + let operation = CollectionUpdateOperations::PayloadOperation( + payload_ops::PayloadOps::ClearPayloadByFilter(filter), + ); + Self(operation) + } + + #[staticmethod] + #[pyo3(signature = (ids, payload, key=None))] + pub fn overwrite_payload( + ids: Vec, + payload: PyPayload, + key: Option, + ) -> Result { + let point_ids = PyPointId::into_rust_vec(ids); + let payload = Payload::from(payload); + + let key = key + .map(|k| JsonPath::from_str(&k).map_err(|_| PyErr::new::(k))) + .transpose()?; + + let operation = CollectionUpdateOperations::PayloadOperation( + payload_ops::PayloadOps::OverwritePayload(payload_ops::SetPayloadOp { + payload, + points: Some(point_ids), + filter: None, + key, + }), + ); + + Ok(Self(operation)) + } + + #[staticmethod] + #[pyo3(signature = (filter, payload, key=None))] + pub fn overwrite_payload_by_filter( + filter: PyFilter, + payload: PyPayload, + key: Option, + ) -> Result { + let filter = Filter::from(filter); + let payload = Payload::from(payload); + + let key = key + .map(|k| JsonPath::from_str(&k).map_err(|_| PyErr::new::(k))) + .transpose()?; + + let operation = CollectionUpdateOperations::PayloadOperation( + payload_ops::PayloadOps::OverwritePayload(payload_ops::SetPayloadOp { + payload, + points: None, + filter: Some(filter), + key, + }), + ); + Ok(Self(operation)) + } +} diff --git a/lib/edge/python/src/update/upsert.rs b/lib/edge/python/src/update/upsert.rs deleted file mode 100644 index d598a373e8..0000000000 --- a/lib/edge/python/src/update/upsert.rs +++ /dev/null @@ -1,42 +0,0 @@ -use pyo3::prelude::*; -use segment::types::Filter; -use shard::operations::point_ops::{PointInsertOperationsInternal, PointStructPersisted}; -use shard::operations::{CollectionUpdateOperations, point_ops}; - -use crate::types::filter::PyFilter; -use crate::types::point::PyPoint; -use crate::update::PyUpdateOperation; - -#[pymethods] -impl PyUpdateOperation { - #[staticmethod] - pub fn upsert_points(points: Vec) -> Self { - let points = points.into_iter().map(PointStructPersisted::from).collect(); - - let operation = - CollectionUpdateOperations::PointOperation(point_ops::PointOperations::UpsertPoints( - PointInsertOperationsInternal::PointsList(points), - )); - - Self(operation) - } - - #[staticmethod] - pub fn update_conditional(points: Vec, condition: PyFilter) -> Self { - let points = points.into_iter().map(PointStructPersisted::from).collect(); - let points_op = PointInsertOperationsInternal::PointsList(points); - - let condition = Filter::from(condition); - - let operation = CollectionUpdateOperations::PointOperation( - point_ops::PointOperations::UpsertPointsConditional( - point_ops::ConditionalInsertOperationInternal { - points_op, - condition, - }, - ), - ); - - Self(operation) - } -}