From e41a480c1457fbd77dfd6ffb39e5e662efccf90f Mon Sep 17 00:00:00 2001 From: Roman Titov Date: Fri, 29 Aug 2025 10:15:58 +0200 Subject: [PATCH] Move `CoreSearchRequest` and `QueryEnum` from collection into `shard` crate (#7169) --- lib/collection/src/operations/conversions.rs | 205 +-------- lib/collection/src/operations/mod.rs | 5 +- lib/collection/src/operations/query_enum.rs | 119 ------ lib/collection/src/operations/types.rs | 61 +-- .../operations/universal_query/shard_query.rs | 132 +++--- .../verification/operation_rate_cost.rs | 12 +- lib/shard/src/lib.rs | 1 + lib/shard/src/search.rs | 390 ++++++++++++++++++ 8 files changed, 464 insertions(+), 461 deletions(-) delete mode 100644 lib/collection/src/operations/query_enum.rs create mode 100644 lib/shard/src/search.rs diff --git a/lib/collection/src/operations/conversions.rs b/lib/collection/src/operations/conversions.rs index 98747108cd..090375eb02 100644 --- a/lib/collection/src/operations/conversions.rs +++ b/lib/collection/src/operations/conversions.rs @@ -13,15 +13,13 @@ use api::grpc::qdrant::update_collection_cluster_setup_request::{ }; use api::rest::schema::ShardKeySelector; use api::rest::{BaseGroupRequest, MaxOptimizationThreads}; -use common::types::ScoreType; use itertools::Itertools; use segment::common::operation_error::OperationError; -use segment::data_types::vectors::{NamedQuery, VectorInternal, VectorStructInternal}; +use segment::data_types::vectors::{VectorInternal, VectorStructInternal}; use segment::types::{ - Distance, Filter, HnswConfig, MultiVectorConfig, QuantizationConfig, SearchParams, - StrictModeConfigOutput, WithPayloadInterface, WithVector, + Distance, HnswConfig, MultiVectorConfig, QuantizationConfig, StrictModeConfigOutput, + WithPayloadInterface, }; -use segment::vector_storage::query::{ContextPair, ContextQuery, DiscoveryQuery, RecoQuery}; use sparse::common::sparse_vector::{SparseVector, validate_sparse_vector_impl}; use tonic::Status; @@ -51,7 +49,6 @@ use crate::operations::config_diff::{ WalConfigDiff, }; use crate::operations::point_ops::{FilterSelector, PointIdsList, PointsSelector, WriteOrdering}; -use crate::operations::query_enum::QueryEnum; use crate::operations::shard_selector_internal::ShardSelectorInternal; use crate::operations::types::{ AliasDescription, CollectionClusterInfo, CollectionInfo, CollectionStatus, CountResult, @@ -1011,106 +1008,6 @@ impl From for api::grpc::qdrant::CountResult { } } -impl TryFrom for CoreSearchRequest { - type Error = Status; - fn try_from(value: api::grpc::qdrant::SearchPoints) -> Result { - let api::grpc::qdrant::SearchPoints { - collection_name: _, - vector, - filter, - limit, - with_payload, - params, - score_threshold, - offset, - vector_name, - with_vectors, - read_consistency: _, - timeout: _, - shard_key_selector: _, - sparse_indices, - } = value; - - if let Some(sparse_indices) = &sparse_indices { - let api::grpc::qdrant::SparseIndices { data } = sparse_indices; - validate_sparse_vector_impl(data, &vector).map_err(|e| { - Status::invalid_argument(format!( - "Sparse indices does not match sparse vector conditions: {e}" - )) - })?; - } - - let vector_struct = - api::grpc::conversions::into_named_vector_struct(vector_name, vector, sparse_indices)?; - - Ok(Self { - query: QueryEnum::Nearest(NamedQuery::from(vector_struct)), - filter: filter.map(Filter::try_from).transpose()?, - params: params.map(SearchParams::from), - limit: limit as usize, - offset: offset.map(|v| v as usize).unwrap_or_default(), - with_payload: with_payload - .map(WithPayloadInterface::try_from) - .transpose()?, - with_vector: with_vectors.map(WithVector::from), - score_threshold: score_threshold.map(|s| s as ScoreType), - }) - } -} - -impl From for api::grpc::qdrant::QueryEnum { - fn from(value: QueryEnum) -> Self { - match value { - QueryEnum::Nearest(vector) => api::grpc::qdrant::QueryEnum { - query: Some(api::grpc::qdrant::query_enum::Query::NearestNeighbors( - api::grpc::qdrant::Vector::from(vector.query), - )), - }, - QueryEnum::RecommendBestScore(named) => api::grpc::qdrant::QueryEnum { - query: Some(api::grpc::qdrant::query_enum::Query::RecommendBestScore( - named.query.into(), - )), - }, - QueryEnum::RecommendSumScores(named) => api::grpc::qdrant::QueryEnum { - query: Some(api::grpc::qdrant::query_enum::Query::RecommendSumScores( - named.query.into(), - )), - }, - QueryEnum::Discover(named) => api::grpc::qdrant::QueryEnum { - query: Some(api::grpc::qdrant::query_enum::Query::Discover( - api::grpc::qdrant::DiscoveryQuery { - target: Some(named.query.target.into()), - context: named - .query - .pairs - .into_iter() - .map(|pair| api::grpc::qdrant::ContextPair { - positive: { Some(pair.positive.into()) }, - negative: { Some(pair.negative.into()) }, - }) - .collect(), - }, - )), - }, - QueryEnum::Context(named) => api::grpc::qdrant::QueryEnum { - query: Some(api::grpc::qdrant::query_enum::Query::Context( - api::grpc::qdrant::ContextQuery { - context: named - .query - .pairs - .into_iter() - .map(|pair| api::grpc::qdrant::ContextPair { - positive: { Some(pair.positive.into()) }, - negative: { Some(pair.negative.into()) }, - }) - .collect(), - }, - )), - }, - } - } -} - impl<'a> From> for api::grpc::qdrant::CoreSearchPoints { fn from(value: CollectionCoreSearchRequest<'a>) -> Self { let (collection_id, request) = value.0; @@ -1184,102 +1081,6 @@ impl TryFrom for RecommendExample { } } -fn try_context_pair_from_grpc( - pair: api::grpc::qdrant::ContextPair, -) -> Result, Status> { - let api::grpc::qdrant::ContextPair { positive, negative } = pair; - match (positive, negative) { - (Some(positive), Some(negative)) => Ok(ContextPair { - positive: positive.try_into()?, - negative: negative.try_into()?, - }), - _ => Err(Status::invalid_argument( - "All context pairs must have both positive and negative parts", - )), - } -} - -impl TryFrom for CoreSearchRequest { - type Error = Status; - - fn try_from(value: api::grpc::qdrant::CoreSearchPoints) -> Result { - let query = value - .query - .and_then(|query| query.query) - .map(|query| { - Ok(match query { - api::grpc::qdrant::query_enum::Query::NearestNeighbors(vector) => { - QueryEnum::Nearest(NamedQuery::from( - api::grpc::conversions::into_named_vector_struct( - value.vector_name, - vector.data, - vector.indices, - )?, - )) - } - api::grpc::qdrant::query_enum::Query::RecommendBestScore(query) => { - QueryEnum::RecommendBestScore(NamedQuery { - query: RecoQuery::try_from(query)?, - using: value.vector_name, - }) - } - api::grpc::qdrant::query_enum::Query::RecommendSumScores(query) => { - QueryEnum::RecommendSumScores(NamedQuery { - query: RecoQuery::try_from(query)?, - using: value.vector_name, - }) - } - api::grpc::qdrant::query_enum::Query::Discover(query) => { - let Some(target) = query.target else { - return Err(Status::invalid_argument("Target is not specified")); - }; - - let pairs = query - .context - .into_iter() - .map(try_context_pair_from_grpc) - .try_collect()?; - - QueryEnum::Discover(NamedQuery { - query: DiscoveryQuery::new(target.try_into()?, pairs), - using: value.vector_name, - }) - } - api::grpc::qdrant::query_enum::Query::Context(query) => { - let pairs = query - .context - .into_iter() - .map(try_context_pair_from_grpc) - .try_collect()?; - - QueryEnum::Context(NamedQuery { - query: ContextQuery::new(pairs), - using: value.vector_name, - }) - } - }) - }) - .transpose()? - .ok_or_else(|| Status::invalid_argument("Query is not specified"))?; - - Ok(Self { - query, - filter: value.filter.map(|f| f.try_into()).transpose()?, - params: value.params.map(|p| p.into()), - limit: value.limit as usize, - offset: value.offset.unwrap_or_default() as usize, - with_payload: value.with_payload.map(|wp| wp.try_into()).transpose()?, - with_vector: Some( - value - .with_vectors - .map(|with_vectors| with_vectors.into()) - .unwrap_or_default(), - ), - score_threshold: value.score_threshold, - }) - } -} - impl TryFrom for api::grpc::qdrant::PointGroup { type Error = OperationError; fn try_from(group: PointGroup) -> Result { diff --git a/lib/collection/src/operations/mod.rs b/lib/collection/src/operations/mod.rs index 41ccf8af83..a377677e3d 100644 --- a/lib/collection/src/operations/mod.rs +++ b/lib/collection/src/operations/mod.rs @@ -6,7 +6,6 @@ pub mod conversions_rest; pub mod operation_effect; pub mod payload_ops; pub mod point_ops; -pub mod query_enum; pub mod shard_selector_internal; pub mod shared_storage_config; pub mod snapshot_ops; @@ -18,6 +17,10 @@ pub mod vector_ops; pub mod vector_params_builder; pub mod verification; +pub mod query_enum { + pub use shard::search::QueryEnum; +} + use std::collections::HashMap; use segment::types::ExtendedPointId; diff --git a/lib/collection/src/operations/query_enum.rs b/lib/collection/src/operations/query_enum.rs deleted file mode 100644 index bffc267a39..0000000000 --- a/lib/collection/src/operations/query_enum.rs +++ /dev/null @@ -1,119 +0,0 @@ -use std::fmt::Debug; -use std::iter; - -use segment::data_types::vectors::{DenseVector, Named, NamedQuery, VectorInternal}; -use segment::types::VectorName; -use segment::vector_storage::query::{ContextQuery, DiscoveryQuery, RecoQuery}; -use sparse::common::sparse_vector::SparseVector; - -impl QueryEnum { - pub fn get_vector_name(&self) -> &VectorName { - match self { - QueryEnum::Nearest(vector) => vector.get_name(), - QueryEnum::RecommendBestScore(reco_query) => reco_query.get_name(), - QueryEnum::RecommendSumScores(reco_query) => reco_query.get_name(), - QueryEnum::Discover(discovery_query) => discovery_query.get_name(), - QueryEnum::Context(context_query) => context_query.get_name(), - } - } - - /// Only when the distance is the scoring, this will return true. - pub fn is_distance_scored(&self) -> bool { - match self { - QueryEnum::Nearest(_) => true, - QueryEnum::RecommendBestScore(_) - | QueryEnum::RecommendSumScores(_) - | QueryEnum::Discover(_) - | QueryEnum::Context(_) => false, - } - } - - pub fn iterate_sparse(&self, mut f: impl FnMut(&VectorName, &SparseVector)) { - match self { - QueryEnum::Nearest(named) => match &named.query { - VectorInternal::Sparse(sparse_vector) => f(named.get_name(), sparse_vector), - VectorInternal::Dense(_) | VectorInternal::MultiDense(_) => {} - }, - QueryEnum::RecommendBestScore(reco_query) - | QueryEnum::RecommendSumScores(reco_query) => { - let name = reco_query.get_name(); - for vector in reco_query.query.flat_iter() { - match vector { - VectorInternal::Sparse(sparse_vector) => f(name, sparse_vector), - VectorInternal::Dense(_) | VectorInternal::MultiDense(_) => {} - } - } - } - QueryEnum::Discover(discovery_query) => { - let name = discovery_query.get_name(); - for pair in discovery_query.query.flat_iter() { - match pair { - VectorInternal::Sparse(sparse_vector) => f(name, sparse_vector), - VectorInternal::Dense(_) | VectorInternal::MultiDense(_) => {} - } - } - } - QueryEnum::Context(context_query) => { - let name = context_query.get_name(); - for pair in context_query.query.flat_iter() { - match pair { - VectorInternal::Sparse(sparse_vector) => f(name, sparse_vector), - VectorInternal::Dense(_) | VectorInternal::MultiDense(_) => {} - } - } - } - } - } -} - -/// Every kind of vector query that can be performed on segment level. -#[derive(Debug, Clone, PartialEq)] -pub enum QueryEnum { - Nearest(NamedQuery), - RecommendBestScore(NamedQuery>), - RecommendSumScores(NamedQuery>), - Discover(NamedQuery>), - Context(NamedQuery>), -} - -impl QueryEnum { - /// Iterate over all vectors in the query. - fn vectors(&self) -> Box + '_> { - match self { - QueryEnum::Nearest(named_query) => Box::new(iter::once(&named_query.query)), - QueryEnum::RecommendBestScore(named_query) => Box::new(named_query.query.flat_iter()), - QueryEnum::RecommendSumScores(named_query) => Box::new(named_query.query.flat_iter()), - QueryEnum::Discover(named_query) => Box::new(named_query.query.flat_iter()), - QueryEnum::Context(named_query) => Box::new(named_query.query.flat_iter()), - } - } - - /// Returns the estimated cost of using this query in terms of number of vectors. - /// The cost approximates how many similarity comparisons this query will make against one point. - pub fn search_cost(&self) -> usize { - self.vectors() - .map(|vector_internal| vector_internal.similarity_cost()) - .sum() - } -} - -impl From for QueryEnum { - fn from(vector: DenseVector) -> Self { - QueryEnum::Nearest(NamedQuery { - query: VectorInternal::Dense(vector), - using: None, - }) - } -} - -impl From>> for QueryEnum { - fn from(query: NamedQuery>) -> Self { - QueryEnum::Discover(query) - } -} - -impl AsRef for QueryEnum { - fn as_ref(&self) -> &QueryEnum { - self - } -} diff --git a/lib/collection/src/operations/types.rs b/lib/collection/src/operations/types.rs index facf015fe5..7856605d36 100644 --- a/lib/collection/src/operations/types.rs +++ b/lib/collection/src/operations/types.rs @@ -25,7 +25,7 @@ use segment::common::operation_error::{CancelledError, OperationError}; use segment::data_types::groups::GroupId; use segment::data_types::order_by::{OrderBy, OrderValue}; use segment::data_types::vectors::{ - DEFAULT_VECTOR_NAME, DenseVector, NamedQuery, NamedVectorStruct, QueryVector, VectorRef, + DEFAULT_VECTOR_NAME, DenseVector, NamedQuery, NamedVectorStruct, VectorRef, VectorStructInternal, }; use segment::types::{ @@ -38,6 +38,7 @@ use semver::Version; use serde; use serde::{Deserialize, Serialize}; use serde_json::{Error as JsonError, Map, Value}; +pub use shard::search::CoreSearchRequest; use shard::wal::WalError; use sparse::common::sparse_vector::SparseVector; use thiserror::Error; @@ -582,27 +583,6 @@ pub struct SearchRequestBatch { pub searches: Vec, } -#[derive(Debug, Clone, PartialEq)] -pub struct CoreSearchRequest { - /// Every kind of query that can be performed on segment level - pub query: QueryEnum, - /// Look only for points which satisfies this conditions - pub filter: Option, - /// Additional search params - pub params: Option, - /// Max number of result to return - pub limit: usize, - /// Offset of the first result to return. - /// May be used to paginate results. - /// Note: large offset values may cause performance issues. - pub offset: usize, - /// Select which payload to return with the response. Default is false. - pub with_payload: Option, - /// Options for specifying which vectors to include into response. Default is false. - pub with_vector: Option, - pub score_threshold: Option, -} - #[derive(Debug, Clone)] pub struct CoreSearchRequestBatch { pub searches: Vec, @@ -2013,31 +1993,6 @@ pub enum NodeType { Listener, } -impl From for CoreSearchRequest { - fn from(request: SearchRequestInternal) -> Self { - let SearchRequestInternal { - vector, - filter, - score_threshold, - limit, - offset, - params, - with_vector, - with_payload, - } = request; - Self { - query: QueryEnum::Nearest(NamedQuery::from(NamedVectorStruct::from(vector))), - filter, - params, - limit, - offset: offset.unwrap_or_default(), - with_payload, - with_vector, - score_threshold, - } - } -} - impl From for ShardQueryRequest { fn from(value: SearchRequestInternal) -> Self { let SearchRequestInternal { @@ -2094,18 +2049,6 @@ impl From for ShardQueryRequest { } } -impl From for QueryVector { - fn from(query: QueryEnum) -> Self { - match query { - QueryEnum::Nearest(named) => QueryVector::Nearest(named.query), - QueryEnum::RecommendBestScore(named) => QueryVector::RecommendBestScore(named.query), - QueryEnum::RecommendSumScores(named) => QueryVector::RecommendSumScores(named.query), - QueryEnum::Discover(named) => QueryVector::Discovery(named.query), - QueryEnum::Context(named) => QueryVector::Context(named.query), - } - } -} - /// All the unresolved issues in a Qdrant instance #[derive(Serialize, JsonSchema, Debug)] pub struct IssuesReport { diff --git a/lib/collection/src/operations/universal_query/shard_query.rs b/lib/collection/src/operations/universal_query/shard_query.rs index 7a99740577..1f3cd1864c 100644 --- a/lib/collection/src/operations/universal_query/shard_query.rs +++ b/lib/collection/src/operations/universal_query/shard_query.rs @@ -297,55 +297,53 @@ impl TryFrom for ShardPrefetch { } } -impl QueryEnum { - fn try_from_grpc_raw_query( - raw_query: grpc::RawQuery, - using: Option, - ) -> Result { - use grpc::raw_query::Variant; +fn query_enum_from_grpc_raw_query( + raw_query: grpc::RawQuery, + using: Option, +) -> Result { + use grpc::raw_query::Variant; - let variant = raw_query - .variant - .ok_or_else(|| Status::invalid_argument("missing field: variant"))?; + let variant = raw_query + .variant + .ok_or_else(|| Status::invalid_argument("missing field: variant"))?; - let query_enum = match variant { - Variant::Nearest(nearest) => { - let vector = VectorInternal::try_from(nearest)?; - let name = match (using, &vector) { - (None, VectorInternal::Sparse(_)) => { - return Err(Status::invalid_argument("Sparse vector must have a name")); - } - ( - Some(name), - VectorInternal::MultiDense(_) - | VectorInternal::Sparse(_) - | VectorInternal::Dense(_), - ) => name, - (None, VectorInternal::MultiDense(_) | VectorInternal::Dense(_)) => { - DEFAULT_VECTOR_NAME.to_owned() - } - }; - let named_vector = NamedQuery::new_from_vector(vector, name); - QueryEnum::Nearest(named_vector) - } - Variant::RecommendBestScore(recommend) => QueryEnum::RecommendBestScore( - NamedQuery::new(RecoQuery::try_from(recommend)?, using), - ), - Variant::RecommendSumScores(recommend) => QueryEnum::RecommendSumScores( - NamedQuery::new(RecoQuery::try_from(recommend)?, using), - ), - Variant::Discover(discovery) => QueryEnum::Discover(NamedQuery { - query: DiscoveryQuery::try_from(discovery)?, - using, - }), - Variant::Context(context) => QueryEnum::Context(NamedQuery { - query: ContextQuery::try_from(context)?, - using, - }), - }; + let query_enum = match variant { + Variant::Nearest(nearest) => { + let vector = VectorInternal::try_from(nearest)?; + let name = match (using, &vector) { + (None, VectorInternal::Sparse(_)) => { + return Err(Status::invalid_argument("Sparse vector must have a name")); + } + ( + Some(name), + VectorInternal::MultiDense(_) + | VectorInternal::Sparse(_) + | VectorInternal::Dense(_), + ) => name, + (None, VectorInternal::MultiDense(_) | VectorInternal::Dense(_)) => { + DEFAULT_VECTOR_NAME.to_owned() + } + }; + let named_vector = NamedQuery::new_from_vector(vector, name); + QueryEnum::Nearest(named_vector) + } + Variant::RecommendBestScore(recommend) => { + QueryEnum::RecommendBestScore(NamedQuery::new(RecoQuery::try_from(recommend)?, using)) + } + Variant::RecommendSumScores(recommend) => { + QueryEnum::RecommendSumScores(NamedQuery::new(RecoQuery::try_from(recommend)?, using)) + } + Variant::Discover(discovery) => QueryEnum::Discover(NamedQuery { + query: DiscoveryQuery::try_from(discovery)?, + using, + }), + Variant::Context(context) => QueryEnum::Context(NamedQuery { + query: ContextQuery::try_from(context)?, + using, + }), + }; - Ok(query_enum) - } + Ok(query_enum) } impl TryFrom for FusionInternal { @@ -620,7 +618,7 @@ impl ScoringQuery { .ok_or_else(|| Status::invalid_argument("missing field: score"))?; let scoring_query = match score { grpc::query_shard_points::query::Score::Vector(query) => { - ScoringQuery::Vector(QueryEnum::try_from_grpc_raw_query(query, using)?) + ScoringQuery::Vector(query_enum_from_grpc_raw_query(query, using)?) } grpc::query_shard_points::query::Score::Fusion(fusion) => { ScoringQuery::Fusion(FusionInternal::try_from(fusion)?) @@ -660,29 +658,25 @@ impl ScoringQuery { } } -impl From for grpc::RawQuery { - fn from(value: QueryEnum) -> Self { - use api::grpc::qdrant::raw_query::Variant; +fn query_enum_into_grpc_raw_query(query: QueryEnum) -> grpc::RawQuery { + use api::grpc::qdrant::raw_query::Variant; - let variant = match value { - QueryEnum::Nearest(named) => Variant::Nearest(grpc::RawVector::from(named.query)), - QueryEnum::RecommendBestScore(named) => { - Variant::RecommendBestScore(grpc::raw_query::Recommend::from(named.query)) - } - QueryEnum::RecommendSumScores(named) => { - Variant::RecommendSumScores(grpc::raw_query::Recommend::from(named.query)) - } - QueryEnum::Discover(named) => { - Variant::Discover(grpc::raw_query::Discovery::from(named.query)) - } - QueryEnum::Context(named) => { - Variant::Context(grpc::raw_query::Context::from(named.query)) - } - }; - - Self { - variant: Some(variant), + let variant = match query { + QueryEnum::Nearest(named) => Variant::Nearest(grpc::RawVector::from(named.query)), + QueryEnum::RecommendBestScore(named) => { + Variant::RecommendBestScore(grpc::raw_query::Recommend::from(named.query)) } + QueryEnum::RecommendSumScores(named) => { + Variant::RecommendSumScores(grpc::raw_query::Recommend::from(named.query)) + } + QueryEnum::Discover(named) => { + Variant::Discover(grpc::raw_query::Discovery::from(named.query)) + } + QueryEnum::Context(named) => Variant::Context(grpc::raw_query::Context::from(named.query)), + }; + + grpc::RawQuery { + variant: Some(variant), } } @@ -692,7 +686,7 @@ impl From for grpc::query_shard_points::Query { match value { ScoringQuery::Vector(query) => Self { - score: Some(Score::Vector(grpc::RawQuery::from(query))), + score: Some(Score::Vector(query_enum_into_grpc_raw_query(query))), }, ScoringQuery::Fusion(fusion) => Self::from(fusion), ScoringQuery::OrderBy(order_by) => Self { diff --git a/lib/collection/src/operations/verification/operation_rate_cost.rs b/lib/collection/src/operations/verification/operation_rate_cost.rs index 263d069354..1e6c3a0005 100644 --- a/lib/collection/src/operations/verification/operation_rate_cost.rs +++ b/lib/collection/src/operations/verification/operation_rate_cost.rs @@ -1,6 +1,6 @@ use segment::types::Filter; -use crate::operations::types::{CoreSearchRequest, QueryScrollRequestInternal}; +use crate::operations::types::QueryScrollRequestInternal; pub fn filter_rate_cost(filter: &Filter) -> usize { filter.total_conditions_count() @@ -9,16 +9,6 @@ pub fn filter_rate_cost(filter: &Filter) -> usize { /// Base cost for a read operation pub const BASE_COST: usize = 1; -impl CoreSearchRequest { - pub fn search_rate_cost(&self) -> usize { - let mut cost = self.query.search_cost(); - if let Some(filter) = &self.filter { - cost += filter_rate_cost(filter); - } - cost - } -} - impl QueryScrollRequestInternal { pub fn scroll_rate_cost(&self) -> usize { let mut cost = BASE_COST; diff --git a/lib/shard/src/lib.rs b/lib/shard/src/lib.rs index 955e5c3d98..b4d38710da 100644 --- a/lib/shard/src/lib.rs +++ b/lib/shard/src/lib.rs @@ -2,6 +2,7 @@ pub mod locked_segment; pub mod operations; pub mod payload_index_schema; pub mod proxy_segment; +pub mod search; pub mod segment_holder; pub mod update; pub mod wal; diff --git a/lib/shard/src/search.rs b/lib/shard/src/search.rs new file mode 100644 index 0000000000..14995b02ad --- /dev/null +++ b/lib/shard/src/search.rs @@ -0,0 +1,390 @@ +use api::rest::SearchRequestInternal; +use common::types::ScoreType; +use itertools::Itertools as _; +use segment::data_types::vectors::{ + DenseVector, Named as _, NamedQuery, NamedVectorStruct, QueryVector, VectorInternal, +}; +use segment::types::{Filter, SearchParams, VectorName, WithPayloadInterface, WithVector}; +use segment::vector_storage::query::{ContextPair, ContextQuery, DiscoveryQuery, RecoQuery}; +use sparse::common::sparse_vector::{SparseVector, validate_sparse_vector_impl}; + +#[derive(Clone, Debug, PartialEq)] +pub struct CoreSearchRequest { + /// Every kind of query that can be performed on segment level + pub query: QueryEnum, + /// Look only for points which satisfies this conditions + pub filter: Option, + /// Additional search params + pub params: Option, + /// Max number of result to return + pub limit: usize, + /// Offset of the first result to return. + /// May be used to paginate results. + /// Note: large offset values may cause performance issues. + pub offset: usize, + /// Select which payload to return with the response. Default is false. + pub with_payload: Option, + /// Options for specifying which vectors to include into response. Default is false. + pub with_vector: Option, + pub score_threshold: Option, +} + +impl CoreSearchRequest { + pub fn search_rate_cost(&self) -> usize { + let mut cost = self.query.search_cost(); + + if let Some(filter) = &self.filter { + cost += filter.total_conditions_count(); + } + + cost + } +} + +impl From for CoreSearchRequest { + fn from(request: SearchRequestInternal) -> Self { + let SearchRequestInternal { + vector, + filter, + score_threshold, + limit, + offset, + params, + with_vector, + with_payload, + } = request; + Self { + query: QueryEnum::Nearest(NamedQuery::from(NamedVectorStruct::from(vector))), + filter, + params, + limit, + offset: offset.unwrap_or_default(), + with_payload, + with_vector, + score_threshold, + } + } +} + +impl TryFrom for CoreSearchRequest { + type Error = tonic::Status; + + fn try_from(value: api::grpc::qdrant::CoreSearchPoints) -> Result { + let query = value + .query + .and_then(|query| query.query) + .map(|query| { + Ok(match query { + api::grpc::qdrant::query_enum::Query::NearestNeighbors(vector) => { + QueryEnum::Nearest(NamedQuery::from( + api::grpc::conversions::into_named_vector_struct( + value.vector_name, + vector.data, + vector.indices, + )?, + )) + } + api::grpc::qdrant::query_enum::Query::RecommendBestScore(query) => { + QueryEnum::RecommendBestScore(NamedQuery { + query: RecoQuery::try_from(query)?, + using: value.vector_name, + }) + } + api::grpc::qdrant::query_enum::Query::RecommendSumScores(query) => { + QueryEnum::RecommendSumScores(NamedQuery { + query: RecoQuery::try_from(query)?, + using: value.vector_name, + }) + } + api::grpc::qdrant::query_enum::Query::Discover(query) => { + let Some(target) = query.target else { + return Err(tonic::Status::invalid_argument("Target is not specified")); + }; + + let pairs = query + .context + .into_iter() + .map(try_context_pair_from_grpc) + .try_collect()?; + + QueryEnum::Discover(NamedQuery { + query: DiscoveryQuery::new(target.try_into()?, pairs), + using: value.vector_name, + }) + } + api::grpc::qdrant::query_enum::Query::Context(query) => { + let pairs = query + .context + .into_iter() + .map(try_context_pair_from_grpc) + .try_collect()?; + + QueryEnum::Context(NamedQuery { + query: ContextQuery::new(pairs), + using: value.vector_name, + }) + } + }) + }) + .transpose()? + .ok_or_else(|| tonic::Status::invalid_argument("Query is not specified"))?; + + Ok(Self { + query, + filter: value.filter.map(|f| f.try_into()).transpose()?, + params: value.params.map(|p| p.into()), + limit: value.limit as usize, + offset: value.offset.unwrap_or_default() as usize, + with_payload: value.with_payload.map(|wp| wp.try_into()).transpose()?, + with_vector: Some( + value + .with_vectors + .map(|with_vectors| with_vectors.into()) + .unwrap_or_default(), + ), + score_threshold: value.score_threshold, + }) + } +} + +fn try_context_pair_from_grpc( + pair: api::grpc::qdrant::ContextPair, +) -> Result, tonic::Status> { + let api::grpc::qdrant::ContextPair { positive, negative } = pair; + match (positive, negative) { + (Some(positive), Some(negative)) => Ok(ContextPair { + positive: positive.try_into()?, + negative: negative.try_into()?, + }), + _ => Err(tonic::Status::invalid_argument( + "All context pairs must have both positive and negative parts", + )), + } +} + +impl TryFrom for CoreSearchRequest { + type Error = tonic::Status; + + fn try_from(value: api::grpc::qdrant::SearchPoints) -> Result { + let api::grpc::qdrant::SearchPoints { + collection_name: _, + vector, + filter, + limit, + with_payload, + params, + score_threshold, + offset, + vector_name, + with_vectors, + read_consistency: _, + timeout: _, + shard_key_selector: _, + sparse_indices, + } = value; + + if let Some(sparse_indices) = &sparse_indices { + let api::grpc::qdrant::SparseIndices { data } = sparse_indices; + validate_sparse_vector_impl(data, &vector).map_err(|e| { + tonic::Status::invalid_argument(format!( + "Sparse indices does not match sparse vector conditions: {e}" + )) + })?; + } + + let vector_struct = + api::grpc::conversions::into_named_vector_struct(vector_name, vector, sparse_indices)?; + + Ok(Self { + query: QueryEnum::Nearest(NamedQuery::from(vector_struct)), + filter: filter.map(Filter::try_from).transpose()?, + params: params.map(SearchParams::from), + limit: limit as usize, + offset: offset.map(|v| v as usize).unwrap_or_default(), + with_payload: with_payload + .map(WithPayloadInterface::try_from) + .transpose()?, + with_vector: with_vectors.map(WithVector::from), + score_threshold: score_threshold.map(|s| s as ScoreType), + }) + } +} + +/// Every kind of vector query that can be performed on segment level. +#[derive(Clone, Debug, PartialEq)] +pub enum QueryEnum { + Nearest(NamedQuery), + RecommendBestScore(NamedQuery>), + RecommendSumScores(NamedQuery>), + Discover(NamedQuery>), + Context(NamedQuery>), +} + +impl QueryEnum { + pub fn get_vector_name(&self) -> &VectorName { + match self { + QueryEnum::Nearest(vector) => vector.get_name(), + QueryEnum::RecommendBestScore(reco_query) => reco_query.get_name(), + QueryEnum::RecommendSumScores(reco_query) => reco_query.get_name(), + QueryEnum::Discover(discovery_query) => discovery_query.get_name(), + QueryEnum::Context(context_query) => context_query.get_name(), + } + } + + /// Only when the distance is the scoring, this will return true. + pub fn is_distance_scored(&self) -> bool { + match self { + QueryEnum::Nearest(_) => true, + QueryEnum::RecommendBestScore(_) + | QueryEnum::RecommendSumScores(_) + | QueryEnum::Discover(_) + | QueryEnum::Context(_) => false, + } + } + + pub fn iterate_sparse(&self, mut f: impl FnMut(&VectorName, &SparseVector)) { + match self { + QueryEnum::Nearest(named) => match &named.query { + VectorInternal::Sparse(sparse_vector) => f(named.get_name(), sparse_vector), + VectorInternal::Dense(_) | VectorInternal::MultiDense(_) => {} + }, + QueryEnum::RecommendBestScore(reco_query) + | QueryEnum::RecommendSumScores(reco_query) => { + let name = reco_query.get_name(); + for vector in reco_query.query.flat_iter() { + match vector { + VectorInternal::Sparse(sparse_vector) => f(name, sparse_vector), + VectorInternal::Dense(_) | VectorInternal::MultiDense(_) => {} + } + } + } + QueryEnum::Discover(discovery_query) => { + let name = discovery_query.get_name(); + for pair in discovery_query.query.flat_iter() { + match pair { + VectorInternal::Sparse(sparse_vector) => f(name, sparse_vector), + VectorInternal::Dense(_) | VectorInternal::MultiDense(_) => {} + } + } + } + QueryEnum::Context(context_query) => { + let name = context_query.get_name(); + for pair in context_query.query.flat_iter() { + match pair { + VectorInternal::Sparse(sparse_vector) => f(name, sparse_vector), + VectorInternal::Dense(_) | VectorInternal::MultiDense(_) => {} + } + } + } + } + } + + /// Returns the estimated cost of using this query in terms of number of vectors. + /// The cost approximates how many similarity comparisons this query will make against one point. + fn search_cost(&self) -> usize { + match self { + QueryEnum::Nearest(named_query) => search_cost([&named_query.query]), + QueryEnum::RecommendBestScore(named_query) => { + search_cost(named_query.query.flat_iter()) + } + QueryEnum::RecommendSumScores(named_query) => { + search_cost(named_query.query.flat_iter()) + } + QueryEnum::Discover(named_query) => search_cost(named_query.query.flat_iter()), + QueryEnum::Context(named_query) => search_cost(named_query.query.flat_iter()), + } + } +} + +fn search_cost<'a>(vectors: impl IntoIterator) -> usize { + vectors + .into_iter() + .map(VectorInternal::similarity_cost) + .sum() +} + +impl AsRef for QueryEnum { + fn as_ref(&self) -> &QueryEnum { + self + } +} + +impl From for QueryEnum { + fn from(vector: DenseVector) -> Self { + QueryEnum::Nearest(NamedQuery { + query: VectorInternal::Dense(vector), + using: None, + }) + } +} + +impl From>> for QueryEnum { + fn from(query: NamedQuery>) -> Self { + QueryEnum::Discover(query) + } +} + +impl From for QueryVector { + fn from(query: QueryEnum) -> Self { + match query { + QueryEnum::Nearest(named) => QueryVector::Nearest(named.query), + QueryEnum::RecommendBestScore(named) => QueryVector::RecommendBestScore(named.query), + QueryEnum::RecommendSumScores(named) => QueryVector::RecommendSumScores(named.query), + QueryEnum::Discover(named) => QueryVector::Discovery(named.query), + QueryEnum::Context(named) => QueryVector::Context(named.query), + } + } +} + +impl From for api::grpc::qdrant::QueryEnum { + fn from(value: QueryEnum) -> Self { + match value { + QueryEnum::Nearest(vector) => api::grpc::qdrant::QueryEnum { + query: Some(api::grpc::qdrant::query_enum::Query::NearestNeighbors( + api::grpc::qdrant::Vector::from(vector.query), + )), + }, + QueryEnum::RecommendBestScore(named) => api::grpc::qdrant::QueryEnum { + query: Some(api::grpc::qdrant::query_enum::Query::RecommendBestScore( + named.query.into(), + )), + }, + QueryEnum::RecommendSumScores(named) => api::grpc::qdrant::QueryEnum { + query: Some(api::grpc::qdrant::query_enum::Query::RecommendSumScores( + named.query.into(), + )), + }, + QueryEnum::Discover(named) => api::grpc::qdrant::QueryEnum { + query: Some(api::grpc::qdrant::query_enum::Query::Discover( + api::grpc::qdrant::DiscoveryQuery { + target: Some(named.query.target.into()), + context: named + .query + .pairs + .into_iter() + .map(|pair| api::grpc::qdrant::ContextPair { + positive: { Some(pair.positive.into()) }, + negative: { Some(pair.negative.into()) }, + }) + .collect(), + }, + )), + }, + QueryEnum::Context(named) => api::grpc::qdrant::QueryEnum { + query: Some(api::grpc::qdrant::query_enum::Query::Context( + api::grpc::qdrant::ContextQuery { + context: named + .query + .pairs + .into_iter() + .map(|pair| api::grpc::qdrant::ContextPair { + positive: { Some(pair.positive.into()) }, + negative: { Some(pair.negative.into()) }, + }) + .collect(), + }, + )), + }, + } + } +}