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

@@ -1556,6 +1556,12 @@ impl From<UpdateResult> for UpdateResultInternal {
     }
 }
 
+impl From<DenseVector> for segment_vectors::DenseVector {
+    fn from(value: DenseVector) -> Self {
+        value.data
+    }
+}
+
 impl From<segment_vectors::DenseVector> for DenseVector {
     fn from(value: segment_vectors::DenseVector) -> Self {
         Self { data: value }
@@ -1570,6 +1576,14 @@ impl From<sparse::common::sparse_vector::SparseVector> for SparseVector {
     }
 }
 
+impl From<SparseVector> for sparse::common::sparse_vector::SparseVector {
+    fn from(value: SparseVector) -> Self {
+        let SparseVector { indices, values } = value;
+
+        Self { indices, values }
+    }
+}
+
 impl From<segment_vectors::MultiDenseVector> for MultiDenseVector {
     fn from(value: segment_vectors::MultiDenseVector) -> Self {
         let vectors = value
@@ -1584,6 +1598,22 @@ impl From<segment_vectors::MultiDenseVector> for MultiDenseVector {
     }
 }
 
+impl From<MultiDenseVector> 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<segment_vectors::Vector> for RawVector {
     fn from(value: segment_vectors::Vector) -> Self {
         use segment_vectors::Vector;
@@ -1602,6 +1632,32 @@ impl From<segment_vectors::Vector> for RawVector {
     }
 }
 
+impl TryFrom<RawVector> for segment_vectors::Vector {
+    type Error = Status;
+
+    fn try_from(value: RawVector) -> Result<Self, Self::Error> {
+        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<segment_vectors::NamedVectorStruct> for RawVector {
     fn from(value: segment_vectors::NamedVectorStruct) -> Self {
         Self::from(value.to_vector())
@@ -1617,6 +1673,24 @@ impl From<segment_query::RecoQuery<segment_vectors::Vector>> for raw_query::Reco
     }
 }
 
+impl TryFrom<raw_query::Recommend> for segment_query::RecoQuery<segment_vectors::Vector> {
+    type Error = Status;
+    fn try_from(value: raw_query::Recommend) -> Result<Self, Self::Error> {
+        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<segment_query::ContextPair<segment_vectors::Vector>> for raw_query::RawContextPair {
     fn from(value: segment_query::ContextPair<segment_vectors::Vector>) -> Self {
         Self {
@@ -1626,6 +1700,28 @@ impl From<segment_query::ContextPair<segment_vectors::Vector>> for raw_query::Ra
     }
 }
 
+impl TryFrom<raw_query::RawContextPair> for segment_query::ContextPair<segment_vectors::Vector> {
+    type Error = Status;
+    fn try_from(value: raw_query::RawContextPair) -> Result<Self, Self::Error> {
+        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<segment_query::ContextQuery<segment_vectors::Vector>> for raw_query::Context {
     fn from(value: segment_query::ContextQuery<segment_vectors::Vector>) -> Self {
         Self {
@@ -1634,6 +1730,19 @@ impl From<segment_query::ContextQuery<segment_vectors::Vector>> for raw_query::C
     }
 }
 
+impl TryFrom<raw_query::Context> for segment_query::ContextQuery<segment_vectors::Vector> {
+    type Error = Status;
+    fn try_from(value: raw_query::Context) -> Result<Self, Self::Error> {
+        Ok(Self {
+            pairs: value
+                .context
+                .into_iter()
+                .map(segment_query::ContextPair::try_from)
+                .try_collect()?,
+        })
+    }
+}
+
 impl From<segment_query::DiscoveryQuery<segment_vectors::Vector>> for raw_query::Discovery {
     fn from(value: segment_query::DiscoveryQuery<segment_vectors::Vector>) -> Self {
         Self {
@@ -1642,3 +1751,21 @@ impl From<segment_query::DiscoveryQuery<segment_vectors::Vector>> for raw_query:
         }
     }
 }
+
+impl TryFrom<raw_query::Discovery> for segment_query::DiscoveryQuery<segment_vectors::Vector> {
+    type Error = Status;
+    fn try_from(value: raw_query::Discovery) -> Result<Self, Self::Error> {
+        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()?,
+        })
+    }
+}

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<String>) -> CollectionError {
+        CollectionError::BadInput {
+            description: description.into(),
+        }
     }
 
     pub fn not_found(what: impl Into<String>) -> CollectionError {

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<Vec<ShardPrefetch>>,
+        pub prefetches: Vec<ShardPrefetch>,
         pub query: ScoringQuery,
         pub limit: usize,
         pub params: Option<SearchParams>,
@@ -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<Vec<ShardPrefetch>>,
+        pub prefetches: Vec<ShardPrefetch>,
         pub query: ScoringQuery,
         pub filter: Option<Filter>,
         pub score_threshold: Option<ScoreType>,
@@ -66,9 +72,144 @@ pub mod shard_query {
         pub with_payload: WithPayloadInterface,
     }
 
-    impl From<grpc::QueryShardPoints> for ShardQueryRequest {
-        fn from(_value: grpc::QueryShardPoints) -> Self {
-            todo!()
+    impl TryFrom<grpc::query_shard_points::Prefetch> for ShardPrefetch {
+        type Error = Status;
+
+        fn try_from(value: grpc::query_shard_points::Prefetch) -> Result<Self, Self::Error> {
+            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<String>,
+        ) -> Result<Self, Status> {
+            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<String>,
+        ) -> Result<Self, Status> {
+            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<grpc::QueryShardPoints> for ShardQueryRequest {
+        type Error = Status;
+
+        fn try_from(value: grpc::QueryShardPoints) -> Result<Self, Self::Error> {
+            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<ShardQueryRequest> 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);
         }

lib/segment/src/data_types/vectors.rs

@@ -630,6 +630,7 @@ impl<T: Validate> Validate for NamedQuery<T> {
 
 impl NamedQuery<RecoQuery<Vector>> {
     pub fn new(query: RecoQuery<Vector>, using: Option<String>) -> Self {
+        // TODO: maybe validate there is no sparse vector without vector name
         NamedQuery { query, using }
     }
 }

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<ShardId>,
+    access: Access,
+) -> Result<Response<QueryResponse>, 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<QueryPointsInternal>,
+        request: Request<QueryPointsInternal>,
     ) -> Result<Response<QueryResponse>, 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
     }
 }