diff --git a/lib/collection/src/collection_manager/optimizers/config_mismatch_optimizer.rs b/lib/collection/src/collection_manager/optimizers/config_mismatch_optimizer.rs index d41f979f10..3dedd82e77 100644 --- a/lib/collection/src/collection_manager/optimizers/config_mismatch_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/config_mismatch_optimizer.rs @@ -336,7 +336,7 @@ mod tests { dir.path().to_owned(), temp_dir.path().to_owned(), collection_params.clone(), - hnsw_config.clone(), + hnsw_config, HnswGlobalConfig::default(), Default::default(), ); @@ -345,7 +345,7 @@ mod tests { dir.path().to_owned(), temp_dir.path().to_owned(), collection_params, - hnsw_config.clone(), + hnsw_config, HnswGlobalConfig::default(), Default::default(), ); @@ -380,7 +380,7 @@ mod tests { let mut changed_hnsw_config = hnsw_config; changed_hnsw_config.m /= 2; changed_hnsw_config.ef_construct /= 5; - config_mismatch_optimizer.hnsw_config = changed_hnsw_config.clone(); + config_mismatch_optimizer.hnsw_config = changed_hnsw_config; // Run mismatch optimizer again, make sure it optimizes now let permit = budget.try_acquire(0, permit_cpu_count).unwrap(); @@ -410,7 +410,7 @@ mod tests { .for_each(|segment| { assert_eq!( segment.config().vector_data[DEFAULT_VECTOR_NAME].index, - Indexes::Hnsw(changed_hnsw_config.clone()), + Indexes::Hnsw(changed_hnsw_config), "segment must be optimized with changed HNSW config", ); }); @@ -501,7 +501,7 @@ mod tests { dir.path().to_owned(), temp_dir.path().to_owned(), collection_params.clone(), - hnsw_config_collection.clone(), + hnsw_config_collection, HnswGlobalConfig::default(), Default::default(), ); @@ -510,7 +510,7 @@ mod tests { dir.path().to_owned(), temp_dir.path().to_owned(), collection_params, - hnsw_config_collection.clone(), + hnsw_config_collection, HnswGlobalConfig::default(), Default::default(), ); diff --git a/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs b/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs index 90c91111ac..26f704d8ec 100644 --- a/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/indexing_optimizer.rs @@ -889,7 +889,7 @@ mod tests { dir.path().to_owned(), temp_dir.path().to_owned(), collection_params.clone(), - hnsw_config.clone(), + hnsw_config, HnswGlobalConfig::default(), Default::default(), ); @@ -898,7 +898,7 @@ mod tests { dir.path().to_owned(), temp_dir.path().to_owned(), collection_params.clone(), - hnsw_config.clone(), + hnsw_config, HnswGlobalConfig::default(), Default::default(), ); @@ -955,7 +955,7 @@ mod tests { dir.path().to_owned(), temp_dir.path().to_owned(), collection_params.clone(), - hnsw_config.clone(), + hnsw_config, HnswGlobalConfig::default(), Default::default(), ); diff --git a/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs b/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs index 79f2728d1c..d47df47abc 100644 --- a/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs +++ b/lib/collection/src/collection_manager/optimizers/vacuum_optimizer.rs @@ -482,7 +482,7 @@ mod tests { dir.path().to_owned(), temp_dir.path().to_owned(), collection_params.clone(), - hnsw_config.clone(), + hnsw_config, HnswGlobalConfig::default(), Default::default(), ); diff --git a/lib/collection/src/optimizers_builder.rs b/lib/collection/src/optimizers_builder.rs index f8bc77089a..444aac6137 100644 --- a/lib/collection/src/optimizers_builder.rs +++ b/lib/collection/src/optimizers_builder.rs @@ -175,7 +175,7 @@ pub fn build_optimizers( segments_path.clone(), temp_segments_path.clone(), collection_params.clone(), - hnsw_config.clone(), + *hnsw_config, hnsw_global_config.clone(), quantization_config.clone(), )), @@ -185,7 +185,7 @@ pub fn build_optimizers( segments_path.clone(), temp_segments_path.clone(), collection_params.clone(), - hnsw_config.clone(), + *hnsw_config, hnsw_global_config.clone(), quantization_config.clone(), )), @@ -196,7 +196,7 @@ pub fn build_optimizers( segments_path.clone(), temp_segments_path.clone(), collection_params.clone(), - hnsw_config.clone(), + *hnsw_config, hnsw_global_config.clone(), quantization_config.clone(), )), @@ -205,7 +205,7 @@ pub fn build_optimizers( segments_path, temp_segments_path, collection_params.clone(), - hnsw_config.clone(), + *hnsw_config, hnsw_global_config.clone(), quantization_config.clone(), )), diff --git a/lib/edge/python/examples/common.py b/lib/edge/python/examples/common.py index 1433ae5026..0d7c0e1919 100644 --- a/lib/edge/python/examples/common.py +++ b/lib/edge/python/examples/common.py @@ -20,16 +20,16 @@ def load_new_shard(): vector_data={ "": VectorDataConfig( size=4, - distance=Distance.DOT, - storage_type=VectorStorageType.CHUNKED_MMAP, - index=Indexes.PLAIN, + distance=Distance.Dot, + storage_type=VectorStorageType.ChunkedMmap, + index=PlainIndexConfig(), quantization_config=None, multivector_config=None, datatype=None, ), }, sparse_vector_data={}, - payload_storage_type=PayloadStorageType.IN_RAM_MMAP, + payload_storage_type=PayloadStorageType.InRamMmap, ) return Shard(DATA_DIRECTORY, config) @@ -48,4 +48,3 @@ def fill_dummy_data(shard: Shard): Point(9, [0.45, 0.55, 0.65, 0.75], {"color": "green", "city": ["Berlin"], "score": 0.92}), Point(10, [0.01, 0.02, 0.03, 0.04], {"color": "yellow", "city": None, "featured": False}), ])) - diff --git a/lib/edge/python/src/config.rs b/lib/edge/python/src/config.rs deleted file mode 100644 index 05c9161d37..0000000000 --- a/lib/edge/python/src/config.rs +++ /dev/null @@ -1,244 +0,0 @@ -use std::collections::HashMap; - -use bytemuck::TransparentWrapper; -use derive_more::Into; -use pyo3::prelude::*; -use segment::types::*; - -#[pyclass(name = "SegmentConfig")] -#[derive(Clone, Debug, Into, TransparentWrapper)] -#[repr(transparent)] -pub struct PySegmentConfig(SegmentConfig); - -#[pymethods] -impl PySegmentConfig { - #[new] - pub fn new( - vector_data: HashMap, - sparse_vector_data: HashMap, - payload_storage_type: PyPayloadStorageType, - ) -> Self { - // TODO: Transmute!? - let vector_data = vector_data - .into_iter() - .map(|(vector, config)| (vector, VectorDataConfig::from(config))) - .collect(); - - // TODO: Transmute!? - let sparse_vector_data = sparse_vector_data - .into_iter() - .map(|(vector, config)| (vector, SparseVectorDataConfig::from(config))) - .collect(); - - Self(SegmentConfig { - vector_data, - sparse_vector_data, - payload_storage_type: PayloadStorageType::from(payload_storage_type), - }) - } -} - -#[pyclass(name = "VectorDataConfig")] -#[derive(Clone, Debug, Into, TransparentWrapper)] -#[repr(transparent)] -pub struct PyVectorDataConfig(VectorDataConfig); - -#[pymethods] -impl PyVectorDataConfig { - #[new] - pub fn new( - size: usize, - distance: PyDistance, - storage_type: PyVectorStorageType, - index: PyIndexes, - quantization_config: Option, - multivector_config: Option, - datatype: Option, - ) -> Self { - Self(VectorDataConfig { - size, - distance: Distance::from(distance), - storage_type: VectorStorageType::from(storage_type), - index: Indexes::from(index), - quantization_config: quantization_config.map(QuantizationConfig::from), - multivector_config: multivector_config.map(MultiVectorConfig::from), - datatype: datatype.map(VectorStorageDatatype::from), - }) - } -} - -#[pyclass(name = "Distance")] -#[derive(Copy, Clone, Debug, Into)] -pub struct PyDistance(Distance); - -#[pymethods] -impl PyDistance { - #[classattr] - pub const COSINE: Self = Self(Distance::Cosine); - - #[classattr] - pub const EUCLID: Self = Self(Distance::Euclid); - - #[classattr] - pub const DOT: Self = Self(Distance::Dot); - - #[classattr] - pub const MANHATTAN: Self = Self(Distance::Manhattan); -} - -impl PyDistance { - fn _variants(distance: Distance) { - match distance { - Distance::Cosine => (), - Distance::Euclid => (), - Distance::Dot => (), - Distance::Manhattan => (), - } - } -} - -#[pyclass(name = "VectorStorageType")] -#[derive(Copy, Clone, Debug, Into)] -pub struct PyVectorStorageType(VectorStorageType); - -#[pymethods] -impl PyVectorStorageType { - #[classattr] - pub const MEMORY: Self = Self(VectorStorageType::Memory); - - #[classattr] - pub const MMAP: Self = Self(VectorStorageType::Mmap); - - #[classattr] - pub const CHUNKED_MMAP: Self = Self(VectorStorageType::ChunkedMmap); - - #[classattr] - pub const IN_RAM_CHUNKED_MMAP: Self = Self(VectorStorageType::InRamChunkedMmap); -} - -impl PyVectorStorageType { - fn _variants(storage_type: VectorStorageType) { - match storage_type { - VectorStorageType::Memory => (), - VectorStorageType::Mmap => (), - VectorStorageType::ChunkedMmap => (), - VectorStorageType::InRamChunkedMmap => (), - } - } -} - -#[pyclass(name = "Indexes")] -#[derive(Clone, Debug, Into)] -pub struct PyIndexes(Indexes); - -#[pymethods] -impl PyIndexes { - #[classattr] - pub const PLAIN: Self = Self(Indexes::Plain {}); - - // TODO: HNSW!? -} - -impl PyIndexes { - fn _variants(indexes: Indexes) { - match indexes { - Indexes::Plain {} => (), - Indexes::Hnsw(_) => (), // TODO!? - } - } -} - -#[pyclass(name = "QuantizationConfig")] -#[derive(Clone, Debug, Into)] -pub struct PyQuantizationConfig(QuantizationConfig); // TODO!? - -#[pyclass(name = "MultiVectorConfig")] -#[derive(Copy, Clone, Debug, Into)] -pub struct PyMultiVectorConfig(MultiVectorConfig); - -#[pymethods] -impl PyMultiVectorConfig { - #[new] - pub fn new(comparator: PyMultiVectorComparator) -> Self { - Self(MultiVectorConfig { - comparator: MultiVectorComparator::from(comparator), - }) - } -} - -#[pyclass(name = "MultiVectorComparator")] -#[derive(Copy, Clone, Debug, Into)] -pub struct PyMultiVectorComparator(MultiVectorComparator); - -#[pymethods] -impl PyMultiVectorComparator { - #[classattr] - pub const MAX_SIM: Self = Self(MultiVectorComparator::MaxSim); -} - -impl PyMultiVectorComparator { - fn _variants(comparator: MultiVectorComparator) { - match comparator { - MultiVectorComparator::MaxSim => (), - } - } -} - -#[pyclass(name = "VectorStorageDatatype")] -#[derive(Copy, Clone, Debug, Into)] -pub struct PyVectorStorageDatatype(VectorStorageDatatype); - -#[pymethods] -impl PyVectorStorageDatatype { - #[classattr] - pub const FLOAT_32: Self = Self(VectorStorageDatatype::Float32); - - #[classattr] - pub const FLOAT_16: Self = Self(VectorStorageDatatype::Float16); - - #[classattr] - pub const UINT_8: Self = Self(VectorStorageDatatype::Uint8); -} - -impl PyVectorStorageDatatype { - fn _variants(storage_datatype: VectorStorageDatatype) { - match storage_datatype { - VectorStorageDatatype::Float32 => (), - VectorStorageDatatype::Float16 => (), - VectorStorageDatatype::Uint8 => (), - } - } -} - -#[pyclass(name = "SparseVectorDataConfig")] -#[derive(Copy, Clone, Debug, Into)] -pub struct PySparseVectorDataConfig(SparseVectorDataConfig); - -#[pymethods] -impl PySparseVectorDataConfig { - // TODO!? -} - -#[pyclass(name = "PayloadStorageType")] -#[derive(Copy, Clone, Debug, Into)] -pub struct PyPayloadStorageType(PayloadStorageType); - -#[pymethods] -impl PyPayloadStorageType { - #[classattr] - pub const MMAP: Self = Self(PayloadStorageType::Mmap); - - #[classattr] - pub const IN_RAM_MMAP: Self = Self(PayloadStorageType::InRamMmap); -} - -impl PyPayloadStorageType { - fn _variants(storage_type: PayloadStorageType) { - #[allow(unreachable_patterns)] - match storage_type { - PayloadStorageType::Mmap => (), - PayloadStorageType::InRamMmap => (), - _ => todo!(), // TODO: Ignore RocksDB storage types - } - } -} diff --git a/lib/edge/python/src/config/mod.rs b/lib/edge/python/src/config/mod.rs new file mode 100644 index 0000000000..23c7528569 --- /dev/null +++ b/lib/edge/python/src/config/mod.rs @@ -0,0 +1,77 @@ +pub mod quantization; +pub mod sparse_vector_data; +pub mod vector_data; + +use std::collections::HashMap; + +use bytemuck::TransparentWrapper; +use derive_more::Into; +use pyo3::prelude::*; +use segment::types::*; + +pub use self::quantization::*; +pub use self::sparse_vector_data::*; +pub use self::vector_data::*; + +#[pyclass(name = "SegmentConfig")] +#[derive(Clone, Debug, Into, TransparentWrapper)] +#[repr(transparent)] +pub struct PySegmentConfig(SegmentConfig); + +#[pymethods] +impl PySegmentConfig { + #[new] + pub fn new( + vector_data: HashMap, + sparse_vector_data: HashMap, + payload_storage_type: PyPayloadStorageType, + ) -> Self { + Self(SegmentConfig { + vector_data: PyVectorDataConfig::peel_map(vector_data), + sparse_vector_data: PySparseVectorDataConfig::peel_map(sparse_vector_data), + payload_storage_type: PayloadStorageType::from(payload_storage_type), + }) + } + + #[getter] + pub fn vector_data(&self) -> &HashMap { + PyVectorDataConfig::wrap_map_ref(&self.0.vector_data) + } + + #[getter] + pub fn sparse_vector_data(&self) -> &HashMap { + PySparseVectorDataConfig::wrap_map_ref(&self.0.sparse_vector_data) + } + + #[getter] + pub fn payload_storage_type(&self) -> PyPayloadStorageType { + PyPayloadStorageType::from(self.0.payload_storage_type) + } +} + +#[pyclass(name = "PayloadStorageType")] +#[derive(Copy, Clone, Debug)] +pub enum PyPayloadStorageType { + Mmap, + InRamMmap, +} + +impl From for PyPayloadStorageType { + fn from(storage_type: PayloadStorageType) -> Self { + #[allow(unreachable_patterns)] + match storage_type { + PayloadStorageType::Mmap => PyPayloadStorageType::Mmap, + PayloadStorageType::InRamMmap => PyPayloadStorageType::InRamMmap, + _ => unimplemented!("RocksDB-backed storage types are not supported by Qdrant Edge"), + } + } +} + +impl From for PayloadStorageType { + fn from(storage_type: PyPayloadStorageType) -> Self { + match storage_type { + PyPayloadStorageType::Mmap => PayloadStorageType::Mmap, + PyPayloadStorageType::InRamMmap => PayloadStorageType::InRamMmap, + } + } +} diff --git a/lib/edge/python/src/config/quantization.rs b/lib/edge/python/src/config/quantization.rs new file mode 100644 index 0000000000..8d7a99c42d --- /dev/null +++ b/lib/edge/python/src/config/quantization.rs @@ -0,0 +1,277 @@ +use derive_more::Into; +use pyo3::IntoPyObjectExt as _; +use pyo3::prelude::*; +use segment::types::*; + +#[derive(Clone, Debug, Into)] +pub struct PyQuantizationConfig(pub QuantizationConfig); + +impl FromPyObject<'_, '_> for PyQuantizationConfig { + type Error = PyErr; + + fn extract(conf: Borrowed<'_, '_, PyAny>) -> PyResult { + #[derive(FromPyObject)] + enum Helper { + Scalar(PyScalarQuantizationConfig), + Product(PyProductQuantizationConfig), + Binary(PyBinaryQuantizationConfig), + } + + let conf = match conf.extract()? { + Helper::Scalar(scalar) => QuantizationConfig::Scalar(ScalarQuantization { + scalar: ScalarQuantizationConfig::from(scalar), + }), + Helper::Product(product) => QuantizationConfig::Product(ProductQuantization { + product: ProductQuantizationConfig::from(product), + }), + Helper::Binary(binary) => QuantizationConfig::Binary(BinaryQuantization { + binary: BinaryQuantizationConfig::from(binary), + }), + }; + + Ok(Self(conf)) + } +} + +impl<'py> IntoPyObject<'py> for PyQuantizationConfig { + type Target = PyAny; + type Output = Bound<'py, Self::Target>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + match self.0 { + QuantizationConfig::Scalar(ScalarQuantization { scalar }) => { + PyScalarQuantizationConfig(scalar).into_bound_py_any(py) + } + QuantizationConfig::Product(ProductQuantization { product }) => { + PyProductQuantizationConfig(product).into_bound_py_any(py) + } + QuantizationConfig::Binary(BinaryQuantization { binary }) => { + PyBinaryQuantizationConfig(binary).into_bound_py_any(py) + } + } + } +} + +#[pyclass(name = "ScalarQuantizationConfig")] +#[derive(Clone, Debug, Into)] +pub struct PyScalarQuantizationConfig(ScalarQuantizationConfig); + +#[pymethods] +impl PyScalarQuantizationConfig { + #[new] + #[pyo3(signature = (r#type, quantile = None, always_ram = None))] + pub fn new(r#type: PyScalarType, quantile: Option, always_ram: Option) -> Self { + Self(ScalarQuantizationConfig { + r#type: ScalarType::from(r#type), + quantile, + always_ram, + }) + } + + #[getter] + pub fn r#type(&self) -> PyScalarType { + PyScalarType::from(self.0.r#type) + } + + #[getter] + pub fn quantile(&self) -> Option { + self.0.quantile + } + + #[getter] + pub fn always_ram(&self) -> Option { + self.0.always_ram + } +} + +#[pyclass(name = "ScalarType")] +#[derive(Copy, Clone, Debug)] +pub enum PyScalarType { + Int8, +} + +impl From for PyScalarType { + fn from(scalar_type: ScalarType) -> Self { + match scalar_type { + ScalarType::Int8 => PyScalarType::Int8, + } + } +} + +impl From for ScalarType { + fn from(scalar_type: PyScalarType) -> Self { + match scalar_type { + PyScalarType::Int8 => ScalarType::Int8, + } + } +} + +#[pyclass(name = "ProductQuantizationConfig")] +#[derive(Clone, Debug, Into)] +pub struct PyProductQuantizationConfig(ProductQuantizationConfig); + +#[pymethods] +impl PyProductQuantizationConfig { + #[new] + #[pyo3(signature = (compression, always_ram = None))] + pub fn new(compression: PyCompressionRatio, always_ram: Option) -> Self { + Self(ProductQuantizationConfig { + compression: CompressionRatio::from(compression), + always_ram, + }) + } + + #[getter] + pub fn compression(&self) -> PyCompressionRatio { + PyCompressionRatio::from(self.0.compression) + } + + #[getter] + pub fn always_ram(&self) -> Option { + self.0.always_ram + } +} + +#[pyclass(name = "CompressionRatio")] +#[derive(Copy, Clone, Debug)] +pub enum PyCompressionRatio { + X4, + X8, + X16, + X32, + X64, +} + +impl From for PyCompressionRatio { + fn from(compression: CompressionRatio) -> Self { + match compression { + CompressionRatio::X4 => PyCompressionRatio::X4, + CompressionRatio::X8 => PyCompressionRatio::X8, + CompressionRatio::X16 => PyCompressionRatio::X16, + CompressionRatio::X32 => PyCompressionRatio::X32, + CompressionRatio::X64 => PyCompressionRatio::X64, + } + } +} + +impl From for CompressionRatio { + fn from(compression: PyCompressionRatio) -> Self { + match compression { + PyCompressionRatio::X4 => CompressionRatio::X4, + PyCompressionRatio::X8 => CompressionRatio::X8, + PyCompressionRatio::X16 => CompressionRatio::X16, + PyCompressionRatio::X32 => CompressionRatio::X32, + PyCompressionRatio::X64 => CompressionRatio::X64, + } + } +} + +#[pyclass(name = "BinaryQuantizationConfig")] +#[derive(Clone, Debug, Into)] +pub struct PyBinaryQuantizationConfig(BinaryQuantizationConfig); + +#[pymethods] +impl PyBinaryQuantizationConfig { + #[new] + #[pyo3(signature = (always_ram = None, encoding = None, query_encoding = None))] + pub fn new( + always_ram: Option, + encoding: Option, + query_encoding: Option, + ) -> Self { + Self(BinaryQuantizationConfig { + always_ram, + encoding: encoding.map(BinaryQuantizationEncoding::from), + query_encoding: query_encoding.map(BinaryQuantizationQueryEncoding::from), + }) + } + + #[getter] + pub fn always_ram(&self) -> Option { + self.0.always_ram + } + + #[getter] + pub fn encoding(&self) -> Option { + self.0.encoding.map(PyBinaryQuantizationEncoding::from) + } + + #[getter] + pub fn query_encoding(&self) -> Option { + self.0 + .query_encoding + .map(PyBinaryQuantizationQueryEncoding::from) + } +} + +#[pyclass(name = "BinaryQuantizationEncoding")] +#[derive(Copy, Clone, Debug)] +pub enum PyBinaryQuantizationEncoding { + OneBit, + TwoBits, + OneAndHalfBits, +} + +impl From for PyBinaryQuantizationEncoding { + fn from(encoding: BinaryQuantizationEncoding) -> Self { + match encoding { + BinaryQuantizationEncoding::OneBit => PyBinaryQuantizationEncoding::OneBit, + BinaryQuantizationEncoding::TwoBits => PyBinaryQuantizationEncoding::TwoBits, + BinaryQuantizationEncoding::OneAndHalfBits => { + PyBinaryQuantizationEncoding::OneAndHalfBits + } + } + } +} + +impl From for BinaryQuantizationEncoding { + fn from(encoding: PyBinaryQuantizationEncoding) -> Self { + match encoding { + PyBinaryQuantizationEncoding::OneBit => BinaryQuantizationEncoding::OneBit, + PyBinaryQuantizationEncoding::TwoBits => BinaryQuantizationEncoding::TwoBits, + PyBinaryQuantizationEncoding::OneAndHalfBits => { + BinaryQuantizationEncoding::OneAndHalfBits + } + } + } +} + +#[pyclass(name = "BinaryQuantizationQueryEncoding")] +#[derive(Copy, Clone, Debug)] +pub enum PyBinaryQuantizationQueryEncoding { + Default, + Binary, + Scalar4Bits, + Scalar8Bits, +} + +impl From for PyBinaryQuantizationQueryEncoding { + fn from(encoding: BinaryQuantizationQueryEncoding) -> Self { + match encoding { + BinaryQuantizationQueryEncoding::Default => PyBinaryQuantizationQueryEncoding::Default, + BinaryQuantizationQueryEncoding::Binary => PyBinaryQuantizationQueryEncoding::Binary, + BinaryQuantizationQueryEncoding::Scalar4Bits => { + PyBinaryQuantizationQueryEncoding::Scalar4Bits + } + BinaryQuantizationQueryEncoding::Scalar8Bits => { + PyBinaryQuantizationQueryEncoding::Scalar8Bits + } + } + } +} + +impl From for BinaryQuantizationQueryEncoding { + fn from(encoding: PyBinaryQuantizationQueryEncoding) -> Self { + match encoding { + PyBinaryQuantizationQueryEncoding::Default => BinaryQuantizationQueryEncoding::Default, + PyBinaryQuantizationQueryEncoding::Binary => BinaryQuantizationQueryEncoding::Binary, + PyBinaryQuantizationQueryEncoding::Scalar4Bits => { + BinaryQuantizationQueryEncoding::Scalar4Bits + } + PyBinaryQuantizationQueryEncoding::Scalar8Bits => { + BinaryQuantizationQueryEncoding::Scalar8Bits + } + } + } +} diff --git a/lib/edge/python/src/config/sparse_vector_data.rs b/lib/edge/python/src/config/sparse_vector_data.rs new file mode 100644 index 0000000000..bcd7d4c2bb --- /dev/null +++ b/lib/edge/python/src/config/sparse_vector_data.rs @@ -0,0 +1,186 @@ +use std::collections::HashMap; +use std::mem; + +use bytemuck::TransparentWrapper; +use derive_more::Into; +use pyo3::prelude::*; +use segment::data_types::modifier::Modifier; +use segment::index::sparse_index::sparse_index_config::{SparseIndexConfig, SparseIndexType}; +use segment::types::*; + +use super::vector_data::*; + +#[pyclass(name = "SparseVectorDataConfig")] +#[derive(Copy, Clone, Debug, Into, TransparentWrapper)] +#[repr(transparent)] +pub struct PySparseVectorDataConfig(pub SparseVectorDataConfig); + +impl PySparseVectorDataConfig { + pub fn peel_map(map: HashMap) -> HashMap + where + Self: TransparentWrapper, + { + unsafe { mem::transmute(map) } + } + + pub fn wrap_map_ref(map: &HashMap) -> &HashMap + where + Self: TransparentWrapper, + { + unsafe { mem::transmute(map) } + } +} + +#[pymethods] +impl PySparseVectorDataConfig { + #[new] + pub fn new( + index: PySparseIndexConfig, + storage_type: PySparseVectorStorageType, + modifier: Option, + ) -> Self { + Self(SparseVectorDataConfig { + index: SparseIndexConfig::from(index), + storage_type: SparseVectorStorageType::from(storage_type), + modifier: modifier.map(Modifier::from), + }) + } + + #[getter] + pub fn index(&self) -> PySparseIndexConfig { + PySparseIndexConfig(self.0.index) + } + + #[getter] + pub fn storage_type(&self) -> PySparseVectorStorageType { + PySparseVectorStorageType::from(self.0.storage_type) + } + + #[getter] + pub fn modifier(&self) -> Option { + self.0.modifier.map(PyModifier::from) + } +} + +impl<'py> IntoPyObject<'py> for &PySparseVectorDataConfig { + type Target = PySparseVectorDataConfig; + type Output = Bound<'py, Self::Target>; + type Error = PyErr; // Infallible + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + IntoPyObject::into_pyobject(*self, py) + } +} + +#[pyclass(name = "SparseIndexConfig")] +#[derive(Copy, Clone, Debug, Into)] +pub struct PySparseIndexConfig(SparseIndexConfig); + +#[pymethods] +impl PySparseIndexConfig { + #[new] + pub fn new( + full_scan_threshold: Option, + index_type: PySparseIndexType, + datatype: Option, + ) -> Self { + Self(SparseIndexConfig { + full_scan_threshold, + index_type: SparseIndexType::from(index_type), + datatype: datatype.map(VectorStorageDatatype::from), + }) + } + + #[getter] + pub fn full_scan_threshold(&self) -> Option { + self.0.full_scan_threshold + } + + #[getter] + pub fn index_type(&self) -> PySparseIndexType { + PySparseIndexType::from(self.0.index_type) + } + + #[getter] + pub fn datatype(&self) -> Option { + self.0.datatype.map(PyVectorStorageDatatype::from) + } +} + +#[pyclass(name = "SparseIndexType")] +#[derive(Copy, Clone, Debug)] +pub enum PySparseIndexType { + MutableRam, + ImmutableRam, + Mmap, +} + +impl From for PySparseIndexType { + fn from(index_type: SparseIndexType) -> Self { + match index_type { + SparseIndexType::MutableRam => PySparseIndexType::MutableRam, + SparseIndexType::ImmutableRam => PySparseIndexType::ImmutableRam, + SparseIndexType::Mmap => PySparseIndexType::Mmap, + } + } +} + +impl From for SparseIndexType { + fn from(index_type: PySparseIndexType) -> Self { + match index_type { + PySparseIndexType::MutableRam => SparseIndexType::MutableRam, + PySparseIndexType::ImmutableRam => SparseIndexType::ImmutableRam, + PySparseIndexType::Mmap => SparseIndexType::Mmap, + } + } +} + +#[pyclass(name = "SparseVectorStorageType")] +#[derive(Copy, Clone, Debug)] +pub enum PySparseVectorStorageType { + Mmap, +} + +impl From for PySparseVectorStorageType { + fn from(storage_type: SparseVectorStorageType) -> Self { + #[allow(unreachable_patterns)] + #[allow(clippy::match_wildcard_for_single_variants)] + match storage_type { + SparseVectorStorageType::Mmap => PySparseVectorStorageType::Mmap, + _ => unimplemented!("RocksDB-backed storage types are not supported by Qdrant Edge"), + } + } +} + +impl From for SparseVectorStorageType { + fn from(storage_type: PySparseVectorStorageType) -> Self { + match storage_type { + PySparseVectorStorageType::Mmap => SparseVectorStorageType::Mmap, + } + } +} + +#[pyclass(name = "Modifier")] +#[derive(Copy, Clone, Debug)] +pub enum PyModifier { + None, + Idf, +} + +impl From for PyModifier { + fn from(modifier: Modifier) -> Self { + match modifier { + Modifier::None => PyModifier::None, + Modifier::Idf => PyModifier::Idf, + } + } +} + +impl From for Modifier { + fn from(modifier: PyModifier) -> Self { + match modifier { + PyModifier::None => Modifier::None, + PyModifier::Idf => Modifier::Idf, + } + } +} diff --git a/lib/edge/python/src/config/vector_data.rs b/lib/edge/python/src/config/vector_data.rs new file mode 100644 index 0000000000..19ace3bb60 --- /dev/null +++ b/lib/edge/python/src/config/vector_data.rs @@ -0,0 +1,343 @@ +use std::collections::HashMap; +use std::mem; + +use bytemuck::TransparentWrapper; +use derive_more::Into; +use pyo3::IntoPyObjectExt as _; +use pyo3::prelude::*; +use segment::types::*; + +use super::quantization::*; + +#[pyclass(name = "VectorDataConfig")] +#[derive(Clone, Debug, Into, TransparentWrapper)] +#[repr(transparent)] +pub struct PyVectorDataConfig(pub VectorDataConfig); + +impl PyVectorDataConfig { + pub fn peel_map(map: HashMap) -> HashMap + where + Self: TransparentWrapper, + { + unsafe { mem::transmute(map) } + } + + pub fn wrap_map_ref(map: &HashMap) -> &HashMap + where + Self: TransparentWrapper, + { + unsafe { mem::transmute(map) } + } +} + +#[pymethods] +impl PyVectorDataConfig { + #[new] + pub fn new( + size: usize, + distance: PyDistance, + storage_type: PyVectorStorageType, + index: PyIndexes, + quantization_config: Option, + multivector_config: Option, + datatype: Option, + ) -> Self { + Self(VectorDataConfig { + size, + distance: Distance::from(distance), + storage_type: VectorStorageType::from(storage_type), + index: Indexes::from(index), + quantization_config: quantization_config.map(QuantizationConfig::from), + multivector_config: multivector_config.map(MultiVectorConfig::from), + datatype: datatype.map(VectorStorageDatatype::from), + }) + } + + #[getter] + pub fn size(&self) -> usize { + self.0.size + } + + #[getter] + pub fn distance(&self) -> PyDistance { + PyDistance::from(self.0.distance) + } + + #[getter] + pub fn storage_type(&self) -> PyVectorStorageType { + PyVectorStorageType::from(self.0.storage_type) + } + + #[getter] + pub fn index(&self) -> PyIndexes { + PyIndexes(self.0.index.clone()) + } + + #[getter] + pub fn quantization_config(&self) -> Option { + self.0.quantization_config.clone().map(PyQuantizationConfig) + } + + #[getter] + pub fn multivector_config(&self) -> Option { + self.0.multivector_config.map(PyMultiVectorConfig) + } + + #[getter] + pub fn datatype(&self) -> Option { + self.0.datatype.map(PyVectorStorageDatatype::from) + } +} + +impl<'py> IntoPyObject<'py> for &PyVectorDataConfig { + type Target = PyVectorDataConfig; + type Output = Bound<'py, Self::Target>; + type Error = PyErr; // Infallible + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + Bound::new(py, self.clone()) + } +} + +#[pyclass(name = "Distance")] +#[derive(Copy, Clone, Debug)] +pub enum PyDistance { + Cosine, + Euclid, + Dot, + Manhattan, +} + +impl From for PyDistance { + fn from(distance: Distance) -> Self { + match distance { + Distance::Cosine => PyDistance::Cosine, + Distance::Euclid => PyDistance::Euclid, + Distance::Dot => PyDistance::Dot, + Distance::Manhattan => PyDistance::Manhattan, + } + } +} + +impl From for Distance { + fn from(distance: PyDistance) -> Self { + match distance { + PyDistance::Cosine => Distance::Cosine, + PyDistance::Euclid => Distance::Euclid, + PyDistance::Dot => Distance::Dot, + PyDistance::Manhattan => Distance::Manhattan, + } + } +} + +#[pyclass(name = "VectorStorageType")] +#[derive(Copy, Clone, Debug)] +pub enum PyVectorStorageType { + Memory, + Mmap, + ChunkedMmap, + InRamChunkedMmap, +} + +impl From for PyVectorStorageType { + fn from(storage_type: VectorStorageType) -> Self { + match storage_type { + VectorStorageType::Memory => PyVectorStorageType::Memory, + VectorStorageType::Mmap => PyVectorStorageType::Mmap, + VectorStorageType::ChunkedMmap => PyVectorStorageType::ChunkedMmap, + VectorStorageType::InRamChunkedMmap => PyVectorStorageType::InRamChunkedMmap, + } + } +} + +impl From for VectorStorageType { + fn from(storage_type: PyVectorStorageType) -> Self { + match storage_type { + PyVectorStorageType::Memory => VectorStorageType::Memory, + PyVectorStorageType::Mmap => VectorStorageType::Mmap, + PyVectorStorageType::ChunkedMmap => VectorStorageType::ChunkedMmap, + PyVectorStorageType::InRamChunkedMmap => VectorStorageType::InRamChunkedMmap, + } + } +} + +#[derive(Clone, Debug, Into)] +pub struct PyIndexes(Indexes); + +impl FromPyObject<'_, '_> for PyIndexes { + type Error = PyErr; + + fn extract(indexes: Borrowed<'_, '_, PyAny>) -> PyResult { + #[derive(FromPyObject)] + enum Helper { + Plain(PyPlainIndexConfig), + Hnsw(PyHnswIndexConfig), + } + + fn _variants(indexes: Indexes) { + match indexes { + Indexes::Plain {} => (), + Indexes::Hnsw(_) => (), + } + } + + let indexes = match indexes.extract()? { + Helper::Plain(_) => Indexes::Plain {}, + Helper::Hnsw(hnsw) => Indexes::Hnsw(HnswConfig::from(hnsw)), + }; + + Ok(Self(indexes)) + } +} + +impl<'py> IntoPyObject<'py> for PyIndexes { + type Target = PyAny; + type Output = Bound<'py, Self::Target>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + match self.0 { + Indexes::Plain {} => PyPlainIndexConfig.into_bound_py_any(py), + Indexes::Hnsw(hnsw) => PyHnswIndexConfig(hnsw).into_bound_py_any(py), + } + } +} + +#[pyclass(name = "PlainIndexConfig")] +#[derive(Copy, Clone, Debug, Default, Into)] +pub struct PyPlainIndexConfig; + +#[pymethods] +impl PyPlainIndexConfig { + #[new] + pub fn new() -> Self { + Self + } +} + +#[pyclass(name = "HnswIndexConfig")] +#[derive(Copy, Clone, Debug, Into)] +pub struct PyHnswIndexConfig(HnswConfig); + +#[pymethods] +impl PyHnswIndexConfig { + #[new] + #[pyo3(signature = (m, ef_construct, full_scan_threshold, on_disk=None, payload_m=None, inline_storage=None))] + pub fn new( + m: usize, + ef_construct: usize, + full_scan_threshold: usize, + on_disk: Option, + payload_m: Option, + inline_storage: Option, + ) -> Self { + Self(HnswConfig { + m, + ef_construct, + full_scan_threshold, + max_indexing_threads: 0, + on_disk, + payload_m, + inline_storage, + }) + } + + #[getter] + pub fn m(&self) -> usize { + self.0.m + } + + #[getter] + pub fn ef_construct(&self) -> usize { + self.0.ef_construct + } + + #[getter] + pub fn full_scan_threshold(&self) -> usize { + self.0.full_scan_threshold + } + + #[getter] + pub fn on_disk(&self) -> Option { + self.0.on_disk + } + + #[getter] + pub fn payload_m(&self) -> Option { + self.0.payload_m + } + + #[getter] + pub fn inline_storage(&self) -> Option { + self.0.inline_storage + } +} + +#[pyclass(name = "MultiVectorConfig")] +#[derive(Copy, Clone, Debug, Into)] +pub struct PyMultiVectorConfig(MultiVectorConfig); + +#[pymethods] +impl PyMultiVectorConfig { + #[new] + pub fn new(comparator: PyMultiVectorComparator) -> Self { + Self(MultiVectorConfig { + comparator: MultiVectorComparator::from(comparator), + }) + } + + #[getter] + pub fn comparator(&self) -> PyMultiVectorComparator { + PyMultiVectorComparator::from(self.0.comparator) + } +} + +#[pyclass(name = "MultiVectorComparator")] +#[derive(Copy, Clone, Debug)] +pub enum PyMultiVectorComparator { + MaxSim, +} + +impl From for PyMultiVectorComparator { + fn from(comparator: MultiVectorComparator) -> Self { + match comparator { + MultiVectorComparator::MaxSim => PyMultiVectorComparator::MaxSim, + } + } +} + +impl From for MultiVectorComparator { + fn from(comparator: PyMultiVectorComparator) -> Self { + match comparator { + PyMultiVectorComparator::MaxSim => MultiVectorComparator::MaxSim, + } + } +} + +#[pyclass(name = "VectorStorageDatatype")] +#[derive(Copy, Clone, Debug)] +pub enum PyVectorStorageDatatype { + Float32, + Float16, + Uint8, +} + +impl From for PyVectorStorageDatatype { + fn from(datatype: VectorStorageDatatype) -> Self { + match datatype { + VectorStorageDatatype::Float32 => PyVectorStorageDatatype::Float32, + VectorStorageDatatype::Float16 => PyVectorStorageDatatype::Float16, + VectorStorageDatatype::Uint8 => PyVectorStorageDatatype::Uint8, + } + } +} + +impl From for VectorStorageDatatype { + fn from(datatype: PyVectorStorageDatatype) -> Self { + match datatype { + PyVectorStorageDatatype::Float32 => VectorStorageDatatype::Float32, + PyVectorStorageDatatype::Float16 => VectorStorageDatatype::Float16, + PyVectorStorageDatatype::Uint8 => VectorStorageDatatype::Uint8, + } + } +} diff --git a/lib/edge/python/src/lib.rs b/lib/edge/python/src/lib.rs index 76b887dee1..8b33763064 100644 --- a/lib/edge/python/src/lib.rs +++ b/lib/edge/python/src/lib.rs @@ -23,12 +23,24 @@ mod qdrant_edge { #[pymodule_export] use super::PyShard; #[pymodule_export] - use super::config::{ - PyDistance, PyIndexes, PyMultiVectorComparator, PyMultiVectorConfig, PyPayloadStorageType, - PyQuantizationConfig, PySegmentConfig, PySparseVectorDataConfig, PyVectorDataConfig, - PyVectorStorageDatatype, PyVectorStorageType, + use super::config::quantization::{ + PyBinaryQuantizationConfig, PyBinaryQuantizationEncoding, + PyBinaryQuantizationQueryEncoding, PyCompressionRatio, PyProductQuantizationConfig, + PyScalarQuantizationConfig, PyScalarType, }; #[pymodule_export] + use super::config::sparse_vector_data::{ + PyModifier, PySparseIndexConfig, PySparseIndexType, PySparseVectorDataConfig, + PySparseVectorStorageType, + }; + #[pymodule_export] + use super::config::vector_data::{ + PyDistance, PyHnswIndexConfig, PyMultiVectorComparator, PyMultiVectorConfig, + PyPlainIndexConfig, PyVectorDataConfig, PyVectorStorageDatatype, PyVectorStorageType, + }; + #[pymodule_export] + use super::config::{PyPayloadStorageType, PySegmentConfig}; + #[pymodule_export] use super::query::{ PyDirection, PyFusion, PyMmr, PyOrderBy, PyPrefetch, PyQueryRequest, PySample, }; diff --git a/lib/edge/python/src/query.rs b/lib/edge/python/src/query.rs index 06340b43c5..8198cbcf68 100644 --- a/lib/edge/python/src/query.rs +++ b/lib/edge/python/src/query.rs @@ -76,10 +76,10 @@ impl PyPrefetch { #[derive(Clone, Debug, Into)] pub struct PyScoringQuery(ScoringQuery); -impl<'py> FromPyObject<'_, 'py> for PyScoringQuery { +impl FromPyObject<'_, '_> for PyScoringQuery { type Error = PyErr; - fn extract(query: Borrowed<'_, 'py, PyAny>) -> PyResult { + fn extract(query: Borrowed<'_, '_, PyAny>) -> PyResult { #[derive(FromPyObject)] enum Helper { Vector(PyQuery), @@ -180,10 +180,10 @@ impl From for Direction { #[derive(Clone, Debug, Into)] pub struct PyStartFrom(StartFrom); -impl<'py> FromPyObject<'_, 'py> for PyStartFrom { +impl FromPyObject<'_, '_> for PyStartFrom { type Error = PyErr; - fn extract(start_from: Borrowed<'_, 'py, PyAny>) -> PyResult { + fn extract(start_from: Borrowed<'_, '_, PyAny>) -> PyResult { #[derive(FromPyObject)] enum Helper { Integer(IntPayloadType), @@ -220,7 +220,7 @@ impl<'py> IntoPyObject<'py> for PyStartFrom { type Output = Bound<'py, Self::Target>; type Error = PyErr; - fn into_pyobject(self, py: Python<'py>) -> std::result::Result { + fn into_pyobject(self, py: Python<'py>) -> PyResult { IntoPyObject::into_pyobject(&self, py) } } @@ -230,7 +230,7 @@ impl<'py> IntoPyObject<'py> for &PyStartFrom { type Output = Bound<'py, Self::Target>; type Error = PyErr; - fn into_pyobject(self, py: Python<'py>) -> std::result::Result { + fn into_pyobject(self, py: Python<'py>) -> PyResult { match &self.0 { StartFrom::Integer(int) => int.into_bound_py_any(py), StartFrom::Float(float) => float.into_bound_py_any(py), diff --git a/lib/edge/python/src/search.rs b/lib/edge/python/src/search.rs index a9400519ef..8cc4a3ea1b 100644 --- a/lib/edge/python/src/search.rs +++ b/lib/edge/python/src/search.rs @@ -97,9 +97,10 @@ impl PyAcornSearchParams { #[derive(Clone, Debug, Into)] pub struct PyWithVector(WithVector); -impl<'py> FromPyObject<'_, 'py> for PyWithVector { +impl FromPyObject<'_, '_> for PyWithVector { type Error = PyErr; - fn extract(with_vector: Borrowed<'_, 'py, PyAny>) -> PyResult { + + fn extract(with_vector: Borrowed<'_, '_, PyAny>) -> PyResult { #[derive(FromPyObject)] enum Helper { Bool(bool), @@ -127,7 +128,7 @@ impl<'py> IntoPyObject<'py> for PyWithVector { type Output = Bound<'py, Self::Target>; type Error = PyErr; // Infallible? - fn into_pyobject(self, py: Python<'py>) -> Result { + fn into_pyobject(self, py: Python<'py>) -> PyResult { IntoPyObject::into_pyobject(&self, py) } } @@ -137,7 +138,7 @@ impl<'py> IntoPyObject<'py> for &PyWithVector { type Output = Bound<'py, Self::Target>; type Error = PyErr; // Infallible? - fn into_pyobject(self, py: Python<'py>) -> Result { + fn into_pyobject(self, py: Python<'py>) -> PyResult { match &self.0 { WithVector::Bool(bool) => bool.into_bound_py_any(py), WithVector::Selector(vectors) => vectors.into_bound_py_any(py), @@ -148,9 +149,10 @@ impl<'py> IntoPyObject<'py> for &PyWithVector { #[derive(Clone, Debug, Into)] pub struct PyWithPayload(WithPayloadInterface); -impl<'py> FromPyObject<'_, 'py> for PyWithPayload { +impl FromPyObject<'_, '_> for PyWithPayload { type Error = PyErr; - fn extract(with_payload: Borrowed<'_, 'py, PyAny>) -> PyResult { + + fn extract(with_payload: Borrowed<'_, '_, PyAny>) -> PyResult { #[derive(FromPyObject)] enum Helper { Bool(bool), @@ -183,7 +185,7 @@ impl<'py> IntoPyObject<'py> for PyWithPayload { type Output = Bound<'py, Self::Target>; type Error = PyErr; // Infallible? - fn into_pyobject(self, py: Python<'py>) -> Result { + fn into_pyobject(self, py: Python<'py>) -> PyResult { IntoPyObject::into_pyobject(&self, py) } } @@ -193,7 +195,7 @@ impl<'py> IntoPyObject<'py> for &PyWithPayload { type Output = Bound<'py, Self::Target>; type Error = PyErr; // Infallible? - fn into_pyobject(self, py: Python<'py>) -> Result { + fn into_pyobject(self, py: Python<'py>) -> PyResult { match &self.0 { WithPayloadInterface::Bool(bool) => bool.into_bound_py_any(py), WithPayloadInterface::Fields(fields) => { @@ -210,9 +212,10 @@ impl<'py> IntoPyObject<'py> for &PyWithPayload { #[repr(transparent)] pub struct PyPayloadSelector(PayloadSelector); -impl<'py> FromPyObject<'_, 'py> for PyPayloadSelector { +impl FromPyObject<'_, '_> for PyPayloadSelector { type Error = PyErr; - fn extract(selector: Borrowed<'_, 'py, PyAny>) -> PyResult { + + fn extract(selector: Borrowed<'_, '_, PyAny>) -> PyResult { let selector = match selector.extract()? { PyPayloadSelectorInterface::Include(keys) => { PayloadSelector::Include(PayloadSelectorInclude { @@ -235,7 +238,7 @@ impl<'py> IntoPyObject<'py> for PyPayloadSelector { type Output = Bound<'py, Self::Target>; type Error = PyErr; // Infallible? - fn into_pyobject(self, py: Python<'py>) -> Result { + fn into_pyobject(self, py: Python<'py>) -> PyResult { let selector = match self.0 { PayloadSelector::Include(PayloadSelectorInclude { include }) => { PyPayloadSelectorInterface::Include(PyJsonPath::wrap_vec(include)) diff --git a/lib/edge/python/src/types/filter/condition.rs b/lib/edge/python/src/types/filter/condition.rs index fa85f26142..12e41251e0 100644 --- a/lib/edge/python/src/types/filter/condition.rs +++ b/lib/edge/python/src/types/filter/condition.rs @@ -2,6 +2,7 @@ use bytemuck::TransparentWrapper; use derive_more::Into; use pyo3::IntoPyObjectExt as _; use pyo3::prelude::*; +use segment::json_path::JsonPath; use segment::types::*; use segment::utils::maybe_arc::MaybeArc; @@ -11,9 +12,10 @@ use crate::types::*; #[repr(transparent)] pub struct PyCondition(pub Condition); -impl<'py> FromPyObject<'_, 'py> for PyCondition { +impl FromPyObject<'_, '_> for PyCondition { type Error = PyErr; - fn extract(condition: Borrowed<'_, 'py, PyAny>) -> PyResult { + + fn extract(condition: Borrowed<'_, '_, PyAny>) -> PyResult { #[derive(FromPyObject)] #[expect(clippy::large_enum_variant)] enum Helper { @@ -45,7 +47,7 @@ impl<'py> IntoPyObject<'py> for PyCondition { type Output = Bound<'py, PyAny>; type Error = PyErr; // Infallible - fn into_pyobject(self, py: Python<'py>) -> Result { + fn into_pyobject(self, py: Python<'py>) -> PyResult { match self.0 { Condition::Field(field) => PyFieldCondition(field).into_bound_py_any(py), Condition::IsEmpty(is_empty) => PyIsEmptyCondition(is_empty).into_bound_py_any(py), @@ -72,7 +74,9 @@ impl PyIsEmptyCondition { #[new] pub fn new(key: PyJsonPath) -> Result { Ok(Self(IsEmptyCondition { - is_empty: PayloadField { key: key.into() }, + is_empty: PayloadField { + key: JsonPath::from(key), + }, })) } } @@ -86,7 +90,9 @@ impl PyIsNullCondition { #[new] pub fn new(key: PyJsonPath) -> Result { Ok(Self(IsNullCondition { - is_null: PayloadField { key: key.into() }, + is_null: PayloadField { + key: JsonPath::from(key), + }, })) } } diff --git a/lib/edge/python/src/types/filter/geo.rs b/lib/edge/python/src/types/filter/geo.rs index 797d7b2a40..932838bdc6 100644 --- a/lib/edge/python/src/types/filter/geo.rs +++ b/lib/edge/python/src/types/filter/geo.rs @@ -1,24 +1,10 @@ +use bytemuck::{TransparentWrapper, TransparentWrapperAlloc as _}; use derive_more::Into; use ordered_float::OrderedFloat; +use pyo3::IntoPyObjectExt as _; use pyo3::exceptions::PyValueError; -use pyo3::{PyErr, pyclass, pymethods}; -use segment::types::{ - GeoBoundingBox, GeoLineString, GeoPoint, GeoPolygon, GeoPolygonShadow, GeoRadius, -}; - -#[pyclass(name = "GeoPoint")] -#[derive(Clone, Debug, Into)] -pub struct PyGeoPoint(pub GeoPoint); - -#[pymethods] -impl PyGeoPoint { - #[new] - pub fn new(lon: f64, lat: f64) -> Result { - Ok(Self(GeoPoint::new(lon, lat).map_err(|err| { - PyErr::new::(err.to_string()) - })?)) - } -} +use pyo3::prelude::*; +use segment::types::*; #[pyclass(name = "GeoBoundingBox")] #[derive(Clone, Debug, Into)] @@ -29,8 +15,8 @@ impl PyGeoBoundingBox { #[new] pub fn new(top_left: PyGeoPoint, bottom_right: PyGeoPoint) -> Self { Self(GeoBoundingBox { - top_left: top_left.0, - bottom_right: bottom_right.0, + top_left: GeoPoint::from(top_left), + bottom_right: GeoPoint::from(bottom_right), }) } } @@ -44,7 +30,7 @@ impl PyGeoRadius { #[new] pub fn new(center: PyGeoPoint, radius: f64) -> Self { Self(GeoRadius { - center: center.0, + center: GeoPoint::from(center), radius: OrderedFloat(radius), }) } @@ -59,30 +45,59 @@ impl PyGeoPolygon { #[new] #[pyo3(signature = (exterior, interiors=None))] pub fn new( - exterior: Vec, - interiors: Option>>, + exterior: PyGeoLineString, + interiors: Option>, ) -> Result { - let exterior = GeoLineString { - points: exterior.into_iter().map(GeoPoint::from).collect(), - }; - - let interiors = interiors.map(|interiors| { - interiors - .into_iter() - .map(|ring| GeoLineString { - points: ring.into_iter().map(GeoPoint::from).collect(), - }) - .collect() - }); - let shadow = GeoPolygonShadow { - exterior, - interiors, + exterior: GeoLineString::from(exterior), + interiors: interiors.map(PyGeoLineString::peel_vec), }; - let polygon = GeoPolygon::try_from(shadow) - .map_err(|err| PyErr::new::(err.to_string()))?; + let polygon = + GeoPolygon::try_from(shadow).map_err(|err| PyValueError::new_err(err.to_string()))?; Ok(Self(polygon)) } } + +#[pyclass(name = "GeoPoint")] +#[derive(Clone, Debug, Into, TransparentWrapper)] +#[repr(transparent)] +pub struct PyGeoPoint(pub GeoPoint); + +#[pymethods] +impl PyGeoPoint { + #[new] + pub fn new(lon: f64, lat: f64) -> Result { + let point = + GeoPoint::new(lon, lat).map_err(|err| PyValueError::new_err(err.to_string()))?; + + Ok(Self(point)) + } +} + +#[derive(Clone, Debug, Into, TransparentWrapper)] +#[repr(transparent)] +pub struct PyGeoLineString(GeoLineString); + +impl FromPyObject<'_, '_> for PyGeoLineString { + type Error = PyErr; + + fn extract(points: Borrowed<'_, '_, PyAny>) -> PyResult { + let points = points.extract()?; + + Ok(Self(GeoLineString { + points: PyGeoPoint::peel_vec(points), + })) + } +} + +impl<'py> IntoPyObject<'py> for PyGeoLineString { + type Target = PyAny; // PyList + type Output = Bound<'py, Self::Target>; + type Error = PyErr; // Infallible + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + PyGeoPoint::wrap_vec(self.0.points).into_bound_py_any(py) + } +} diff --git a/lib/edge/python/src/types/filter/match.rs b/lib/edge/python/src/types/filter/match.rs index 0675803b6d..04a9e76f57 100644 --- a/lib/edge/python/src/types/filter/match.rs +++ b/lib/edge/python/src/types/filter/match.rs @@ -1,11 +1,8 @@ use derive_more::Into; -use pyo3::{FromPyObject, IntoPyObject, pyclass, pymethods}; -use segment::types::{ - AnyVariants, IntPayloadType, Match, MatchAny, MatchExcept, MatchPhrase, MatchText, - MatchTextAny, MatchValue, ValueVariants, -}; +use pyo3::prelude::*; +use segment::types::*; -#[derive(Clone, Debug, IntoPyObject, FromPyObject)] +#[derive(Clone, Debug, FromPyObject, IntoPyObject)] pub enum PyMatch { Value(PyMatchValue), Text(PyMatchText), @@ -18,12 +15,12 @@ pub enum PyMatch { impl From for Match { fn from(value: PyMatch) -> Self { match value { - PyMatch::Value(v) => Match::Value(MatchValue::from(v)), - PyMatch::Text(v) => Match::Text(MatchText::from(v)), - PyMatch::TextAny(v) => Match::TextAny(MatchTextAny::from(v)), - PyMatch::Phrase(v) => Match::Phrase(MatchPhrase::from(v)), - PyMatch::Any(v) => Match::Any(MatchAny::from(v)), - PyMatch::Except(v) => Match::Except(MatchExcept::from(v)), + PyMatch::Value(value) => Match::Value(MatchValue::from(value)), + PyMatch::Text(text) => Match::Text(MatchText::from(text)), + PyMatch::TextAny(text_any) => Match::TextAny(MatchTextAny::from(text_any)), + PyMatch::Phrase(phrase) => Match::Phrase(MatchPhrase::from(phrase)), + PyMatch::Any(any) => Match::Any(MatchAny::from(any)), + PyMatch::Except(except) => Match::Except(MatchExcept::from(except)), } } } @@ -31,12 +28,12 @@ impl From for Match { impl From for PyMatch { fn from(value: Match) -> Self { match value { - Match::Value(v) => PyMatch::Value(PyMatchValue(v)), - Match::Text(v) => PyMatch::Text(PyMatchText(v)), - Match::TextAny(v) => PyMatch::TextAny(PyMatchTextAny(v)), - Match::Phrase(v) => PyMatch::Phrase(PyMatchPhrase(v)), - Match::Any(v) => PyMatch::Any(PyMatchAny(v)), - Match::Except(v) => PyMatch::Except(PyMatchExcept(v)), + Match::Value(value) => PyMatch::Value(PyMatchValue(value)), + Match::Text(text) => PyMatch::Text(PyMatchText(text)), + Match::TextAny(text_any) => PyMatch::TextAny(PyMatchTextAny(text_any)), + Match::Phrase(phrase) => PyMatch::Phrase(PyMatchPhrase(phrase)), + Match::Any(any) => PyMatch::Any(PyMatchAny(any)), + Match::Except(except) => PyMatch::Except(PyMatchExcept(except)), } } } @@ -55,6 +52,23 @@ impl PyMatchValue { } } +#[derive(Clone, Debug, FromPyObject, IntoPyObject)] +pub enum PyValueVariants { + String(String), + Integer(IntPayloadType), + Bool(bool), +} + +impl From for ValueVariants { + fn from(value: PyValueVariants) -> Self { + match value { + PyValueVariants::String(str) => ValueVariants::String(str), + PyValueVariants::Integer(int) => ValueVariants::Integer(int), + PyValueVariants::Bool(bool) => ValueVariants::Bool(bool), + } + } +} + #[pyclass(name = "MatchText")] #[derive(Clone, Debug, Into)] pub struct PyMatchText(pub MatchText); @@ -119,34 +133,17 @@ impl PyMatchExcept { } } -#[derive(Clone, Debug, IntoPyObject, FromPyObject)] -pub enum PyValueVariants { - String(String), - Integer(IntPayloadType), - Bool(bool), -} - -impl From for ValueVariants { - fn from(value: PyValueVariants) -> Self { - match value { - PyValueVariants::String(s) => ValueVariants::String(s), - PyValueVariants::Integer(i) => ValueVariants::Integer(i), - PyValueVariants::Bool(b) => ValueVariants::Bool(b), - } - } -} - -#[derive(Clone, Debug, IntoPyObject, FromPyObject)] +#[derive(Clone, Debug, FromPyObject, IntoPyObject)] pub enum PyAnyVariants { Strings(Vec), Integers(Vec), } impl From for AnyVariants { - fn from(value: PyAnyVariants) -> Self { - match value { - PyAnyVariants::Strings(s) => AnyVariants::Strings(s.into_iter().collect()), - PyAnyVariants::Integers(i) => AnyVariants::Integers(i.into_iter().collect()), + fn from(any: PyAnyVariants) -> Self { + match any { + PyAnyVariants::Strings(str) => AnyVariants::Strings(str.into_iter().collect()), + PyAnyVariants::Integers(int) => AnyVariants::Integers(int.into_iter().collect()), } } } diff --git a/lib/edge/python/src/types/filter/min_should.rs b/lib/edge/python/src/types/filter/min_should.rs index 8c61e9b121..ee7d006e7c 100644 --- a/lib/edge/python/src/types/filter/min_should.rs +++ b/lib/edge/python/src/types/filter/min_should.rs @@ -1,6 +1,7 @@ +use bytemuck::TransparentWrapperAlloc as _; use derive_more::Into; use pyo3::prelude::*; -use segment::types::{Condition, MinShould}; +use segment::types::MinShould; use crate::types::filter::condition::PyCondition; @@ -12,9 +13,8 @@ pub struct PyMinShould(pub MinShould); impl PyMinShould { #[new] pub fn new(conditions: Vec, min_count: usize) -> Self { - let conditions = conditions.into_iter().map(Condition::from).collect(); Self(MinShould { - conditions, + conditions: PyCondition::peel_vec(conditions), min_count, }) } diff --git a/lib/edge/python/src/types/filter/range.rs b/lib/edge/python/src/types/filter/range.rs index 9433913d84..f960506bdf 100644 --- a/lib/edge/python/src/types/filter/range.rs +++ b/lib/edge/python/src/types/filter/range.rs @@ -4,7 +4,7 @@ use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use segment::types::*; -#[derive(Copy, Clone, Debug, IntoPyObject, FromPyObject)] +#[derive(Copy, Clone, Debug, FromPyObject, IntoPyObject)] pub enum PyRange { Float(PyRangeFloat), DateTime(PyRangeDateTime), diff --git a/lib/edge/python/src/types/filter/value_count.rs b/lib/edge/python/src/types/filter/value_count.rs index eaedd679f2..8258fc68ed 100644 --- a/lib/edge/python/src/types/filter/value_count.rs +++ b/lib/edge/python/src/types/filter/value_count.rs @@ -1,5 +1,5 @@ use derive_more::Into; -use pyo3::{pyclass, pymethods}; +use pyo3::prelude::*; use segment::types::ValuesCount; #[pyclass(name = "ValuesCount")] diff --git a/lib/edge/python/src/types/formula.rs b/lib/edge/python/src/types/formula.rs index bcad1d65f1..9235fcecd2 100644 --- a/lib/edge/python/src/types/formula.rs +++ b/lib/edge/python/src/types/formula.rs @@ -35,9 +35,10 @@ impl PyFormula { #[repr(transparent)] pub struct PyExpression(ExpressionInternal); -impl<'py> FromPyObject<'_, 'py> for PyExpression { +impl FromPyObject<'_, '_> for PyExpression { type Error = PyErr; - fn extract(helper: Borrowed<'_, 'py, PyAny>) -> PyResult { + + fn extract(helper: Borrowed<'_, '_, PyAny>) -> PyResult { let expr = match helper.extract()? { PyExpressionInterface::Constant(val) => ExpressionInternal::Constant(val), PyExpressionInterface::Variable(var) => ExpressionInternal::Variable(var), @@ -101,7 +102,7 @@ impl<'py> IntoPyObject<'py> for PyExpression { type Output = Bound<'py, Self::Target>; type Error = PyErr; - fn into_pyobject(self, py: Python<'py>) -> Result { + fn into_pyobject(self, py: Python<'py>) -> PyResult { let helper = match self.0 { ExpressionInternal::Constant(var) => PyExpressionInterface::Constant(var), ExpressionInternal::Variable(var) => PyExpressionInterface::Variable(var), diff --git a/lib/edge/python/src/types/json_path.rs b/lib/edge/python/src/types/json_path.rs index 3b4b164856..a074803929 100644 --- a/lib/edge/python/src/types/json_path.rs +++ b/lib/edge/python/src/types/json_path.rs @@ -12,9 +12,10 @@ use segment::json_path::JsonPath; #[repr(transparent)] pub struct PyJsonPath(pub JsonPath); -impl<'py> FromPyObject<'_, 'py> for PyJsonPath { +impl FromPyObject<'_, '_> for PyJsonPath { type Error = PyErr; - fn extract(json_path: Borrowed<'_, 'py, PyAny>) -> PyResult { + + fn extract(json_path: Borrowed<'_, '_, PyAny>) -> PyResult { let json_path: String = json_path.extract()?; let json_path = JsonPath::from_str(&json_path) .map_err(|_| PyValueError::new_err(format!("invalid JSON path {json_path}")))?; diff --git a/lib/edge/python/src/types/payload.rs b/lib/edge/python/src/types/payload.rs index 67018e0f66..2c8d868ae6 100644 --- a/lib/edge/python/src/types/payload.rs +++ b/lib/edge/python/src/types/payload.rs @@ -9,9 +9,10 @@ use super::value::*; #[repr(transparent)] pub struct PyPayload(pub Payload); -impl<'py> FromPyObject<'_, 'py> for PyPayload { +impl FromPyObject<'_, '_> for PyPayload { type Error = PyErr; - fn extract(payload: Borrowed<'_, 'py, PyAny>) -> PyResult { + + fn extract(payload: Borrowed<'_, '_, PyAny>) -> PyResult { let payload = value_map_from_py(&payload)?; Ok(Self(Payload(payload))) } @@ -22,7 +23,7 @@ impl<'py> IntoPyObject<'py> for PyPayload { type Output = Bound<'py, Self::Target>; type Error = PyErr; - fn into_pyobject(self, py: Python<'py>) -> Result { + fn into_pyobject(self, py: Python<'py>) -> PyResult { IntoPyObject::into_pyobject(&self, py) } } @@ -32,7 +33,7 @@ impl<'py> IntoPyObject<'py> for &PyPayload { type Output = Bound<'py, Self::Target>; type Error = PyErr; - fn into_pyobject(self, py: Python<'py>) -> Result { + fn into_pyobject(self, py: Python<'py>) -> PyResult { value_map_into_py(&self.0.0, py) } } diff --git a/lib/edge/python/src/types/point_id.rs b/lib/edge/python/src/types/point_id.rs index 2bc5a1e52f..29a0d3529f 100644 --- a/lib/edge/python/src/types/point_id.rs +++ b/lib/edge/python/src/types/point_id.rs @@ -21,9 +21,10 @@ impl PyPointId { } } -impl<'py> FromPyObject<'_, 'py> for PyPointId { +impl FromPyObject<'_, '_> for PyPointId { type Error = PyErr; - fn extract(point_id: Borrowed<'_, 'py, PyAny>) -> PyResult { + + fn extract(point_id: Borrowed<'_, '_, PyAny>) -> PyResult { #[derive(FromPyObject)] enum Helper { NumId(u64), @@ -59,7 +60,7 @@ impl<'py> IntoPyObject<'py> for PyPointId { type Output = Bound<'py, PyAny>; type Error = PyErr; // Infallible - fn into_pyobject(self, py: Python<'py>) -> Result { + fn into_pyobject(self, py: Python<'py>) -> PyResult { IntoPyObject::into_pyobject(&self, py) } } @@ -69,7 +70,7 @@ impl<'py> IntoPyObject<'py> for &PyPointId { type Output = Bound<'py, PyAny>; type Error = PyErr; // Infallible - fn into_pyobject(self, py: Python<'py>) -> Result { + fn into_pyobject(self, py: Python<'py>) -> PyResult { 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/point_vectors.rs b/lib/edge/python/src/types/point_vectors.rs index 238ca0625e..59b965509f 100644 --- a/lib/edge/python/src/types/point_vectors.rs +++ b/lib/edge/python/src/types/point_vectors.rs @@ -19,4 +19,14 @@ impl PyPointVectors { vector: VectorStructPersisted::from(vector), }) } + + #[getter] + fn id(&self) -> PyPointId { + PyPointId(self.0.id) + } + + #[getter] + fn vector(&self) -> PyVector { + PyVector::from(self.0.vector.clone()) + } } diff --git a/lib/edge/python/src/types/query/mod.rs b/lib/edge/python/src/types/query/mod.rs index 7706bd9623..7902417355 100644 --- a/lib/edge/python/src/types/query/mod.rs +++ b/lib/edge/python/src/types/query/mod.rs @@ -11,9 +11,10 @@ use crate::types::*; #[derive(Clone, Debug, Into)] pub struct PyQuery(QueryEnum); -impl<'py> FromPyObject<'_, 'py> for PyQuery { +impl FromPyObject<'_, '_> for PyQuery { type Error = PyErr; - fn extract(query: Borrowed<'_, 'py, PyAny>) -> PyResult { + + fn extract(query: Borrowed<'_, '_, PyAny>) -> PyResult { let query = match query.extract()? { PyQueryInterface::Nearest { query, using } => QueryEnum::Nearest(NamedQuery { query: VectorInternal::try_from(query)?, @@ -61,7 +62,7 @@ impl<'py> IntoPyObject<'py> for PyQuery { type Output = Bound<'py, Self::Target>; type Error = PyErr; // Infallible? - fn into_pyobject(self, py: Python<'py>) -> Result { + fn into_pyobject(self, py: Python<'py>) -> PyResult { let query = match self.0 { QueryEnum::Nearest(NamedQuery { query, using }) => PyQueryInterface::Nearest { query: PyNamedVector::from(query), diff --git a/lib/edge/python/src/types/value.rs b/lib/edge/python/src/types/value.rs index f9b289b175..da4076b53b 100644 --- a/lib/edge/python/src/types/value.rs +++ b/lib/edge/python/src/types/value.rs @@ -21,9 +21,10 @@ impl PyValue { } } -impl<'py> FromPyObject<'_, 'py> for PyValue { +impl FromPyObject<'_, '_> for PyValue { type Error = PyErr; - fn extract(value: Borrowed<'_, 'py, PyAny>) -> PyResult { + + fn extract(value: Borrowed<'_, '_, PyAny>) -> PyResult { #[derive(FromPyObject)] enum Helper { Bool(bool), @@ -66,7 +67,7 @@ impl<'py> IntoPyObject<'py> for PyValue { type Output = Bound<'py, Self::Target>; type Error = PyErr; - fn into_pyobject(self, py: Python<'py>) -> Result { + fn into_pyobject(self, py: Python<'py>) -> PyResult { IntoPyObject::into_pyobject(&self, py) } } @@ -76,7 +77,7 @@ impl<'py> IntoPyObject<'py> for &PyValue { type Output = Bound<'py, Self::Target>; type Error = PyErr; - fn into_pyobject(self, py: Python<'py>) -> Result { + fn into_pyobject(self, py: Python<'py>) -> PyResult { match &self.0 { serde_json::Value::Null => Ok(py.None().into_bound(py)), serde_json::Value::Bool(bool) => bool.into_bound_py_any(py), diff --git a/lib/edge/python/src/types/vector.rs b/lib/edge/python/src/types/vector.rs index a231490650..8aed0f745d 100644 --- a/lib/edge/python/src/types/vector.rs +++ b/lib/edge/python/src/types/vector.rs @@ -9,7 +9,7 @@ use shard::operations::point_ops::{VectorPersisted, VectorStructPersisted}; use sparse::common::sparse_vector::SparseVector; use sparse::common::types::{DimId, DimWeight}; -#[derive(Clone, Debug, IntoPyObject, FromPyObject)] +#[derive(Clone, Debug, FromPyObject, IntoPyObject)] pub enum PyVector { // Put Int first so ints don't get parsed as floats (since f64 can extract from ints). Single(DenseVector), @@ -85,7 +85,7 @@ impl TryFrom for VectorStructInternal { } } -#[derive(Clone, Debug, IntoPyObject, FromPyObject)] +#[derive(Clone, Debug, FromPyObject, IntoPyObject)] pub enum PyNamedVector { // Put Int first so ints don't get parsed as floats (since f64 can extract from ints). Dense(DenseVector), diff --git a/lib/segment/src/segment_constructor/segment_constructor_base.rs b/lib/segment/src/segment_constructor/segment_constructor_base.rs index 4e8b5c2113..5ab8ee00f3 100644 --- a/lib/segment/src/segment_constructor/segment_constructor_base.rs +++ b/lib/segment/src/segment_constructor/segment_constructor_base.rs @@ -313,7 +313,7 @@ pub(crate) fn open_vector_index( vector_storage, quantized_vectors, payload_index, - hnsw_config: hnsw_config.clone(), + hnsw_config: *hnsw_config, })?), }) } @@ -344,7 +344,7 @@ pub(crate) fn build_vector_index( vector_storage, quantized_vectors, payload_index, - hnsw_config: hnsw_config.clone(), + hnsw_config: *hnsw_config, }, build_args, )?), diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index a2a0dc57eb..9505e767f7 100644 --- a/lib/segment/src/types.rs +++ b/lib/segment/src/types.rs @@ -629,7 +629,9 @@ impl Indexes { } /// Config of HNSW index -#[derive(Debug, Deserialize, Serialize, JsonSchema, Validate, Anonymize, Clone, PartialEq, Eq)] +#[derive( + Copy, Clone, Debug, Eq, PartialEq, Deserialize, Serialize, JsonSchema, Validate, Anonymize, +)] #[serde(rename_all = "snake_case")] #[anonymize(false)] pub struct HnswConfig { @@ -684,7 +686,7 @@ impl HnswConfig { payload_m, on_disk, inline_storage, - } = self.clone(); + } = *self; m != other.m || ef_construct != other.ef_construct @@ -730,7 +732,7 @@ pub enum CompressionRatio { X64, } -#[derive(Default, Debug, Deserialize, Serialize, JsonSchema, Clone, PartialEq, Eq, Hash)] +#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Hash, Deserialize, Serialize, JsonSchema)] #[serde(rename_all = "lowercase")] pub enum ScalarType { #[default] diff --git a/lib/segment/tests/integration/multivector_hnsw_test.rs b/lib/segment/tests/integration/multivector_hnsw_test.rs index 7c8299ef43..e31058979a 100644 --- a/lib/segment/tests/integration/multivector_hnsw_test.rs +++ b/lib/segment/tests/integration/multivector_hnsw_test.rs @@ -124,7 +124,7 @@ fn test_single_multi_and_dense_hnsw_equivalency() { vector_storage: vector_storage.clone(), quantized_vectors: quantized_vectors.clone(), payload_index: segment.payload_index.clone(), - hnsw_config: hnsw_config.clone(), + hnsw_config, }, VectorIndexBuildArgs { permit: permit.clone(),