diff --git a/examples/search.rs b/examples/search.rs index e3ff35a..252cc7b 100644 --- a/examples/search.rs +++ b/examples/search.rs @@ -1,8 +1,8 @@ use anyhow::Result; use qdrant_client::prelude::*; -use qdrant_client::qdrant::Distance; use qdrant_client::qdrant::{ - Condition, CreateCollectionBuilder, Filter, SearchPoints, VectorParamsBuilder, + Condition, CreateCollectionBuilder, Distance, Filter, QuantizationType, + ScalarQuantizationBuilder, SearchPoints, VectorParamsBuilder, }; use serde_json::json; @@ -31,11 +31,8 @@ async fn main() -> Result<()> { .create_collection( &CreateCollectionBuilder::default() .collection_name(collection_name) - .vectors_config( - VectorParamsBuilder::default() - .distance(Distance::Cosine) - .size(10), - ) + .vectors_config(VectorParamsBuilder::new(300, Distance::Cosine)) + .quantization_config(ScalarQuantizationBuilder::new(QuantizationType::Int8)) .build(), ) .await?; diff --git a/src/grpc_ext.rs b/src/grpc_ext.rs index bedafa0..ee3176b 100644 --- a/src/grpc_ext.rs +++ b/src/grpc_ext.rs @@ -8,17 +8,18 @@ use crate::qdrant::{ alias_operations, condition, group_id, points_update_operation, quantization_config, quantization_config_diff, r#match, read_consistency, shard_key, start_from, target_vector, vector_example, vectors_config, vectors_config_diff, with_payload_selector, - with_vectors_selector, AliasOperations, BinaryQuantization, Condition, CreateAlias, - DeleteAlias, Disabled, FieldCondition, Filter, GeoLineString, GeoPoint, GroupId, - HasIdCondition, IntegerIndexParams, IsEmptyCondition, IsNullCondition, ListValue, Match, - NamedVectors, NestedCondition, PayloadExcludeSelector, PayloadIncludeSelector, + with_vectors_selector, AliasOperations, BinaryQuantization, BinaryQuantizationBuilder, + Condition, CreateAlias, DeleteAlias, Disabled, FieldCondition, Filter, GeoLineString, GeoPoint, + GroupId, HasIdCondition, IntegerIndexParams, IsEmptyCondition, IsNullCondition, ListValue, + Match, NamedVectors, NestedCondition, PayloadExcludeSelector, PayloadIncludeSelector, PayloadIndexParams, PointId, PointStruct, PointsIdsList, PointsSelector, PointsUpdateOperation, - ProductQuantization, QuantizationConfig, QuantizationConfigDiff, ReadConsistency, RenameAlias, - RepeatedIntegers, RepeatedStrings, ScalarQuantization, ShardKey, ShardKeySelector, - SparseIndexConfig, SparseIndices, SparseVectorConfig, SparseVectorParams, StartFrom, Struct, - TargetVector, TextIndexParams, Value, Vector, VectorExample, VectorParams, VectorParamsBuilder, - VectorParamsDiff, VectorParamsDiffMap, VectorParamsMap, Vectors, VectorsConfig, - VectorsConfigDiff, VectorsSelector, WithPayloadSelector, WithVectorsSelector, + ProductQuantization, ProductQuantizationBuilder, QuantizationConfig, QuantizationConfigDiff, + ReadConsistency, RenameAlias, RepeatedIntegers, RepeatedStrings, ScalarQuantization, + ScalarQuantizationBuilder, ShardKey, ShardKeySelector, SparseIndexConfig, SparseIndices, + SparseVectorConfig, SparseVectorParams, StartFrom, Struct, TargetVector, TextIndexParams, + Value, Vector, VectorExample, VectorParams, VectorParamsBuilder, VectorParamsDiff, + VectorParamsDiffMap, VectorParamsMap, Vectors, VectorsConfig, VectorsConfigDiff, + VectorsSelector, WithPayloadSelector, WithVectorsSelector, }; use std::collections::HashMap; @@ -772,8 +773,50 @@ impl From> for PointsIdsList { } } +impl From for vectors_config::Config { + fn from(value: VectorParamsBuilder) -> Self { + value.build().into() + } +} + impl From<&mut VectorParamsBuilder> for vectors_config::Config { fn from(value: &mut VectorParamsBuilder) -> Self { value.build().into() } } + +impl From for quantization_config::Quantization { + fn from(value: ScalarQuantizationBuilder) -> Self { + Self::Scalar(value.build()) + } +} + +impl From<&mut ScalarQuantizationBuilder> for quantization_config::Quantization { + fn from(value: &mut ScalarQuantizationBuilder) -> Self { + Self::Scalar(value.build()) + } +} + +impl From for quantization_config::Quantization { + fn from(value: ProductQuantizationBuilder) -> Self { + Self::Product(value.build()) + } +} + +impl From<&mut ProductQuantizationBuilder> for quantization_config::Quantization { + fn from(value: &mut ProductQuantizationBuilder) -> Self { + Self::Product(value.build()) + } +} + +impl From for quantization_config::Quantization { + fn from(value: BinaryQuantizationBuilder) -> Self { + Self::Binary(value.build()) + } +} + +impl From<&mut BinaryQuantizationBuilder> for quantization_config::Quantization { + fn from(value: &mut BinaryQuantizationBuilder) -> Self { + Self::Binary(value.build()) + } +} diff --git a/src/qdrant.rs b/src/qdrant.rs index c644336..e104eae 100644 --- a/src/qdrant.rs +++ b/src/qdrant.rs @@ -1,6 +1,6 @@ // This file is @generated by prost-build. #[derive(derive_builder::Builder)] -#[builder(build_fn(private, error = "std::convert::Infallible", name = "build_inner"))] +#[builder(build_fn(private, name = "build_inner"), custom_constructor)] #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] pub struct VectorParams { @@ -18,7 +18,13 @@ pub struct VectorParams { pub hnsw_config: ::core::option::Option, /// Configuration of vector quantization config. If omitted - the collection configuration will be used #[prost(message, optional, tag = "4")] - #[builder(default, setter(into, strip_option))] + #[builder( + setter(into, strip_option), + field( + ty = "Option", + build = "convert_option(&self.quantization_config)" + ) + )] pub quantization_config: ::core::option::Option, /// If true - serve vectors from disk. If set to false, the vectors will be loaded in RAM. #[prost(bool, optional, tag = "5")] @@ -226,26 +232,34 @@ pub struct SparseIndexConfig { #[prost(bool, optional, tag = "2")] pub on_disk: ::core::option::Option, } +#[derive(derive_builder::Builder)] +#[builder(build_fn(private, error = "std::convert::Infallible", name = "build_inner"))] #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] pub struct WalConfigDiff { /// Size of a single WAL block file #[prost(uint64, optional, tag = "1")] + #[builder(default, setter(strip_option))] pub wal_capacity_mb: ::core::option::Option, /// Number of segments to create in advance #[prost(uint64, optional, tag = "2")] + #[builder(default, setter(strip_option))] pub wal_segments_ahead: ::core::option::Option, } +#[derive(derive_builder::Builder)] +#[builder(build_fn(private, error = "std::convert::Infallible", name = "build_inner"))] #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] pub struct OptimizersConfigDiff { /// /// The minimal fraction of deleted vectors in a segment, required to perform segment optimization #[prost(double, optional, tag = "1")] + #[builder(default, setter(strip_option))] pub deleted_threshold: ::core::option::Option, /// /// The minimal number of vectors in a segment, required to perform segment optimization #[prost(uint64, optional, tag = "2")] + #[builder(default, setter(strip_option))] pub vacuum_min_vector_number: ::core::option::Option, /// /// Target amount of segments the optimizer will try to keep. @@ -257,6 +271,7 @@ pub struct OptimizersConfigDiff { /// It is recommended to select the default number of segments as a factor of the number of search threads, /// so that each segment would be handled evenly by one of the threads. #[prost(uint64, optional, tag = "3")] + #[builder(default, setter(strip_option))] pub default_segment_number: ::core::option::Option, /// /// Do not create segments larger this size (in kilobytes). @@ -268,6 +283,7 @@ pub struct OptimizersConfigDiff { /// Note: 1Kb = 1 vector of size 256 /// If not set, will be automatically selected considering the number of available CPUs. #[prost(uint64, optional, tag = "4")] + #[builder(default, setter(strip_option))] pub max_segment_size: ::core::option::Option, /// /// Maximum size (in kilobytes) of vectors to store in-memory per segment. @@ -279,6 +295,7 @@ pub struct OptimizersConfigDiff { /// /// Note: 1Kb = 1 vector of size 256 #[prost(uint64, optional, tag = "5")] + #[builder(default, setter(strip_option))] pub memmap_threshold: ::core::option::Option, /// /// Maximum size (in kilobytes) of vectors allowed for plain index, exceeding this threshold will enable vector indexing @@ -289,10 +306,12 @@ pub struct OptimizersConfigDiff { /// /// Note: 1kB = 1 vector of size 256. #[prost(uint64, optional, tag = "6")] + #[builder(default, setter(strip_option))] pub indexing_threshold: ::core::option::Option, /// /// Interval between forced flushes. #[prost(uint64, optional, tag = "7")] + #[builder(default, setter(strip_option))] pub flush_interval_sec: ::core::option::Option, /// /// Max number of threads (jobs) for running optimizations per shard. @@ -300,8 +319,11 @@ pub struct OptimizersConfigDiff { /// If null - have no limit and choose dynamically to saturate CPU. /// If 0 - no optimization threads, optimizations will be disabled. #[prost(uint64, optional, tag = "8")] + #[builder(default, setter(strip_option))] pub max_optimization_threads: ::core::option::Option, } +#[derive(derive_builder::Builder)] +#[builder(build_fn(private, name = "build_inner"), custom_constructor)] #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] pub struct ScalarQuantization { @@ -310,11 +332,15 @@ pub struct ScalarQuantization { pub r#type: i32, /// Number of bits to use for quantization #[prost(float, optional, tag = "2")] + #[builder(default, setter(strip_option))] pub quantile: ::core::option::Option, /// If true - quantized vectors always will be stored in RAM, ignoring the config of main storage #[prost(bool, optional, tag = "3")] + #[builder(default, setter(strip_option))] pub always_ram: ::core::option::Option, } +#[derive(derive_builder::Builder)] +#[builder(build_fn(private, name = "build_inner"), custom_constructor)] #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] pub struct ProductQuantization { @@ -323,13 +349,17 @@ pub struct ProductQuantization { pub compression: i32, /// If true - quantized vectors always will be stored in RAM, ignoring the config of main storage #[prost(bool, optional, tag = "2")] + #[builder(default, setter(strip_option))] pub always_ram: ::core::option::Option, } +#[derive(derive_builder::Builder)] +#[builder(build_fn(private, name = "build_inner"), custom_constructor)] #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] pub struct BinaryQuantization { /// If true - quantized vectors always will be stored in RAM, ignoring the config of main storage #[prost(bool, optional, tag = "1")] + #[builder(default, setter(strip_option))] pub always_ram: ::core::option::Option, } #[allow(clippy::derive_partial_eq_without_eq)] @@ -432,7 +462,13 @@ pub struct CreateCollection { pub init_from_collection: ::core::option::Option<::prost::alloc::string::String>, /// Quantization configuration of vector #[prost(message, optional, tag = "14")] - #[builder(default, setter(into, strip_option))] + #[builder( + setter(into, strip_option), + field( + ty = "Option", + build = "convert_option(&self.quantization_config)" + ) + )] pub quantization_config: ::core::option::Option, /// Sharding method #[prost(enumeration = "ShardingMethod", optional, tag = "15")] @@ -7480,4 +7516,42 @@ where builder_type_conversions!(CreateCollection, CreateCollectionBuilder); builder_type_conversions!(VectorParams, VectorParamsBuilder); builder_type_conversions!(HnswConfigDiff, HnswConfigDiffBuilder); +builder_type_conversions!(ScalarQuantization, ScalarQuantizationBuilder); +builder_type_conversions!(ProductQuantization, ProductQuantizationBuilder); +builder_type_conversions!(BinaryQuantization, BinaryQuantizationBuilder); +builder_type_conversions!(OptimizersConfigDiff, OptimizersConfigDiffBuilder); +builder_type_conversions!(WalConfigDiff, WalConfigDiffBuilder); + +impl VectorParamsBuilder { + pub fn new(size: u64, distance: Distance) -> Self { + let mut builder = Self::create_empty(); + builder.size = Some(size); + builder.distance = Some(distance.into()); + builder + } +} + +impl ScalarQuantizationBuilder { + pub fn new(r#type: QuantizationType) -> Self { + let mut builder = Self::create_empty(); + builder.r#type = Some(r#type.into()); + builder + } +} + +impl ProductQuantizationBuilder { + pub fn new(compression: i32) -> Self { + let mut builder = Self::create_empty(); + builder.compression = Some(compression); + builder + } +} + +impl BinaryQuantizationBuilder { + pub fn new(always_ram: bool) -> Self { + let mut builder = Self::create_empty(); + builder.always_ram = Some(Some(always_ram)); + builder + } +} diff --git a/tests/protos.rs b/tests/protos.rs index ac1c4ab..52a9ce0 100644 --- a/tests/protos.rs +++ b/tests/protos.rs @@ -29,6 +29,8 @@ fn protos() { append_file_to_file(GRPC_OUTPUT_FILE, "./tests/protos_append/grpc_macros.rs"); add_builder_macro_impls(GRPC_OUTPUT_FILE, builder_derive_options()); + append_file_to_file(GRPC_OUTPUT_FILE, "./tests/protos_append/builder_ext.rs"); + panic!("proto definitions changed. Stubs recompiled. Please commit the changes.") } @@ -116,7 +118,11 @@ fn configure_builder(builder: Builder) -> Builder { ("VectorParams.size", DEFAULT), ("VectorParams.distance", DEFAULT_INTO), ("VectorParams.hnsw_config", DEFAULT_OPTION_INTO), - ("VectorParams.quantization_config", DEFAULT_OPTION_INTO), + // ("VectorParams.quantization_config", DEFAULT_OPTION_INTO), + ( + "VectorParams.quantization_config", + builder_custom_into!(quantization_config::Quantization, self.quantization_config), + ), ("VectorParams.on_disk", DEFAULT_OPTION), ("VectorParams.datatype", DEFAULT_OPTION_INTO), // Create collection @@ -135,7 +141,11 @@ fn configure_builder(builder: Builder) -> Builder { ("CreateCollection.replication_factor", DEFAULT_OPTION), ("CreateCollection.write_consistency_factor", DEFAULT_OPTION), ("CreateCollection.init_from_collection", DEFAULT_OPTION_INTO), - ("CreateCollection.quantization_config", DEFAULT_OPTION_INTO), + // ("CreateCollection.quantization_config", DEFAULT_OPTION_INTO), ( + ( + "CreateCollection.quantization_config", + builder_custom_into!(quantization_config::Quantization, self.quantization_config), + ), ("CreateCollection.sharding_method", DEFAULT_OPTION), ( "CreateCollection.sparse_vectors_config", @@ -148,6 +158,34 @@ fn configure_builder(builder: Builder) -> Builder { ("HnswConfigDiff.max_indexing_threads", DEFAULT_OPTION), ("HnswConfigDiff.on_disk", DEFAULT_OPTION), ("HnswConfigDiff.payload_m", DEFAULT_OPTION), + // ScalarQuantization + ("ScalarQuantization.quantile", DEFAULT_OPTION), + ("ScalarQuantization.always_ram", DEFAULT_OPTION), + // ProductQuantization + ("ProductQuantization.always_ram", DEFAULT_OPTION), + // BinaryQuantization + ("BinaryQuantization.always_ram", DEFAULT_OPTION), + // OptimizersConfigDiff + ("OptimizersConfigDiff.deleted_threshold", DEFAULT_OPTION), + ( + "OptimizersConfigDiff.vacuum_min_vector_number", + DEFAULT_OPTION, + ), + ( + "OptimizersConfigDiff.default_segment_number", + DEFAULT_OPTION, + ), + ("OptimizersConfigDiff.max_segment_size", DEFAULT_OPTION), + ("OptimizersConfigDiff.memmap_threshold", DEFAULT_OPTION), + ("OptimizersConfigDiff.indexing_threshold", DEFAULT_OPTION), + ("OptimizersConfigDiff.flush_interval_sec", DEFAULT_OPTION), + ( + "OptimizersConfigDiff.max_optimization_threads", + DEFAULT_OPTION, + ), + //WalConfigDiff + ("WalConfigDiff.wal_capacity_mb", DEFAULT_OPTION), + ("WalConfigDiff.wal_segments_ahead", DEFAULT_OPTION), ], builder_derive_options(), ) @@ -160,12 +198,31 @@ fn builder_derive_options() -> &'static [BuildDeriveOptions] { // Infallible allows secure unwrapping and compiler errors on missing fields. const DEFAULT_BUILDER_DERIVE_OPTIONS: &str = "build_fn(private, error = \"std::convert::Infallible\", name = \"build_inner\")"; + const NO_DEFAULT_BUILDER_DERIVE_OPTIONS: &str = + "build_fn(private, name = \"build_inner\"), custom_constructor"; // Tuple structure: (Path, build attributes, 'from' macro generation enabled) &[ ("CreateCollection", DEFAULT_BUILDER_DERIVE_OPTIONS, true), - ("VectorParams", DEFAULT_BUILDER_DERIVE_OPTIONS, true), + ("VectorParams", NO_DEFAULT_BUILDER_DERIVE_OPTIONS, true), ("HnswConfigDiff", DEFAULT_BUILDER_DERIVE_OPTIONS, true), + ( + "ScalarQuantization", + NO_DEFAULT_BUILDER_DERIVE_OPTIONS, + true, + ), + ( + "ProductQuantization", + NO_DEFAULT_BUILDER_DERIVE_OPTIONS, + true, + ), + ( + "BinaryQuantization", + NO_DEFAULT_BUILDER_DERIVE_OPTIONS, + true, + ), + ("OptimizersConfigDiff", DEFAULT_BUILDER_DERIVE_OPTIONS, true), + ("WalConfigDiff", DEFAULT_BUILDER_DERIVE_OPTIONS, true), ] } diff --git a/tests/protos_append/builder_ext.rs b/tests/protos_append/builder_ext.rs new file mode 100644 index 0000000..5827b9f --- /dev/null +++ b/tests/protos_append/builder_ext.rs @@ -0,0 +1,32 @@ +impl VectorParamsBuilder { + pub fn new(size: u64, distance: Distance) -> Self { + let mut builder = Self::create_empty(); + builder.size = Some(size); + builder.distance = Some(distance.into()); + builder + } +} + +impl ScalarQuantizationBuilder { + pub fn new(r#type: QuantizationType) -> Self { + let mut builder = Self::create_empty(); + builder.r#type = Some(r#type.into()); + builder + } +} + +impl ProductQuantizationBuilder { + pub fn new(compression: i32) -> Self { + let mut builder = Self::create_empty(); + builder.compression = Some(compression); + builder + } +} + +impl BinaryQuantizationBuilder { + pub fn new(always_ram: bool) -> Self { + let mut builder = Self::create_empty(); + builder.always_ram = Some(Some(always_ram)); + builder + } +}