diff --git a/Cargo.lock b/Cargo.lock index f7309b1245..099f9732bc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1829,7 +1829,9 @@ dependencies = [ "edge", "pyo3", "segment", + "serde_json", "shard", + "sparse", "uuid", ] diff --git a/lib/edge/python/Cargo.toml b/lib/edge/python/Cargo.toml index acc35face4..905958fb8f 100644 --- a/lib/edge/python/Cargo.toml +++ b/lib/edge/python/Cargo.toml @@ -15,7 +15,9 @@ crate-type = ["cdylib"] edge = { path = ".." } segment = { path = "../../segment", default-features = false } shard = { path = "../../shard" } +sparse = { path = "../../sparse" } +serde_json = { workspace = true } uuid = { workspace = true } [dependencies.derive_more] diff --git a/lib/edge/python/examples/qdrant-edge.py b/lib/edge/python/examples/qdrant-edge.py index 3b70b1ae2e..25da68af58 100644 --- a/lib/edge/python/examples/qdrant-edge.py +++ b/lib/edge/python/examples/qdrant-edge.py @@ -19,7 +19,29 @@ config = SegmentConfig( shard = Shard(".", config) shard.update(UpdateOperation.upsert_points([ - Point(PointId.num(1), Vector.single([6.0, 9.0, 4.0, 2.0]), None), + Point( + PointId.num(1), + Vector.single([6.0, 9.0, 4.0, 2.0]), + Payload({ + "null": None, + "str": "string", + "uint": 42, + "int": -69, + "float": 4.20, + "bool": True, + "obj": { + "null": None, + "str": "string", + "uint": 42, + "int": -69, + "float": 4.20, + "bool": True, + "obj": {}, + "arr": [], + }, + "arr": [None, "string", 42, -69, 4.20, True, {}, []], + }), + ), ])) points = shard.search(SearchRequest( @@ -34,3 +56,4 @@ points = shard.search(SearchRequest( )) print(points[0].vector) +print(points[0].payload) diff --git a/lib/edge/python/src/lib.rs b/lib/edge/python/src/lib.rs index 79eadbe386..0fbe182d9a 100644 --- a/lib/edge/python/src/lib.rs +++ b/lib/edge/python/src/lib.rs @@ -82,11 +82,82 @@ pub struct PyPayload(Payload); #[pymethods] impl PyPayload { #[new] - fn new(_dict: Py) -> Self { - Self(Payload::default()) // TODO! + fn new(dict: Bound<'_, pyo3::types::PyDict>) -> pyo3::PyResult { + let obj = payload_object_from_py(&dict)?; + Ok(Self(Payload(obj))) } } +fn payload_object_from_py( + dict: &Bound<'_, pyo3::types::PyDict>, +) -> pyo3::PyResult> { + 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> { + 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 { + 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::(format!( + "failed to convert Python object {val} into payload value" + ))) +} + pub type PyResult = std::result::Result; #[derive(Debug)] diff --git a/lib/edge/python/src/search.rs b/lib/edge/python/src/search.rs index 4b0cb4c380..af174af7d1 100644 --- a/lib/edge/python/src/search.rs +++ b/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) -> 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, + exact: bool, + quantization: Option, + 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, oversampling: Option) -> 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) -> Self { - Self(WithVector::Selector(vectors)) // TODO? + pub fn selector(vectors: Vec) -> 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) -> 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) -> Self { + let include = PayloadSelectorInclude { + include: fields.into_iter().map(Into::into).collect(), // TODO: Transmute!? + }; + + Self(PayloadSelector::Include(include)) + } + + #[staticmethod] + pub fn exclude(fields: Vec) -> 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 { + 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 { + pub fn vector(&self) -> Option { let Some(vector) = &self.0.vector else { return None; }; @@ -187,7 +283,7 @@ impl PyScoredPoint { } #[getter] - fn payload(&self) -> Option { + pub fn payload(&self) -> Option { self.0.payload.clone().map(PyPayload) } } diff --git a/lib/edge/python/src/update.rs b/lib/edge/python/src/update.rs index d7a6beb33b..1720e108a1 100644 --- a/lib/edge/python/src/update.rs +++ b/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) -> Self { + Self(VectorStructPersisted::MultiDense(vectors)) + } + + #[staticmethod] + pub fn named(vectors: HashMap) -> 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) -> 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, values: Vec) -> Self { + Self(SparseVector { indices, values }) + } + + #[getter] + pub fn indices(&self) -> Vec { + self.0.indices.clone() + } + + #[getter] + pub fn values(&self) -> Vec { + self.0.values.clone() + } +}