WIP: Extend Qdrant Edge Python bindings (#7343)

lib/edge/python/src/lib.rs

@@ -82,11 +82,82 @@ pub struct PyPayload(Payload);
 #[pymethods]
 impl PyPayload {
     #[new]
-    fn new(_dict: Py<pyo3::types::PyDict>) -> Self {
-        Self(Payload::default()) // TODO!
+    fn new(dict: Bound<'_, pyo3::types::PyDict>) -> pyo3::PyResult<Self> {
+        let obj = payload_object_from_py(&dict)?;
+        Ok(Self(Payload(obj)))
     }
 }
 
+fn payload_object_from_py(
+    dict: &Bound<'_, pyo3::types::PyDict>,
+) -> pyo3::PyResult<serde_json::Map<String, serde_json::Value>> {
+    let mut object = serde_json::Map::with_capacity(dict.len());
+
+    for (key, value) in dict {
+        let key = key.extract()?;
+        let value = payload_value_from_py(&value)?;
+        object.insert(key, value);
+    }
+
+    Ok(object)
+}
+
+fn payload_array_from_py(
+    list: &Bound<'_, pyo3::types::PyList>,
+) -> pyo3::PyResult<Vec<serde_json::Value>> {
+    let mut array = Vec::with_capacity(list.len());
+
+    for value in list {
+        let value = payload_value_from_py(&value)?;
+        array.push(value);
+    }
+
+    Ok(array)
+}
+
+fn payload_value_from_py(val: &Bound<'_, PyAny>) -> pyo3::PyResult<serde_json::Value> {
+    if val.is_none() {
+        return Ok(serde_json::Value::Null);
+    }
+
+    if let Ok(dict) = val.cast() {
+        let obj = payload_object_from_py(dict)?;
+        return Ok(serde_json::Value::Object(obj));
+    }
+
+    if let Ok(list) = val.cast() {
+        let arr = payload_array_from_py(list)?;
+        return Ok(serde_json::Value::Array(arr));
+    }
+
+    if let Ok(str) = val.extract() {
+        return Ok(serde_json::Value::String(str));
+    }
+
+    if let Ok(uint) = val.extract() {
+        let num = serde_json::Number::from_u128(uint).unwrap(); // TODO?
+        return Ok(serde_json::Value::Number(num));
+    }
+
+    if let Ok(int) = val.extract() {
+        let num = serde_json::Number::from_i128(int).unwrap(); // TODO?
+        return Ok(serde_json::Value::Number(num));
+    }
+
+    if let Ok(float) = val.extract() {
+        let num = serde_json::Number::from_f64(float).unwrap(); // TODO?
+        return Ok(serde_json::Value::Number(num));
+    }
+
+    if let Ok(bool) = val.extract() {
+        return Ok(serde_json::Value::Bool(bool));
+    }
+
+    Err(PyErr::new::<PyException, _>(format!(
+        "failed to convert Python object {val} into payload value"
+    )))
+}
+
 pub type PyResult<T, E = PyError> = std::result::Result<T, E>;
 
 #[derive(Debug)]

lib/edge/python/src/search.rs

@@ -1,6 +1,7 @@
 use derive_more::Into;
 use pyo3::prelude::*;
 use segment::data_types::vectors::*;
+use segment::json_path::JsonPath;
 use segment::types::*;
 use shard::query::query_enum::QueryEnum;
 use shard::search::*;
@@ -75,13 +76,18 @@ impl PyQueryVector {
     pub fn dense(vector: Vec<f32>) -> Self {
         Self(VectorInternal::Dense(vector))
     }
+
+    #[staticmethod]
+    pub fn sparse(vector: PySparseVector) -> Self {
+        Self(VectorInternal::Sparse(vector.into()))
+    }
 }
 
 impl PyQueryVector {
     fn _variants(query: VectorInternal) {
         match query {
             VectorInternal::Dense(_) => (),
-            VectorInternal::Sparse(_) => todo!(),     // TODO!
+            VectorInternal::Sparse(_) => (),
             VectorInternal::MultiDense(_) => todo!(), // TODO!
         }
     }
@@ -97,12 +103,41 @@ impl PyFilter {
 }
 
 #[pyclass(name = "SearchParams")]
-#[derive(Clone, Debug, Into)]
+#[derive(Copy, Clone, Debug, Into)]
 pub struct PySearchParams(SearchParams);
 
 #[pymethods]
 impl PySearchParams {
-    // TODO!
+    #[new]
+    pub fn new(
+        hnsw_ef: Option<usize>,
+        exact: bool,
+        quantization: Option<PyQuantizationSearchParams>,
+        indexed_only: bool,
+    ) -> Self {
+        Self(SearchParams {
+            hnsw_ef,
+            exact,
+            quantization: quantization.map(Into::into),
+            indexed_only,
+        })
+    }
+}
+
+#[pyclass(name = "QuantizationSearchParams")]
+#[derive(Copy, Clone, Debug, Into)]
+pub struct PyQuantizationSearchParams(QuantizationSearchParams);
+
+#[pymethods]
+impl PyQuantizationSearchParams {
+    #[new]
+    pub fn new(ignore: bool, rescore: Option<bool>, oversampling: Option<f64>) -> Self {
+        Self(QuantizationSearchParams {
+            ignore,
+            rescore,
+            oversampling,
+        })
+    }
 }
 
 #[pyclass(name = "WithVector")]
@@ -112,13 +147,13 @@ pub struct PyWithVector(WithVector);
 #[pymethods]
 impl PyWithVector {
     #[new]
-    fn new(with_vector: bool) -> Self {
+    pub fn new(with_vector: bool) -> Self {
         Self(WithVector::Bool(with_vector))
     }
 
     #[staticmethod]
-    fn selector(vectors: Vec<VectorNameBuf>) -> Self {
-        Self(WithVector::Selector(vectors)) // TODO?
+    pub fn selector(vectors: Vec<VectorNameBuf>) -> Self {
+        Self(WithVector::Selector(vectors))
     }
 }
 
@@ -126,7 +161,7 @@ impl PyWithVector {
     fn _variants(with_vector: WithVector) {
         match with_vector {
             WithVector::Bool(_) => (),
-            WithVector::Selector(_) => (), // TODO?
+            WithVector::Selector(_) => (),
         }
     }
 }
@@ -141,18 +176,79 @@ impl PyWithPayload {
     pub fn new(with_payload: bool) -> Self {
         Self(WithPayloadInterface::Bool(with_payload))
     }
+
+    #[staticmethod]
+    pub fn fields(fields: Vec<PyJsonPath>) -> Self {
+        let fields = fields.into_iter().map(Into::into).collect(); // TODO: Transmute!?
+        Self(WithPayloadInterface::Fields(fields))
+    }
+
+    #[staticmethod]
+    pub fn selector(selector: PyPayloadSelector) -> Self {
+        Self(WithPayloadInterface::Selector(selector.into()))
+    }
 }
 
 impl PyWithPayload {
     fn _variants(with_payload: WithPayloadInterface) {
         match with_payload {
             WithPayloadInterface::Bool(_) => (),
-            WithPayloadInterface::Fields(_) => todo!(), // TODO!
-            WithPayloadInterface::Selector(_) => todo!(), // TODO!
+            WithPayloadInterface::Fields(_) => (),
+            WithPayloadInterface::Selector(_) => (),
         }
     }
 }
 
+#[pyclass(name = "PayloadSelector")]
+#[derive(Clone, Debug, Into)]
+pub struct PyPayloadSelector(PayloadSelector);
+
+#[pymethods]
+impl PyPayloadSelector {
+    #[staticmethod]
+    pub fn include(fields: Vec<PyJsonPath>) -> Self {
+        let include = PayloadSelectorInclude {
+            include: fields.into_iter().map(Into::into).collect(), // TODO: Transmute!?
+        };
+
+        Self(PayloadSelector::Include(include))
+    }
+
+    #[staticmethod]
+    pub fn exclude(fields: Vec<PyJsonPath>) -> Self {
+        let exclude = PayloadSelectorExclude {
+            exclude: fields.into_iter().map(Into::into).collect(), // TODO: Transmute!?
+        };
+
+        Self(PayloadSelector::Exclude(exclude))
+    }
+}
+
+impl PyPayloadSelector {
+    fn _variants(payload_selector: PayloadSelector) {
+        match payload_selector {
+            PayloadSelector::Include(_) => (),
+            PayloadSelector::Exclude(_) => (),
+        }
+    }
+}
+
+#[pyclass(name = "JsonPath")]
+#[derive(Clone, Debug, Into)]
+pub struct PyJsonPath(JsonPath);
+
+#[pymethods]
+impl PyJsonPath {
+    #[new]
+    pub fn new(json_path: &str) -> super::PyResult<Self> {
+        let json_path = json_path.parse().map_err(|()| {
+            OperationError::validation_error(format!("{json_path} is not a valid JSON path"))
+        })?;
+
+        Ok(Self(json_path))
+    }
+}
+
 #[pyclass(name = "ScoredPoint")]
 #[derive(Clone, Debug, Into)]
 pub struct PyScoredPoint(pub ScoredPoint);
@@ -160,22 +256,22 @@ pub struct PyScoredPoint(pub ScoredPoint);
 #[pymethods]
 impl PyScoredPoint {
     #[getter]
-    fn id(&self) -> PyPointId {
+    pub fn id(&self) -> PyPointId {
         PyPointId(self.0.id)
     }
 
     #[getter]
-    fn version(&self) -> u64 {
+    pub fn version(&self) -> u64 {
         self.0.version
     }
 
     #[getter]
-    fn score(&self) -> f32 {
+    pub fn score(&self) -> f32 {
         self.0.score
     }
 
     #[getter]
-    fn vector(&self) -> Option<DenseVector> {
+    pub fn vector(&self) -> Option<DenseVector> {
         let Some(vector) = &self.0.vector else {
             return None;
         };
@@ -187,7 +283,7 @@ impl PyScoredPoint {
     }
 
     #[getter]
-    fn payload(&self) -> Option<PyPayload> {
+    pub fn payload(&self) -> Option<PyPayload> {
         self.0.payload.clone().map(PyPayload)
     }
 }

lib/edge/python/src/update.rs

@@ -1,8 +1,11 @@
+use std::collections::HashMap;
+
 use derive_more::Into;
 use pyo3::prelude::*;
 use segment::data_types::vectors::*;
 use shard::operations::point_ops::*;
 use shard::operations::{CollectionUpdateOperations, point_ops};
+use sparse::common::sparse_vector::SparseVector;
 
 use super::*;
 
@@ -50,8 +53,24 @@ pub struct PyVector(VectorStructPersisted);
 #[pymethods]
 impl PyVector {
     #[staticmethod]
-    fn single(vec: DenseVector) -> Self {
-        Self(VectorStructPersisted::Single(vec))
+    pub fn single(vector: DenseVector) -> Self {
+        Self(VectorStructPersisted::Single(vector))
+    }
+
+    #[staticmethod]
+    pub fn multi_dense(vectors: Vec<DenseVector>) -> Self {
+        Self(VectorStructPersisted::MultiDense(vectors))
+    }
+
+    #[staticmethod]
+    pub fn named(vectors: HashMap<VectorNameBuf, PyNamedVector>) -> Self {
+        // TODO: Transmute!?
+        let vectors = vectors
+            .into_iter()
+            .map(|(name, vector)| (name, vector.into()))
+            .collect();
+
+        Self(VectorStructPersisted::Named(vectors))
     }
 }
 
@@ -59,8 +78,62 @@ impl PyVector {
     fn _variants(vector: VectorStructPersisted) {
         match vector {
             VectorStructPersisted::Single(_) => (),
-            VectorStructPersisted::MultiDense(_) => todo!(), // TODO!
-            VectorStructPersisted::Named(_) => todo!(),      // TODO!
+            VectorStructPersisted::MultiDense(_) => (),
+            VectorStructPersisted::Named(_) => (),
         }
     }
 }
+
+#[pyclass(name = "NamedVector")]
+#[derive(Clone, Debug, Into)]
+pub struct PyNamedVector(VectorPersisted);
+
+#[pymethods]
+impl PyNamedVector {
+    #[staticmethod]
+    pub fn dense(vector: DenseVector) -> Self {
+        Self(VectorPersisted::Dense(vector))
+    }
+
+    #[staticmethod]
+    pub fn multi_dense(vectors: Vec<DenseVector>) -> Self {
+        Self(VectorPersisted::MultiDense(vectors))
+    }
+
+    #[staticmethod]
+    pub fn sparse(vector: PySparseVector) -> Self {
+        Self(VectorPersisted::Sparse(vector.into()))
+    }
+}
+
+impl PyNamedVector {
+    fn _variants(vector: VectorPersisted) {
+        match vector {
+            VectorPersisted::Dense(_) => (),
+            VectorPersisted::MultiDense(_) => (),
+            VectorPersisted::Sparse(_) => (),
+        }
+    }
+}
+
+#[pyclass(name = "SparseVector")]
+#[derive(Clone, Debug, Into)]
+pub struct PySparseVector(SparseVector);
+
+#[pymethods]
+impl PySparseVector {
+    #[new]
+    pub fn new(indices: Vec<u32>, values: Vec<f32>) -> Self {
+        Self(SparseVector { indices, values })
+    }
+
+    #[getter]
+    pub fn indices(&self) -> Vec<u32> {
+        self.0.indices.clone()
+    }
+
+    #[getter]
+    pub fn values(&self) -> Vec<f32> {
+        self.0.values.clone()
+    }
+}