Move CoreSearchRequest and QueryEnum from collection into shard crate (#7169)

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<CountResult> for api::grpc::qdrant::CountResult {
     }
 }
 
-impl TryFrom<api::grpc::qdrant::SearchPoints> for CoreSearchRequest {
-    type Error = Status;
-    fn try_from(value: api::grpc::qdrant::SearchPoints) -> Result<Self, Self::Error> {
-        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<QueryEnum> 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<CollectionCoreSearchRequest<'a>> for api::grpc::qdrant::CoreSearchPoints {
     fn from(value: CollectionCoreSearchRequest<'a>) -> Self {
         let (collection_id, request) = value.0;
@@ -1184,102 +1081,6 @@ impl TryFrom<api::grpc::qdrant::TargetVector> for RecommendExample {
     }
 }
 
-fn try_context_pair_from_grpc(
-    pair: api::grpc::qdrant::ContextPair,
-) -> Result<ContextPair<VectorInternal>, 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<api::grpc::qdrant::CoreSearchPoints> for CoreSearchRequest {
-    type Error = Status;
-
-    fn try_from(value: api::grpc::qdrant::CoreSearchPoints) -> Result<Self, Self::Error> {
-        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<PointGroup> for api::grpc::qdrant::PointGroup {
     type Error = OperationError;
     fn try_from(group: PointGroup) -> Result<Self, Self::Error> {

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;

lib/collection/src/operations/query_enum.rs

@@ -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<VectorInternal>),
-    RecommendBestScore(NamedQuery<RecoQuery<VectorInternal>>),
-    RecommendSumScores(NamedQuery<RecoQuery<VectorInternal>>),
-    Discover(NamedQuery<DiscoveryQuery<VectorInternal>>),
-    Context(NamedQuery<ContextQuery<VectorInternal>>),
-}
-
-impl QueryEnum {
-    /// Iterate over all vectors in the query.
-    fn vectors(&self) -> Box<dyn Iterator<Item = &VectorInternal> + '_> {
-        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<DenseVector> for QueryEnum {
-    fn from(vector: DenseVector) -> Self {
-        QueryEnum::Nearest(NamedQuery {
-            query: VectorInternal::Dense(vector),
-            using: None,
-        })
-    }
-}
-
-impl From<NamedQuery<DiscoveryQuery<VectorInternal>>> for QueryEnum {
-    fn from(query: NamedQuery<DiscoveryQuery<VectorInternal>>) -> Self {
-        QueryEnum::Discover(query)
-    }
-}
-
-impl AsRef<QueryEnum> for QueryEnum {
-    fn as_ref(&self) -> &QueryEnum {
-        self
-    }
-}

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<SearchRequest>,
 }
 
-#[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<Filter>,
-    /// Additional search params
-    pub params: Option<SearchParams>,
-    /// 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<WithPayloadInterface>,
-    /// Options for specifying which vectors to include into response. Default is false.
-    pub with_vector: Option<WithVector>,
-    pub score_threshold: Option<ScoreType>,
-}
-
 #[derive(Debug, Clone)]
 pub struct CoreSearchRequestBatch {
     pub searches: Vec<CoreSearchRequest>,
@@ -2013,31 +1993,6 @@ pub enum NodeType {
     Listener,
 }
 
-impl From<SearchRequestInternal> 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<SearchRequestInternal> for ShardQueryRequest {
     fn from(value: SearchRequestInternal) -> Self {
         let SearchRequestInternal {
@@ -2094,18 +2049,6 @@ impl From<CoreSearchRequest> for ShardQueryRequest {
     }
 }
 
-impl From<QueryEnum> 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 {

lib/collection/src/operations/universal_query/shard_query.rs

@@ -297,55 +297,53 @@ impl TryFrom<grpc::query_shard_points::Prefetch> for ShardPrefetch {
     }
 }
 
-impl QueryEnum {
-    fn try_from_grpc_raw_query(
-        raw_query: grpc::RawQuery,
-        using: Option<VectorNameBuf>,
-    ) -> Result<Self, Status> {
-        use grpc::raw_query::Variant;
+fn query_enum_from_grpc_raw_query(
+    raw_query: grpc::RawQuery,
+    using: Option<VectorNameBuf>,
+) -> Result<QueryEnum, Status> {
+    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<i32> 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<QueryEnum> 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<ScoringQuery> 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 {

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;

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;

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<Filter>,
+    /// Additional search params
+    pub params: Option<SearchParams>,
+    /// 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<WithPayloadInterface>,
+    /// Options for specifying which vectors to include into response. Default is false.
+    pub with_vector: Option<WithVector>,
+    pub score_threshold: Option<ScoreType>,
+}
+
+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<SearchRequestInternal> 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<api::grpc::qdrant::CoreSearchPoints> for CoreSearchRequest {
+    type Error = tonic::Status;
+
+    fn try_from(value: api::grpc::qdrant::CoreSearchPoints) -> Result<Self, Self::Error> {
+        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<ContextPair<VectorInternal>, 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<api::grpc::qdrant::SearchPoints> for CoreSearchRequest {
+    type Error = tonic::Status;
+
+    fn try_from(value: api::grpc::qdrant::SearchPoints) -> Result<Self, Self::Error> {
+        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<VectorInternal>),
+    RecommendBestScore(NamedQuery<RecoQuery<VectorInternal>>),
+    RecommendSumScores(NamedQuery<RecoQuery<VectorInternal>>),
+    Discover(NamedQuery<DiscoveryQuery<VectorInternal>>),
+    Context(NamedQuery<ContextQuery<VectorInternal>>),
+}
+
+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<Item = &'a VectorInternal>) -> usize {
+    vectors
+        .into_iter()
+        .map(VectorInternal::similarity_cost)
+        .sum()
+}
+
+impl AsRef<QueryEnum> for QueryEnum {
+    fn as_ref(&self) -> &QueryEnum {
+        self
+    }
+}
+
+impl From<DenseVector> for QueryEnum {
+    fn from(vector: DenseVector) -> Self {
+        QueryEnum::Nearest(NamedQuery {
+            query: VectorInternal::Dense(vector),
+            using: None,
+        })
+    }
+}
+
+impl From<NamedQuery<DiscoveryQuery<VectorInternal>>> for QueryEnum {
+    fn from(query: NamedQuery<DiscoveryQuery<VectorInternal>>) -> Self {
+        QueryEnum::Discover(query)
+    }
+}
+
+impl From<QueryEnum> 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<QueryEnum> 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(),
+                    },
+                )),
+            },
+        }
+    }
+}