Skip to content

Commit 75434ca

Browse files
authored
implement all point update operation constructors (#7424)
1 parent e01957a commit 75434ca

5 files changed

Lines changed: 311 additions & 45 deletions

File tree

lib/edge/python/src/lib.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ mod qdrant_edge {
4848
#[pymodule_export]
4949
use super::types::filter::{PyFilter, field_condition::PyFieldCondition};
5050
#[pymodule_export]
51-
use super::types::{PyPoint, PyRecord, PySparseVector};
51+
use super::types::{PyPoint, PyPointVectors, PyRecord, PySparseVector};
5252
#[pymodule_export]
5353
use super::update::PyUpdateOperation;
5454
}

lib/edge/python/src/types/mod.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,14 @@ pub mod filter;
22
pub mod payload;
33
pub mod point;
44
pub mod point_id;
5+
pub mod point_vectors;
56
pub mod record;
67
pub mod vector;
78

89
pub use filter::*;
910
pub use payload::*;
1011
pub use point::*;
1112
pub use point_id::*;
13+
pub use point_vectors::*;
1214
pub use record::*;
1315
pub use vector::*;
Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
use derive_more::Into;
2+
use pyo3::{pyclass, pymethods};
3+
use segment::types::PointIdType;
4+
use shard::operations::point_ops::VectorStructPersisted;
5+
use shard::operations::vector_ops::PointVectorsPersisted;
6+
7+
use crate::types::{PyPointId, PyVector};
8+
9+
#[pyclass(name = "PointVectors")]
10+
#[derive(Clone, Debug, Into)]
11+
#[repr(transparent)]
12+
pub struct PyPointVectors(pub PointVectorsPersisted);
13+
14+
#[pymethods]
15+
impl PyPointVectors {
16+
#[new]
17+
pub fn new(id: PyPointId, vector: PyVector) -> Self {
18+
Self(PointVectorsPersisted {
19+
id: PointIdType::from(id),
20+
vector: VectorStructPersisted::from(vector),
21+
})
22+
}
23+
}

lib/edge/python/src/update.rs

Lines changed: 285 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,292 @@
1-
pub mod upsert;
1+
use std::str::FromStr;
22

33
use derive_more::Into;
4+
use pyo3::exceptions::PyValueError;
45
use pyo3::prelude::*;
5-
use shard::operations::CollectionUpdateOperations;
6+
use segment::json_path::JsonPath;
7+
use segment::types::{Filter, Payload, VectorNameBuf};
8+
use shard::operations::point_ops::{
9+
PointIdsList, PointInsertOperationsInternal, PointStructPersisted,
10+
};
11+
use shard::operations::vector_ops::PointVectorsPersisted;
12+
use shard::operations::{CollectionUpdateOperations, payload_ops, point_ops, vector_ops};
13+
14+
use crate::types::{PyFilter, PyPayload, PyPoint, PyPointId, PyPointVectors};
615

716
#[pyclass(name = "UpdateOperation")]
817
#[derive(Clone, Debug, Into)]
918
pub struct PyUpdateOperation(CollectionUpdateOperations);
19+
20+
#[pymethods] // Can't split impl block due to pyo3 limitations, so all constructors go here
21+
impl PyUpdateOperation {
22+
#[staticmethod]
23+
pub fn upsert_points(points: Vec<PyPoint>) -> Self {
24+
let points = points.into_iter().map(PointStructPersisted::from).collect();
25+
26+
let operation =
27+
CollectionUpdateOperations::PointOperation(point_ops::PointOperations::UpsertPoints(
28+
PointInsertOperationsInternal::PointsList(points),
29+
));
30+
31+
Self(operation)
32+
}
33+
34+
#[staticmethod]
35+
pub fn upsert_points_conditional(points: Vec<PyPoint>, condition: PyFilter) -> Self {
36+
let points = points.into_iter().map(PointStructPersisted::from).collect();
37+
let points_op = PointInsertOperationsInternal::PointsList(points);
38+
39+
let condition = Filter::from(condition);
40+
41+
let operation = CollectionUpdateOperations::PointOperation(
42+
point_ops::PointOperations::UpsertPointsConditional(
43+
point_ops::ConditionalInsertOperationInternal {
44+
points_op,
45+
condition,
46+
},
47+
),
48+
);
49+
50+
Self(operation)
51+
}
52+
53+
#[staticmethod]
54+
pub fn delete_points(ids: Vec<PyPointId>) -> Self {
55+
let point_ids = PyPointId::into_rust_vec(ids);
56+
57+
let operation =
58+
CollectionUpdateOperations::PointOperation(point_ops::PointOperations::DeletePoints {
59+
ids: point_ids,
60+
});
61+
62+
Self(operation)
63+
}
64+
65+
#[staticmethod]
66+
pub fn delete_points_by_filter(filter: PyFilter) -> Self {
67+
let filter = Filter::from(filter);
68+
69+
let operation = CollectionUpdateOperations::PointOperation(
70+
point_ops::PointOperations::DeletePointsByFilter(filter),
71+
);
72+
73+
Self(operation)
74+
}
75+
76+
#[staticmethod]
77+
pub fn update_vectors(point_vectors: Vec<PyPointVectors>) -> Self {
78+
let points = point_vectors
79+
.into_iter()
80+
.map(PointVectorsPersisted::from)
81+
.collect();
82+
83+
let operation = CollectionUpdateOperations::VectorOperation(
84+
vector_ops::VectorOperations::UpdateVectors(vector_ops::UpdateVectorsOp {
85+
points,
86+
update_filter: None,
87+
}),
88+
);
89+
90+
Self(operation)
91+
}
92+
93+
#[staticmethod]
94+
pub fn update_vectors_conditional(
95+
point_vectors: Vec<PyPointVectors>,
96+
filter: PyFilter,
97+
) -> Self {
98+
let points = point_vectors
99+
.into_iter()
100+
.map(PointVectorsPersisted::from)
101+
.collect();
102+
let filter = Filter::from(filter);
103+
let operation = CollectionUpdateOperations::VectorOperation(
104+
vector_ops::VectorOperations::UpdateVectors(vector_ops::UpdateVectorsOp {
105+
points,
106+
update_filter: Some(filter),
107+
}),
108+
);
109+
Self(operation)
110+
}
111+
112+
#[staticmethod]
113+
pub fn delete_vectors(ids: Vec<PyPointId>, vector_names: Vec<VectorNameBuf>) -> Self {
114+
let point_ids = PyPointId::into_rust_vec(ids);
115+
let operation = CollectionUpdateOperations::VectorOperation(
116+
vector_ops::VectorOperations::DeleteVectors(
117+
PointIdsList::from(point_ids),
118+
vector_names,
119+
),
120+
);
121+
Self(operation)
122+
}
123+
124+
#[staticmethod]
125+
pub fn delete_vectors_by_filter(filter: PyFilter, vector_names: Vec<VectorNameBuf>) -> Self {
126+
let filter = Filter::from(filter);
127+
let operation = CollectionUpdateOperations::VectorOperation(
128+
vector_ops::VectorOperations::DeleteVectorsByFilter(filter, vector_names),
129+
);
130+
Self(operation)
131+
}
132+
133+
#[staticmethod]
134+
#[pyo3(signature = (ids, payload, key=None))]
135+
pub fn set_payload(
136+
ids: Vec<PyPointId>,
137+
payload: PyPayload,
138+
key: Option<String>,
139+
) -> Result<Self, PyErr> {
140+
let point_ids = PyPointId::into_rust_vec(ids);
141+
let payload = Payload::from(payload);
142+
143+
let key = key
144+
.map(|k| JsonPath::from_str(&k).map_err(|_| PyErr::new::<PyValueError, _>(k)))
145+
.transpose()?;
146+
147+
let operation = CollectionUpdateOperations::PayloadOperation(
148+
payload_ops::PayloadOps::SetPayload(payload_ops::SetPayloadOp {
149+
payload,
150+
points: Some(point_ids),
151+
filter: None,
152+
key,
153+
}),
154+
);
155+
156+
Ok(Self(operation))
157+
}
158+
159+
#[staticmethod]
160+
#[pyo3(signature = (filter, payload, key=None))]
161+
pub fn set_payload_by_filter(
162+
filter: PyFilter,
163+
payload: PyPayload,
164+
key: Option<String>,
165+
) -> Result<Self, PyErr> {
166+
let filter = Filter::from(filter);
167+
let payload = Payload::from(payload);
168+
169+
let key = key
170+
.map(|k| JsonPath::from_str(&k).map_err(|_| PyErr::new::<PyValueError, _>(k)))
171+
.transpose()?;
172+
173+
let operation = CollectionUpdateOperations::PayloadOperation(
174+
payload_ops::PayloadOps::SetPayload(payload_ops::SetPayloadOp {
175+
payload,
176+
points: None,
177+
filter: Some(filter),
178+
key,
179+
}),
180+
);
181+
Ok(Self(operation))
182+
}
183+
184+
#[staticmethod]
185+
#[pyo3(signature = (ids, keys))]
186+
pub fn delete_payload(ids: Vec<PyPointId>, keys: Vec<String>) -> Result<Self, PyErr> {
187+
let point_ids = PyPointId::into_rust_vec(ids);
188+
189+
let keys: Vec<_> = keys
190+
.into_iter()
191+
.map(|k| JsonPath::from_str(&k).map_err(|_| PyErr::new::<PyValueError, _>(k)))
192+
.collect::<Result<_, _>>()?;
193+
194+
let operation = CollectionUpdateOperations::PayloadOperation(
195+
payload_ops::PayloadOps::DeletePayload(payload_ops::DeletePayloadOp {
196+
keys,
197+
points: Some(point_ids),
198+
filter: None,
199+
}),
200+
);
201+
Ok(Self(operation))
202+
}
203+
204+
#[staticmethod]
205+
#[pyo3(signature = (filter, keys))]
206+
pub fn delete_payload_by_filter(filter: PyFilter, keys: Vec<String>) -> Result<Self, PyErr> {
207+
let filter = Filter::from(filter);
208+
let keys: Vec<_> = keys
209+
.into_iter()
210+
.map(|k| JsonPath::from_str(&k).map_err(|_| PyErr::new::<PyValueError, _>(k)))
211+
.collect::<Result<_, _>>()?;
212+
213+
let operation = CollectionUpdateOperations::PayloadOperation(
214+
payload_ops::PayloadOps::DeletePayload(payload_ops::DeletePayloadOp {
215+
keys,
216+
points: None,
217+
filter: Some(filter),
218+
}),
219+
);
220+
Ok(Self(operation))
221+
}
222+
223+
#[staticmethod]
224+
pub fn clear_payload(ids: Vec<PyPointId>) -> Self {
225+
let point_ids = PyPointId::into_rust_vec(ids);
226+
let operation =
227+
CollectionUpdateOperations::PayloadOperation(payload_ops::PayloadOps::ClearPayload {
228+
points: point_ids,
229+
});
230+
Self(operation)
231+
}
232+
233+
#[staticmethod]
234+
pub fn clear_payload_by_filter(filter: PyFilter) -> Self {
235+
let filter = Filter::from(filter);
236+
let operation = CollectionUpdateOperations::PayloadOperation(
237+
payload_ops::PayloadOps::ClearPayloadByFilter(filter),
238+
);
239+
Self(operation)
240+
}
241+
242+
#[staticmethod]
243+
#[pyo3(signature = (ids, payload, key=None))]
244+
pub fn overwrite_payload(
245+
ids: Vec<PyPointId>,
246+
payload: PyPayload,
247+
key: Option<String>,
248+
) -> Result<Self, PyErr> {
249+
let point_ids = PyPointId::into_rust_vec(ids);
250+
let payload = Payload::from(payload);
251+
252+
let key = key
253+
.map(|k| JsonPath::from_str(&k).map_err(|_| PyErr::new::<PyValueError, _>(k)))
254+
.transpose()?;
255+
256+
let operation = CollectionUpdateOperations::PayloadOperation(
257+
payload_ops::PayloadOps::OverwritePayload(payload_ops::SetPayloadOp {
258+
payload,
259+
points: Some(point_ids),
260+
filter: None,
261+
key,
262+
}),
263+
);
264+
265+
Ok(Self(operation))
266+
}
267+
268+
#[staticmethod]
269+
#[pyo3(signature = (filter, payload, key=None))]
270+
pub fn overwrite_payload_by_filter(
271+
filter: PyFilter,
272+
payload: PyPayload,
273+
key: Option<String>,
274+
) -> Result<Self, PyErr> {
275+
let filter = Filter::from(filter);
276+
let payload = Payload::from(payload);
277+
278+
let key = key
279+
.map(|k| JsonPath::from_str(&k).map_err(|_| PyErr::new::<PyValueError, _>(k)))
280+
.transpose()?;
281+
282+
let operation = CollectionUpdateOperations::PayloadOperation(
283+
payload_ops::PayloadOps::OverwritePayload(payload_ops::SetPayloadOp {
284+
payload,
285+
points: None,
286+
filter: Some(filter),
287+
key,
288+
}),
289+
);
290+
Ok(Self(operation))
291+
}
292+
}

lib/edge/python/src/update/upsert.rs

Lines changed: 0 additions & 42 deletions
This file was deleted.

0 commit comments

Comments
 (0)