From d7186e80949ec0d8d2e453d32198358face324cb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20Coss=C3=ADo?= Date: Wed, 15 May 2024 11:53:33 -0400 Subject: [PATCH] universal-query: Connect grpc internal query service (#4223) * add grpc -> internal conversions * return `Status` instead of `CollectionError` in grpc -> internal conversions * sketch points internal api connection of `query` --- lib/api/src/grpc/conversions.rs | 127 ++++++++++++ lib/collection/src/operations/types.rs | 6 +- .../src/operations/universal_query.rs | 195 +++++++++++++++--- lib/segment/src/data_types/vectors.rs | 1 + src/tonic/api/points_internal_api.rs | 78 ++++++- 5 files changed, 369 insertions(+), 38 deletions(-) diff --git a/lib/api/src/grpc/conversions.rs b/lib/api/src/grpc/conversions.rs index 1678400070..3c23b9a603 100644 --- a/lib/api/src/grpc/conversions.rs +++ b/lib/api/src/grpc/conversions.rs @@ -1556,6 +1556,12 @@ impl From for UpdateResultInternal { } } +impl From for segment_vectors::DenseVector { + fn from(value: DenseVector) -> Self { + value.data + } +} + impl From for DenseVector { fn from(value: segment_vectors::DenseVector) -> Self { Self { data: value } @@ -1570,6 +1576,14 @@ impl From for SparseVector { } } +impl From for sparse::common::sparse_vector::SparseVector { + fn from(value: SparseVector) -> Self { + let SparseVector { indices, values } = value; + + Self { indices, values } + } +} + impl From for MultiDenseVector { fn from(value: segment_vectors::MultiDenseVector) -> Self { let vectors = value @@ -1584,6 +1598,22 @@ impl From for MultiDenseVector { } } +impl From for segment_vectors::MultiDenseVector { + /// Uses the equivalent of [new_unchecked()](segment_vectors::MultiDenseVector::new_unchecked), but rewritten to avoid collecting twice + fn from(value: MultiDenseVector) -> Self { + let dim = value.vectors[0].data.len(); + let inner_vector = value + .vectors + .into_iter() + .flat_map(segment_vectors::DenseVector::from) + .collect(); + Self { + flattened_vectors: inner_vector, + dim, + } + } +} + impl From for RawVector { fn from(value: segment_vectors::Vector) -> Self { use segment_vectors::Vector; @@ -1602,6 +1632,32 @@ impl From for RawVector { } } +impl TryFrom for segment_vectors::Vector { + type Error = Status; + + fn try_from(value: RawVector) -> Result { + use crate::grpc::qdrant::raw_vector::Variant; + + let variant = value + .variant + .ok_or_else(|| Status::invalid_argument("No vector variant provided"))?; + + let vector = match variant { + Variant::Dense(dense) => { + segment_vectors::Vector::Dense(segment_vectors::DenseVector::from(dense)) + } + Variant::Sparse(sparse) => segment_vectors::Vector::Sparse( + sparse::common::sparse_vector::SparseVector::from(sparse), + ), + Variant::MultiDense(multi_dense) => segment_vectors::Vector::MultiDense( + segment_vectors::MultiDenseVector::from(multi_dense), + ), + }; + + Ok(vector) + } +} + impl From for RawVector { fn from(value: segment_vectors::NamedVectorStruct) -> Self { Self::from(value.to_vector()) @@ -1617,6 +1673,24 @@ impl From> for raw_query::Reco } } +impl TryFrom for segment_query::RecoQuery { + type Error = Status; + fn try_from(value: raw_query::Recommend) -> Result { + Ok(Self { + positives: value + .positives + .into_iter() + .map(segment_vectors::Vector::try_from) + .try_collect()?, + negatives: value + .negatives + .into_iter() + .map(segment_vectors::Vector::try_from) + .try_collect()?, + }) + } +} + impl From> for raw_query::RawContextPair { fn from(value: segment_query::ContextPair) -> Self { Self { @@ -1626,6 +1700,28 @@ impl From> for raw_query::Ra } } +impl TryFrom for segment_query::ContextPair { + type Error = Status; + fn try_from(value: raw_query::RawContextPair) -> Result { + Ok(Self { + positive: value + .positive + .map(segment_vectors::Vector::try_from) + .transpose()? + .ok_or_else(|| { + Status::invalid_argument("No positive part of context pair provided") + })?, + negative: value + .negative + .map(segment_vectors::Vector::try_from) + .transpose()? + .ok_or_else(|| { + Status::invalid_argument("No negative part of context pair provided") + })?, + }) + } +} + impl From> for raw_query::Context { fn from(value: segment_query::ContextQuery) -> Self { Self { @@ -1634,6 +1730,19 @@ impl From> for raw_query::C } } +impl TryFrom for segment_query::ContextQuery { + type Error = Status; + fn try_from(value: raw_query::Context) -> Result { + Ok(Self { + pairs: value + .context + .into_iter() + .map(segment_query::ContextPair::try_from) + .try_collect()?, + }) + } +} + impl From> for raw_query::Discovery { fn from(value: segment_query::DiscoveryQuery) -> Self { Self { @@ -1642,3 +1751,21 @@ impl From> for raw_query: } } } + +impl TryFrom for segment_query::DiscoveryQuery { + type Error = Status; + fn try_from(value: raw_query::Discovery) -> Result { + Ok(Self { + target: value + .target + .map(segment_vectors::Vector::try_from) + .transpose()? + .ok_or_else(|| Status::invalid_argument("No target provided"))?, + pairs: value + .context + .into_iter() + .map(segment_query::ContextPair::try_from) + .try_collect()?, + }) + } +} diff --git a/lib/collection/src/operations/types.rs b/lib/collection/src/operations/types.rs index 8f250e3c4a..6c08c0d853 100644 --- a/lib/collection/src/operations/types.rs +++ b/lib/collection/src/operations/types.rs @@ -924,8 +924,10 @@ impl CollectionError { } } - pub fn bad_input(description: String) -> CollectionError { - CollectionError::BadInput { description } + pub fn bad_input(description: impl Into) -> CollectionError { + CollectionError::BadInput { + description: description.into(), + } } pub fn not_found(what: impl Into) -> CollectionError { diff --git a/lib/collection/src/operations/universal_query.rs b/lib/collection/src/operations/universal_query.rs index fce6896797..cc347e0810 100644 --- a/lib/collection/src/operations/universal_query.rs +++ b/lib/collection/src/operations/universal_query.rs @@ -15,7 +15,13 @@ pub mod shard_query { use api::grpc::qdrant as grpc; use common::types::ScoreType; + use itertools::Itertools; + use segment::data_types::vectors::{ + NamedQuery, NamedVectorStruct, Vector, DEFAULT_VECTOR_NAME, + }; use segment::types::{Filter, ScoredPoint, SearchParams, WithPayloadInterface, WithVector}; + use segment::vector_storage::query::{ContextQuery, DiscoveryQuery, RecoQuery}; + use tonic::Status; use crate::operations::query_enum::QueryEnum; @@ -42,7 +48,7 @@ pub mod shard_query { #[derive(Clone)] pub struct ShardPrefetch { - pub prefetches: Option>, + pub prefetches: Vec, pub query: ScoringQuery, pub limit: usize, pub params: Option, @@ -55,7 +61,7 @@ pub mod shard_query { /// Direct translation of the user-facing request, but with all point ids substituted with their corresponding vectors. #[derive(Clone)] pub struct ShardQueryRequest { - pub prefetch: Option>, + pub prefetches: Vec, pub query: ScoringQuery, pub filter: Option, pub score_threshold: Option, @@ -66,9 +72,144 @@ pub mod shard_query { pub with_payload: WithPayloadInterface, } - impl From for ShardQueryRequest { - fn from(_value: grpc::QueryShardPoints) -> Self { - todo!() + impl TryFrom for ShardPrefetch { + type Error = Status; + + fn try_from(value: grpc::query_shard_points::Prefetch) -> Result { + let grpc::query_shard_points::Prefetch { + prefetch, + query, + limit, + params, + filter, + score_threshold, + using, + } = value; + + let shard_prefetch = Self { + prefetches: prefetch + .into_iter() + .map(ShardPrefetch::try_from) + .try_collect()?, + query: query + .map(|query| ScoringQuery::try_from_grpc_query(query, using)) + .transpose()? + .ok_or_else(|| Status::invalid_argument("missing field: query"))?, + limit: limit as usize, + params: params.map(SearchParams::from), + filter: filter.map(Filter::try_from).transpose()?, + score_threshold, + }; + + Ok(shard_prefetch) + } + } + + impl QueryEnum { + fn try_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 query_enum = match variant { + Variant::Nearest(nearest) => { + let vector = Vector::try_from(nearest)?; + let name = match (using, &vector) { + (None, Vector::Sparse(_)) => { + return Err(Status::invalid_argument("Sparse vector must have a name")) + } + ( + Some(name), + Vector::MultiDense(_) | Vector::Sparse(_) | Vector::Dense(_), + ) => name, + (None, Vector::MultiDense(_) | Vector::Dense(_)) => { + DEFAULT_VECTOR_NAME.to_string() + } + }; + let named_vector = NamedVectorStruct::new_from_vector(vector, name); + QueryEnum::Nearest(named_vector) + } + Variant::RecommendBestScore(recommend) => QueryEnum::RecommendBestScore( + 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) + } + } + + impl ScoringQuery { + fn try_from_grpc_query( + query: grpc::query_shard_points::Query, + using: Option, + ) -> Result { + let score = query + .score + .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)?) + } + }; + + Ok(scoring_query) + } + } + + impl TryFrom for ShardQueryRequest { + type Error = Status; + + fn try_from(value: grpc::QueryShardPoints) -> Result { + let grpc::QueryShardPoints { + prefetch, + query, + using, + filter, + limit, + params, + score_threshold, + offset, + with_payload, + with_vectors, + } = value; + + let request = Self { + prefetches: prefetch + .into_iter() + .map(ShardPrefetch::try_from) + .try_collect()?, + query: query + .map(|query| ScoringQuery::try_from_grpc_query(query, using)) + .transpose()? + .ok_or_else(|| Status::invalid_argument("missing field: query"))?, + filter: filter.map(Filter::try_from).transpose()?, + score_threshold, + limit: limit as usize, + offset: offset as usize, + params: params.map(SearchParams::from), + with_vector: with_vectors + .map(WithVector::from) + .unwrap_or(WithVector::Bool(false)), + with_payload: with_payload + .map(WithPayloadInterface::try_from) + .transpose()? + .unwrap_or(WithPayloadInterface::Bool(true)), + }; + + Ok(request) } } @@ -120,11 +261,7 @@ pub mod shard_query { score_threshold, } = value; Self { - prefetch: prefetches - .into_iter() - .flat_map(IntoIterator::into_iter) - .map(Self::from) - .collect(), + prefetch: prefetches.into_iter().map(Self::from).collect(), using: query.get_vector_name().map(ToOwned::to_owned), query: Some(grpc::query_shard_points::Query::from(query)), filter: filter.map(grpc::Filter::from), @@ -138,7 +275,7 @@ pub mod shard_query { impl From for grpc::QueryShardPoints { fn from(value: ShardQueryRequest) -> Self { let ShardQueryRequest { - prefetch, + prefetches, query, filter, score_threshold, @@ -150,9 +287,8 @@ pub mod shard_query { } = value; Self { - prefetch: prefetch + prefetch: prefetches .into_iter() - .flat_map(IntoIterator::into_iter) .map(grpc::query_shard_points::Prefetch::from) .collect(), using: query.get_vector_name().map(ToOwned::to_owned), @@ -226,7 +362,7 @@ pub mod planned_query { offset: req_offset, with_vector, with_payload, - prefetch, + prefetches: prefetch, params, } = request; @@ -235,7 +371,7 @@ pub mod planned_query { let rescore; let offset; - if let Some(prefetch) = prefetch { + if !prefetch.is_empty() { sources = recurse_prefetches(&mut core_searches, prefetch); rescore = Some(query); offset = req_offset; @@ -294,20 +430,8 @@ pub mod planned_query { score_threshold, } = prefetch; - let source = match prefetches { - Some(inner_prefetches) => { - let sources = recurse_prefetches(core_searches, inner_prefetches); - - let prefetch_plan = PrefetchPlan { - sources, - merge: PrefetchMerge { - rescore: Some(query), - limit, - }, - }; - PrefetchSource::Prefetch(prefetch_plan) - } - None => match query { + let source = if prefetches.is_empty() { + match query { ScoringQuery::Vector(query_enum) => { let core_search = CoreSearchRequest { query: query_enum, @@ -325,7 +449,18 @@ pub mod planned_query { PrefetchSource::BatchIdx(idx) } - }, + } + } else { + let sources = recurse_prefetches(core_searches, prefetches); + + let prefetch_plan = PrefetchPlan { + sources, + merge: PrefetchMerge { + rescore: Some(query), + limit, + }, + }; + PrefetchSource::Prefetch(prefetch_plan) }; sources.push(source); } diff --git a/lib/segment/src/data_types/vectors.rs b/lib/segment/src/data_types/vectors.rs index b791004992..209f16b26a 100644 --- a/lib/segment/src/data_types/vectors.rs +++ b/lib/segment/src/data_types/vectors.rs @@ -630,6 +630,7 @@ impl Validate for NamedQuery { impl NamedQuery> { pub fn new(query: RecoQuery, using: Option) -> Self { + // TODO: maybe validate there is no sparse vector without vector name NamedQuery { query, using } } } diff --git a/src/tonic/api/points_internal_api.rs b/src/tonic/api/points_internal_api.rs index 97d52a3e35..2b5e55e31d 100644 --- a/src/tonic/api/points_internal_api.rs +++ b/src/tonic/api/points_internal_api.rs @@ -7,12 +7,16 @@ use api::grpc::qdrant::{ CreateFieldIndexCollectionInternal, DeleteFieldIndexCollectionInternal, DeletePayloadPointsInternal, DeletePointsInternal, DeleteVectorsInternal, GetPointsInternal, GetResponse, PointsOperationResponseInternal, QueryPointsInternal, QueryResponse, - RecommendPointsInternal, RecommendResponse, ScrollPointsInternal, ScrollResponse, - SearchBatchResponse, SetPayloadPointsInternal, SyncPointsInternal, UpdateVectorsInternal, - UpsertPointsInternal, + QueryShardPoints, RecommendPointsInternal, RecommendResponse, ScrollPointsInternal, + ScrollResponse, SearchBatchResponse, SetPayloadPointsInternal, SyncPointsInternal, + UpdateVectorsInternal, UpsertPointsInternal, }; +use collection::operations::shard_selector_internal::ShardSelectorInternal; +use collection::operations::universal_query::shard_query::ShardQueryRequest; +use collection::shards::shard::ShardId; use storage::content_manager::toc::TableOfContent; use storage::rbac::Access; +use tokio::time::Instant; use tonic::{Request, Response, Status}; use super::points_common::core_search_list; @@ -36,6 +40,45 @@ impl PointsInternalService { } } +#[allow(unused_variables)] // TODO(universal-query): remove +pub async fn query( + toc: &TableOfContent, + collection_name: String, + query_points: QueryShardPoints, + shard_selection: Option, + access: Access, +) -> Result, Status> { + let request = ShardQueryRequest::try_from(query_points).map_err(Status::from)?; + + let timing = Instant::now(); + + // As this function is handling an internal request, + // we can assume that shard_key is already resolved + let shard_selection = match shard_selection { + None => { + debug_assert!(false, "Shard selection is expected for internal request"); + ShardSelectorInternal::All + } + Some(shard_id) => ShardSelectorInternal::ShardId(shard_id), + }; + + // TODO(universal-query): add `query()` to TableOfContent + // let scored_points = toc + // .query( + // &collection_name, + // request, + // shard_selection, + // access, + // ) + // .await + // .map_err(error_to_status)?; + + // TODO(universal-query): convert response to grpc + todo!() + + // Ok(Response::new(response)) +} + #[tonic::async_trait] impl PointsInternal for PointsInternalService { async fn upsert( @@ -431,9 +474,32 @@ impl PointsInternal for PointsInternalService { async fn query( &self, - _request: Request, + request: Request, ) -> Result, Status> { - // TODO(universal-query): Implement this - return Err(Status::unimplemented("Not yet implemented")); + // TODO(universal-query): validate + // validate_and_log(request.get_ref()); + + let QueryPointsInternal { + collection_name, + shard_id, + query_points, + } = request.into_inner(); + + let query_points = + query_points.ok_or_else(|| Status::invalid_argument("QueryPoints is missing"))?; + + // TODO(universal-query): add timeout + // let timeout = timeout.map(Duration::from_secs); + + // Individual `read_consistency` values are ignored + + query( + self.toc.as_ref(), + collection_name, + query_points, + shard_id, + FULL_ACCESS.clone(), + ) + .await } }