diff --git a/lib/edge/python/examples/qdrant-edge.py b/lib/edge/python/examples/qdrant-edge.py index 7dd360fbfb..7fdc2c3162 100644 --- a/lib/edge/python/examples/qdrant-edge.py +++ b/lib/edge/python/examples/qdrant-edge.py @@ -1,13 +1,25 @@ import os import shutil import uuid + from qdrant_edge import * +print("---- Load shard ----") + +DATA_DIRECTORY = "./data" + +# Clear and recreate data directory +if os.path.exists(DATA_DIRECTORY): + shutil.rmtree(DATA_DIRECTORY) + +os.makedirs(DATA_DIRECTORY) + +# Load Qdrant Edge shard config = SegmentConfig( vector_data={ "": VectorDataConfig( size=4, - distance=Distance.COSINE, + distance=Distance.DOT, storage_type=VectorStorageType.CHUNKED_MMAP, index=Indexes.PLAIN, quantization_config=None, @@ -19,20 +31,14 @@ config = SegmentConfig( payload_storage_type=PayloadStorageType.IN_RAM_MMAP, ) -DATA_DIRECTORY = "./data" - -# Clear and recreate data directory - -if os.path.exists(DATA_DIRECTORY): - shutil.rmtree(DATA_DIRECTORY) -os.makedirs(DATA_DIRECTORY) - shard = Shard(DATA_DIRECTORY, config) +print("---- Upsert ----") + shard.update(UpdateOperation.upsert_points([ Point( 1, - Vector.single([6.0, 9.0, 4.0, 2.0]), + [6.0, 9.0, 4.0, 2.0], { "null": None, "str": "string", @@ -55,35 +61,39 @@ shard.update(UpdateOperation.upsert_points([ ), Point( "e9408f2b-b917-4af1-ab75-d97ac6b2c047", - Vector.single([6.0, 9.0, 4.0, 2.0]), + [6.0, 9.0, 3.0, -2.0], { "hello": "world" }, ), Point( uuid.uuid4(), - Vector.single([6.0, 9.0, 4.0, 2.0]), + [1.0, 6.0, 4.0, 2.0], { "hello": "world" }, ), ])) +print("---- Search ----") + points = shard.search(SearchRequest( - query=Query.nearest(QueryVector.dense([1.0, 1.0, 1.0, 1.0]), None), + query=[1.0, 1.0, 1.0, 1.0], filter=None, params=None, limit=10, offset=0, - with_vector=WithVector(True), - with_payload=WithPayload(True), + with_vector=True, + with_payload=True, score_threshold=None, )) for point in points: print(f"Point: {point.id}, vector: {point.vector}, payload: {point.payload}, score: {point.score}") -retrieve = shard.retrieve(ids=[1], with_vector=WithVector(True), with_payload=WithPayload(True)) +print("---- Retrieve ----") -for point in retrieve: +points = shard.retrieve(point_ids=[1], with_vector=True, with_payload=True) + +for point in points: print(f"Point: {point.id}, vector: {point.vector}, payload: {point.payload}") diff --git a/lib/edge/python/src/lib.rs b/lib/edge/python/src/lib.rs index 71c199016b..1a1a4dbd86 100644 --- a/lib/edge/python/src/lib.rs +++ b/lib/edge/python/src/lib.rs @@ -26,12 +26,9 @@ mod qdrant_edge { PyVectorStorageDatatype, PyVectorStorageType, }; #[pymodule_export] - use super::search::{ - PyFilter, PyQuery, PyQueryVector, PyScoredPoint, PySearchParams, PySearchRequest, - PyWithPayload, PyWithVector, - }; + use super::search::{PyScoredPoint, PySearchParams, PySearchRequest}; #[pymodule_export] - use super::types::{PyRecord, PyVector}; + use super::types::{PyRecord, PySparseVector}; #[pymodule_export] use super::update::{PyPoint, PyUpdateOperation}; } @@ -55,29 +52,23 @@ impl PyShard { pub fn search(&self, search: PySearchRequest) -> Result> { let points = self.0.search(search.into())?; - let points = points.into_iter().map(PyScoredPoint).collect(); + let points = PyScoredPoint::from_rust_vec(points); Ok(points) } pub fn retrieve( &self, - ids: Vec, + point_ids: Vec, with_payload: Option, with_vector: Option, - ) -> Result, PyErr> { - let ids_res: Result, _> = ids.into_iter().map(PointIdType::try_from).collect(); - let ids = ids_res?; - - let records = self - .0 - .retrieve( - &ids, - with_payload.map(WithPayloadInterface::from), - with_vector.map(WithVector::from), - ) - .map_err(PyError::from)?; - - let points = records.into_iter().map(PyRecord).collect(); + ) -> Result> { + let point_ids = PyPointId::into_rust_vec(point_ids); + let points = self.0.retrieve( + &point_ids, + with_payload.map(WithPayloadInterface::from), + with_vector.map(WithVector::from), + )?; + let points = PyRecord::from_rust_vec(points); Ok(points) } } @@ -95,6 +86,6 @@ impl From for PyError { impl From for PyErr { fn from(err: PyError) -> Self { - PyErr::new::(err.0.to_string()) + PyException::new_err(err.0.to_string()) } } diff --git a/lib/edge/python/src/search.rs b/lib/edge/python/src/search.rs index dc8a0a920f..cf44eb207e 100644 --- a/lib/edge/python/src/search.rs +++ b/lib/edge/python/src/search.rs @@ -1,7 +1,10 @@ +use std::mem; + use derive_more::Into; +use pyo3::IntoPyObjectExt as _; +use pyo3::exceptions::PyValueError; use pyo3::prelude::*; -use segment::data_types::vectors::*; -use segment::json_path::JsonPath; +use segment::data_types::vectors::NamedQuery; use segment::types::*; use shard::query::query_enum::QueryEnum; use shard::search::*; @@ -39,67 +42,30 @@ impl PySearchRequest { } } -#[pyclass(name = "Query")] #[derive(Clone, Debug, Into)] pub struct PyQuery(QueryEnum); -#[pymethods] -impl PyQuery { - #[staticmethod] - pub fn nearest(query: PyQueryVector, using: Option) -> Self { - Self(QueryEnum::Nearest(NamedQuery { - query: query.into(), - using, - })) +impl<'py> FromPyObject<'py> for PyQuery { + fn extract_bound(query: &Bound<'py, PyAny>) -> PyResult { + let query = if let Ok(single) = query.extract() { + QueryEnum::Nearest(NamedQuery::default_dense(single)) + } else { + return Err(PyValueError::new_err(format!( + "failed to convert Python object {query} into query" + ))); + }; + + Ok(Self(query)) } } -impl PyQuery { - fn _variants(query: QueryEnum) { - match query { - QueryEnum::Nearest(_) => (), - QueryEnum::RecommendBestScore(_) => todo!(), // TODO! - QueryEnum::RecommendSumScores(_) => todo!(), // TODO! - QueryEnum::Discover(_) => todo!(), // TODO! - QueryEnum::Context(_) => todo!(), // TODO! - } - } -} - -#[pyclass(name = "QueryVector")] -#[derive(Clone, Debug, Into)] -pub struct PyQueryVector(VectorInternal); - -#[pymethods] -impl PyQueryVector { - #[staticmethod] - 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(_) => (), - VectorInternal::MultiDense(_) => todo!(), // TODO! - } - } -} - -#[pyclass(name = "Filter")] #[derive(Clone, Debug, Into)] pub struct PyFilter(Filter); -#[pymethods] -impl PyFilter { - // TODO! +impl<'py> FromPyObject<'py> for PyFilter { + fn extract_bound(_filter: &Bound<'py, PyAny>) -> PyResult { + todo!() + } } #[pyclass(name = "SearchParams")] @@ -140,124 +106,110 @@ impl PyQuantizationSearchParams { } } -#[pyclass(name = "WithVector")] #[derive(Clone, Debug, Into)] pub struct PyWithVector(WithVector); -#[pymethods] -impl PyWithVector { - #[new] - pub fn new(with_vector: bool) -> Self { - Self(WithVector::Bool(with_vector)) - } +impl<'py> FromPyObject<'py> for PyWithVector { + fn extract_bound(with_vector: &Bound<'py, PyAny>) -> PyResult { + #[derive(FromPyObject)] + enum Helper { + Bool(bool), + Selector(Vec), + } - #[staticmethod] - pub fn selector(vectors: Vec) -> Self { - Self(WithVector::Selector(vectors)) + let with_vector = match with_vector.extract()? { + Helper::Bool(bool) => WithVector::Bool(bool), + Helper::Selector(vectors) => WithVector::Selector(vectors), + }; + + Ok(Self(with_vector)) } } -impl PyWithVector { - fn _variants(with_vector: WithVector) { - match with_vector { - WithVector::Bool(_) => (), - WithVector::Selector(_) => (), +impl<'py> IntoPyObject<'py> for PyWithVector { + type Target = PyAny; + type Output = Bound<'py, Self::Target>; + type Error = PyErr; // Infallible? + + fn into_pyobject(self, py: Python<'py>) -> Result { + IntoPyObject::into_pyobject(&self, py) + } +} + +impl<'py> IntoPyObject<'py> for &PyWithVector { + type Target = PyAny; + type Output = Bound<'py, Self::Target>; + type Error = PyErr; // Infallible? + + fn into_pyobject(self, py: Python<'py>) -> Result { + match &self.0 { + WithVector::Bool(bool) => bool.into_bound_py_any(py), + WithVector::Selector(vectors) => vectors.into_bound_py_any(py), } } } -#[pyclass(name = "WithPayload")] #[derive(Clone, Debug, Into)] pub struct PyWithPayload(WithPayloadInterface); -#[pymethods] -impl PyWithPayload { - #[new] - 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(_) => (), - WithPayloadInterface::Selector(_) => (), +impl<'py> FromPyObject<'py> for PyWithPayload { + fn extract_bound(with_payload: &Bound<'py, PyAny>) -> PyResult { + #[derive(FromPyObject)] + enum Helper { + Bool(bool), + // TODO: `Fields(Vec)`! + // TODO: `Selector(PayloadSelector)`! } - } -} -#[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!? + let with_payload = match with_payload.extract()? { + Helper::Bool(bool) => WithPayloadInterface::Bool(bool), }; - 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)) + Ok(Self(with_payload)) } } -impl PyPayloadSelector { - fn _variants(payload_selector: PayloadSelector) { - match payload_selector { - PayloadSelector::Include(_) => (), - PayloadSelector::Exclude(_) => (), +impl<'py> IntoPyObject<'py> for PyWithPayload { + type Target = PyAny; + type Output = Bound<'py, Self::Target>; + type Error = PyErr; // Infallible? + + fn into_pyobject(self, py: Python<'py>) -> Result { + IntoPyObject::into_pyobject(&self, py) + } +} + +impl<'py> IntoPyObject<'py> for &PyWithPayload { + type Target = PyAny; + type Output = Bound<'py, Self::Target>; + type Error = PyErr; // Infallible? + + fn into_pyobject(self, py: Python<'py>) -> Result { + match &self.0 { + WithPayloadInterface::Bool(bool) => bool.into_bound_py_any(py), + WithPayloadInterface::Fields(_fields) => todo!(), + WithPayloadInterface::Selector(_selector) => todo!(), } } } -#[pyclass(name = "JsonPath")] -#[derive(Clone, Debug, Into)] -pub struct PyJsonPath(JsonPath); - -#[pymethods] -impl PyJsonPath { - #[new] - pub fn new(json_path: &str) -> super::Result { - 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)] +#[repr(transparent)] pub struct PyScoredPoint(pub ScoredPoint); +impl PyScoredPoint { + pub fn from_rust_vec(points: Vec) -> Vec { + // `PyScoredPoint` has transparent representation, so transmuting is safe + unsafe { mem::transmute(points) } + } +} + #[pymethods] impl PyScoredPoint { #[getter] pub fn id(&self) -> PyPointId { - PyPointId::from(self.0.id) + PyPointId(self.0.id) } #[getter] @@ -271,15 +223,8 @@ impl PyScoredPoint { } #[getter] - pub fn vector(&self) -> Option { - let Some(vector) = &self.0.vector else { - return None; - }; - - match vector { - VectorStructInternal::Single(vec) => Some(vec.clone()), - _ => None, // TODO! - } + pub fn vector(&self) -> Option<&PyVectorInternal> { + self.0.vector.as_ref().map(PyVectorInternal::from_ref) } #[getter] diff --git a/lib/edge/python/src/types/payload.rs b/lib/edge/python/src/types/payload.rs index 8303841efc..83ca366b32 100644 --- a/lib/edge/python/src/types/payload.rs +++ b/lib/edge/python/src/types/payload.rs @@ -2,7 +2,7 @@ use std::mem; use derive_more::Into; use pyo3::IntoPyObjectExt as _; -use pyo3::exceptions::PyException; +use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::types::{PyDict, PyFloat, PyInt, PyList}; use segment::types::*; @@ -62,8 +62,8 @@ fn value_from_py(val: &Bound) -> PyResult { let obj = object_from_py(dict)?; serde_json::Value::Object(obj) } else { - return Err(PyErr::new::(format!( - "failed to convert Python object {val} into payload value" + return Err(PyValueError::new_err(format!( + "failed to convert Python object {val} into payload value type" ))); }; @@ -88,11 +88,13 @@ fn number_from_py(num: &Bound) -> PyResult { serde_json::Number::from(int) } else if let Ok(float) = num.extract() { serde_json::Number::from_f64(float).ok_or_else(|| { - PyErr::new::(format!("failed to convert {float} into payload number")) + PyValueError::new_err(format!( + "failed to convert {float} into payload number type" + )) })? } else { - return Err(PyErr::new::(format!( - "failed to convert Python object {num} into payload number" + return Err(PyValueError::new_err(format!( + "failed to convert Python object {num} into payload number type" ))); }; diff --git a/lib/edge/python/src/types/point_id.rs b/lib/edge/python/src/types/point_id.rs index b74b169ff4..7b5b62775e 100644 --- a/lib/edge/python/src/types/point_id.rs +++ b/lib/edge/python/src/types/point_id.rs @@ -1,38 +1,67 @@ -use pyo3::exceptions::PyException; +use std::mem; + +use derive_more::Into; +use pyo3::IntoPyObjectExt as _; +use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use segment::types::PointIdType; use uuid::Uuid; -#[derive(Clone, Debug, IntoPyObject, FromPyObject)] -pub enum PyPointId { - NumId(u64), - UuidString(String), - Uuid(Uuid), +#[derive(Copy, Clone, Debug, Into)] +#[repr(transparent)] +pub struct PyPointId(pub PointIdType); + +impl PyPointId { + pub fn into_rust_vec(point_ids: Vec) -> Vec { + // `PyPointId` has transparent representation, so transmuting is safe + unsafe { mem::transmute(point_ids) } + } } -impl TryFrom for PointIdType { - type Error = PyErr; - fn try_from(value: PyPointId) -> Result { - match value { - PyPointId::NumId(id) => Ok(PointIdType::NumId(id)), - PyPointId::UuidString(uuid_string) => { - let uuid = Uuid::parse_str(&uuid_string).map_err(|_| { - PyErr::new::(format!( - "failed to parse string {uuid_string} into UUID for point ID" - )) +impl<'py> FromPyObject<'py> for PyPointId { + fn extract_bound(point_id: &Bound<'py, PyAny>) -> PyResult { + #[derive(FromPyObject)] + enum Helper { + NumId(u64), + Uuid(Uuid), + UuidStr(String), + } + + let point_id = match point_id.extract()? { + Helper::NumId(id) => PointIdType::NumId(id), + Helper::Uuid(uuid) => PointIdType::Uuid(uuid), + Helper::UuidStr(uuid_str) => { + let uuid = Uuid::parse_str(&uuid_str).map_err(|err| { + PyValueError::new_err(format!("failed to parse {uuid_str} as UUID: {err}")) })?; - Ok(PointIdType::Uuid(uuid)) + + PointIdType::Uuid(uuid) } - PyPointId::Uuid(uuid) => Ok(PointIdType::Uuid(uuid)), - } + }; + + Ok(Self(point_id)) } } -impl From for PyPointId { - fn from(value: PointIdType) -> Self { - match value { - PointIdType::NumId(id) => PyPointId::NumId(id), - PointIdType::Uuid(uuid) => PyPointId::Uuid(uuid), +impl<'py> IntoPyObject<'py> for PyPointId { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; // Infallible + + fn into_pyobject(self, py: Python<'py>) -> Result { + IntoPyObject::into_pyobject(&self, py) + } +} + +impl<'py> IntoPyObject<'py> for &PyPointId { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; // Infallible + + fn into_pyobject(self, py: Python<'py>) -> Result { + match &self.0 { + PointIdType::NumId(id) => id.into_bound_py_any(py), + PointIdType::Uuid(uuid) => uuid.into_bound_py_any(py), } } } diff --git a/lib/edge/python/src/types/record.rs b/lib/edge/python/src/types/record.rs index a859bb078f..8c0b589cc5 100644 --- a/lib/edge/python/src/types/record.rs +++ b/lib/edge/python/src/types/record.rs @@ -1,31 +1,34 @@ +use std::mem; + use derive_more::Into; use pyo3::prelude::*; use segment::data_types::order_by::OrderValue; -use shard::operations::point_ops::VectorStructPersisted; use shard::retrieve::record_internal::RecordInternal; use crate::*; #[pyclass(name = "Record")] #[derive(Clone, Debug, Into)] +#[repr(transparent)] pub struct PyRecord(pub RecordInternal); +impl PyRecord { + pub fn from_rust_vec(records: Vec) -> Vec { + // `PyRecord` has transparent representation, so transmuting is safe + unsafe { mem::transmute(records) } + } +} + #[pymethods] impl PyRecord { #[getter] pub fn id(&self) -> PyPointId { - PyPointId::from(self.0.id) + PyPointId(self.0.id) } #[getter] - pub fn vector(&self) -> Option { - let Some(vector) = &self.0.vector else { - return None; - }; - - let vector_persisted = VectorStructPersisted::from(vector.clone()); - - Some(PyVector(vector_persisted)) + pub fn vector(&self) -> Option<&PyVectorInternal> { + self.0.vector.as_ref().map(PyVectorInternal::from_ref) } #[getter] @@ -41,7 +44,6 @@ impl PyRecord { #[derive(IntoPyObject)] pub enum PyOrderValue { - // Put Int first so ints don't get parsed as floats (since f64 can extract from ints). Int(i64), Float(f64), } @@ -49,8 +51,8 @@ pub enum PyOrderValue { impl From for PyOrderValue { fn from(value: OrderValue) -> Self { match value { - OrderValue::Int(int) => PyOrderValue::Int(int), - OrderValue::Float(float) => PyOrderValue::Float(float), + OrderValue::Int(int) => Self::Int(int), + OrderValue::Float(float) => Self::Float(float), } } } diff --git a/lib/edge/python/src/types/vector.rs b/lib/edge/python/src/types/vector.rs index 69814cff8d..797c629aad 100644 --- a/lib/edge/python/src/types/vector.rs +++ b/lib/edge/python/src/types/vector.rs @@ -1,87 +1,85 @@ use std::collections::HashMap; +use std::mem; use derive_more::Into; -use pyo3::{pyclass, pymethods}; -use segment::data_types::vectors::DenseVector; +use pyo3::IntoPyObjectExt as _; +use pyo3::prelude::*; +use segment::data_types::vectors::{DenseVector, VectorStructInternal}; use segment::types::VectorNameBuf; use shard::operations::point_ops::{VectorPersisted, VectorStructPersisted}; use sparse::common::sparse_vector::SparseVector; use sparse::common::types::{DimId, DimWeight}; -#[pyclass(name = "Vector")] #[derive(Clone, Debug, Into)] -pub struct PyVector(pub VectorStructPersisted); +pub struct PyVector(VectorStructPersisted); -#[pymethods] -impl PyVector { - #[staticmethod] - 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)) - } -} - -impl PyVector { - fn _variants(vector: VectorStructPersisted) { - match vector { - VectorStructPersisted::Single(_) => (), - VectorStructPersisted::MultiDense(_) => (), - VectorStructPersisted::Named(_) => (), +impl<'py> FromPyObject<'py> for PyVector { + fn extract_bound(vector: &Bound<'py, PyAny>) -> PyResult { + #[derive(FromPyObject)] + enum Helper { + Single(DenseVector), + Multi(Vec), + Named(HashMap), } + + let vector = match vector.extract()? { + Helper::Single(single) => VectorStructPersisted::Single(single), + Helper::Multi(multi) => VectorStructPersisted::MultiDense(multi), + Helper::Named(named) => { + let named = PyNamedVector::into_rust_map(named); + VectorStructPersisted::Named(named) + } + }; + + Ok(Self(vector)) } } -#[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())) - } -} +#[repr(transparent)] +struct PyNamedVector(VectorPersisted); impl PyNamedVector { - fn _variants(vector: VectorPersisted) { - match vector { - VectorPersisted::Dense(_) => (), - VectorPersisted::MultiDense(_) => (), - VectorPersisted::Sparse(_) => (), + pub fn into_rust_map( + vectors: HashMap, + ) -> HashMap { + unsafe { mem::transmute(vectors) } + } +} + +impl<'py> FromPyObject<'py> for PyNamedVector { + fn extract_bound(vector: &Bound<'py, PyAny>) -> PyResult { + #[derive(FromPyObject)] + enum Helper { + Dense(DenseVector), + MultiDense(Vec), + Sparse(PySparseVector), } + + let vector = match vector.extract()? { + Helper::Dense(dense) => VectorPersisted::Dense(dense), + Helper::MultiDense(multi) => VectorPersisted::MultiDense(multi), + Helper::Sparse(sparse) => { + let sparse = PySparseVector::into_rust(sparse); + VectorPersisted::Sparse(sparse) + } + }; + + Ok(Self(vector)) } } #[pyclass(name = "SparseVector")] #[derive(Clone, Debug, Into)] +#[repr(transparent)] pub struct PySparseVector(SparseVector); +impl PySparseVector { + pub fn into_rust(self) -> SparseVector { + unsafe { mem::transmute(self) } + } +} + #[pymethods] impl PySparseVector { #[new] @@ -90,12 +88,46 @@ impl PySparseVector { } #[getter] - pub fn indices(&self) -> Vec { - self.0.indices.clone() + pub fn indices(&self) -> &[DimId] { + &self.0.indices } #[getter] - pub fn values(&self) -> Vec { - self.0.values.clone() + pub fn values(&self) -> &[DimWeight] { + &self.0.values + } +} + +#[derive(Clone, Debug, Into)] +#[repr(transparent)] +pub struct PyVectorInternal(pub VectorStructInternal); + +impl PyVectorInternal { + pub fn from_ref(vector: &VectorStructInternal) -> &Self { + unsafe { mem::transmute(vector) } + } +} + +impl<'py> IntoPyObject<'py> for PyVectorInternal { + type Target = PyAny; + type Output = Bound<'py, Self::Target>; + type Error = PyErr; // Infallible + + fn into_pyobject(self, py: Python<'py>) -> Result { + IntoPyObject::into_pyobject(&self, py) + } +} + +impl<'py> IntoPyObject<'py> for &PyVectorInternal { + type Target = PyAny; + type Output = Bound<'py, Self::Target>; + type Error = PyErr; // Infallible + + fn into_pyobject(self, py: Python<'py>) -> Result { + match &self.0 { + VectorStructInternal::Single(single) => single.into_bound_py_any(py), + VectorStructInternal::MultiDense(_multi) => todo!(), + VectorStructInternal::Named(_named) => todo!(), + } } } diff --git a/lib/edge/python/src/update.rs b/lib/edge/python/src/update.rs index c1081bcb8a..8878796860 100644 --- a/lib/edge/python/src/update.rs +++ b/lib/edge/python/src/update.rs @@ -33,7 +33,7 @@ impl PyPoint { #[new] pub fn new(id: PyPointId, vector: PyVector, payload: Option) -> Result { let point = PointStructPersisted { - id: PointIdType::try_from(id)?, + id: PointIdType::from(id), vector: VectorStructPersisted::from(vector), payload: payload.map(Payload::from), }; diff --git a/lib/edge/src/retrieve.rs b/lib/edge/src/retrieve.rs index c189b13545..6a9a0ecd29 100644 --- a/lib/edge/src/retrieve.rs +++ b/lib/edge/src/retrieve.rs @@ -1,3 +1,5 @@ +use std::sync::atomic::AtomicBool; + use common::counter::hardware_accumulator::HwMeasurementAcc; use segment::common::operation_error::OperationResult; use segment::types::{ExtendedPointId, WithPayload, WithPayloadInterface, WithVector}; @@ -9,7 +11,7 @@ use crate::Shard; impl Shard { pub fn retrieve( &self, - ids: &[ExtendedPointId], + point_ids: &[ExtendedPointId], with_payload: Option, with_vector: Option, ) -> OperationResult> { @@ -17,24 +19,20 @@ impl Shard { WithPayload::from(with_payload.unwrap_or(WithPayloadInterface::Bool(true))); let with_vector = with_vector.unwrap_or(WithVector::Bool(false)); - let never_stopped = std::sync::atomic::AtomicBool::new(false); - - let hw_measurement = HwMeasurementAcc::disposable(); - - let mut retrieve_result = retrieve_blocking( + let mut points = retrieve_blocking( self.segments.clone(), - ids, + point_ids, &with_payload, &with_vector, - &never_stopped, - hw_measurement, + &AtomicBool::new(false), + HwMeasurementAcc::disposable(), )?; - let response: Vec<_> = ids + let points: Vec<_> = point_ids .iter() - .filter_map(|idx| retrieve_result.remove(idx)) + .filter_map(|id| points.remove(id)) .collect(); - Ok(response) + Ok(points) } }