qdrant_edge/shard/
search.rs

1use crate::common::types::ScoreType;
2#[cfg(feature = "api")]
3use itertools::Itertools as _;
4#[cfg(feature = "api")]
5use crate::segment::data_types::vectors::NamedQuery;
6use crate::segment::types::{Filter, SearchParams, WithPayloadInterface, WithVector};
7#[cfg(feature = "api")]
8use crate::segment::{data_types::vectors::VectorInternal, vector_storage::query::ContextPair};
9
10use crate::shard::query::query_enum::QueryEnum;
11
12/// DEPRECATED: Search method should be removed and replaced with `ShardQueryRequest`
13#[derive(Clone, Debug, PartialEq)]
14pub struct CoreSearchRequest {
15    /// Every kind of query that can be performed on segment level
16    pub query: QueryEnum,
17    /// Look only for points which satisfies this conditions
18    pub filter: Option<Filter>,
19    /// Additional search params
20    pub params: Option<SearchParams>,
21    /// Max number of result to return
22    pub limit: usize,
23    /// Offset of the first result to return.
24    /// May be used to paginate results.
25    /// Note: large offset values may cause performance issues.
26    pub offset: usize,
27    /// Select which payload to return with the response. Default is false.
28    pub with_payload: Option<WithPayloadInterface>,
29    /// Options for specifying which vectors to include into response. Default is false.
30    pub with_vector: Option<WithVector>,
31    pub score_threshold: Option<ScoreType>,
32}
33
34impl CoreSearchRequest {
35    pub fn search_rate_cost(&self) -> usize {
36        let mut cost = self.query.search_cost();
37
38        if let Some(filter) = &self.filter {
39            cost += filter.total_conditions_count();
40        }
41
42        cost
43    }
44}
45
46#[cfg(feature = "api")]
47impl From<api::rest::SearchRequestInternal> for CoreSearchRequest {
48    fn from(request: api::rest::SearchRequestInternal) -> Self {
49        #[cfg(feature = "api")]
50        use crate::segment::data_types::vectors::NamedVectorStruct;
51
52        let api::rest::SearchRequestInternal {
53            vector,
54            filter,
55            score_threshold,
56            limit,
57            offset,
58            params,
59            with_vector,
60            with_payload,
61        } = request;
62        Self {
63            query: QueryEnum::Nearest(NamedQuery::from(NamedVectorStruct::from(vector))),
64            filter,
65            params,
66            limit,
67            offset: offset.unwrap_or_default(),
68            with_payload,
69            with_vector,
70            score_threshold,
71        }
72    }
73}
74
75#[cfg(feature = "api")]
76impl TryFrom<api::grpc::qdrant::CoreSearchPoints> for CoreSearchRequest {
77    type Error = tonic::Status;
78
79    fn try_from(value: api::grpc::qdrant::CoreSearchPoints) -> Result<Self, Self::Error> {
80        use crate::segment::data_types::vectors::VectorInternal;
81        use crate::segment::vector_storage::query::{ContextQuery, DiscoverQuery, RecoQuery};
82
83        let query = value
84            .query
85            .and_then(|query| query.query)
86            .map(|query| {
87                Ok(match query {
88                    api::grpc::qdrant::query_enum::Query::NearestNeighbors(vector) => {
89                        let vector_internal = VectorInternal::try_from(vector)?;
90                        QueryEnum::Nearest(NamedQuery::from(
91                            api::grpc::conversions::into_named_vector_struct(
92                                value.vector_name,
93                                vector_internal,
94                            )?,
95                        ))
96                    }
97                    api::grpc::qdrant::query_enum::Query::RecommendBestScore(query) => {
98                        QueryEnum::RecommendBestScore(NamedQuery {
99                            query: RecoQuery::try_from(query)?,
100                            using: value.vector_name,
101                        })
102                    }
103                    api::grpc::qdrant::query_enum::Query::RecommendSumScores(query) => {
104                        QueryEnum::RecommendSumScores(NamedQuery {
105                            query: RecoQuery::try_from(query)?,
106                            using: value.vector_name,
107                        })
108                    }
109                    api::grpc::qdrant::query_enum::Query::Discover(query) => {
110                        let Some(target) = query.target else {
111                            return Err(tonic::Status::invalid_argument("Target is not specified"));
112                        };
113
114                        let pairs = query
115                            .context
116                            .into_iter()
117                            .map(try_context_pair_from_grpc)
118                            .try_collect()?;
119
120                        QueryEnum::Discover(NamedQuery {
121                            query: DiscoverQuery::new(target.try_into()?, pairs),
122                            using: value.vector_name,
123                        })
124                    }
125                    api::grpc::qdrant::query_enum::Query::Context(query) => {
126                        let pairs = query
127                            .context
128                            .into_iter()
129                            .map(try_context_pair_from_grpc)
130                            .try_collect()?;
131
132                        QueryEnum::Context(NamedQuery {
133                            query: ContextQuery::new(pairs),
134                            using: value.vector_name,
135                        })
136                    }
137                })
138            })
139            .transpose()?
140            .ok_or_else(|| tonic::Status::invalid_argument("Query is not specified"))?;
141
142        Ok(Self {
143            query,
144            filter: value.filter.map(|f| f.try_into()).transpose()?,
145            params: value.params.map(Into::into),
146            limit: value.limit as usize,
147            offset: value.offset.unwrap_or_default() as usize,
148            with_payload: value.with_payload.map(|wp| wp.try_into()).transpose()?,
149            with_vector: Some(value.with_vectors.map(Into::into).unwrap_or_default()),
150            score_threshold: value.score_threshold,
151        })
152    }
153}
154
155#[cfg(feature = "api")]
156fn try_context_pair_from_grpc(
157    pair: api::grpc::qdrant::ContextPair,
158) -> Result<ContextPair<VectorInternal>, tonic::Status> {
159    let api::grpc::qdrant::ContextPair { positive, negative } = pair;
160    match (positive, negative) {
161        (Some(positive), Some(negative)) => Ok(ContextPair {
162            positive: positive.try_into()?,
163            negative: negative.try_into()?,
164        }),
165        _ => Err(tonic::Status::invalid_argument(
166            "All context pairs must have both positive and negative parts",
167        )),
168    }
169}
170
171#[cfg(feature = "api")]
172impl TryFrom<api::grpc::qdrant::SearchPoints> for CoreSearchRequest {
173    type Error = tonic::Status;
174
175    fn try_from(value: api::grpc::qdrant::SearchPoints) -> Result<Self, Self::Error> {
176        use crate::sparse::common::sparse_vector::validate_sparse_vector_impl;
177
178        let api::grpc::qdrant::SearchPoints {
179            collection_name: _,
180            vector,
181            filter,
182            limit,
183            with_payload,
184            params,
185            score_threshold,
186            offset,
187            vector_name,
188            with_vectors,
189            read_consistency: _,
190            timeout: _,
191            shard_key_selector: _,
192            sparse_indices,
193        } = value;
194
195        if let Some(sparse_indices) = &sparse_indices {
196            let api::grpc::qdrant::SparseIndices { data } = sparse_indices;
197            validate_sparse_vector_impl(data, &vector).map_err(|e| {
198                tonic::Status::invalid_argument(format!(
199                    "Sparse indices does not match sparse vector conditions: {e}"
200                ))
201            })?;
202        }
203
204        let vector_internal =
205            VectorInternal::from_vector_and_indices(vector, sparse_indices.map(|v| v.data));
206
207        let vector_struct =
208            api::grpc::conversions::into_named_vector_struct(vector_name, vector_internal)?;
209
210        Ok(Self {
211            query: QueryEnum::Nearest(NamedQuery::from(vector_struct)),
212            filter: filter.map(Filter::try_from).transpose()?,
213            params: params.map(SearchParams::from),
214            limit: limit as usize,
215            offset: offset.map(|v| v as usize).unwrap_or_default(),
216            with_payload: with_payload
217                .map(WithPayloadInterface::try_from)
218                .transpose()?,
219            with_vector: with_vectors.map(WithVector::from),
220            score_threshold: score_threshold.map(|s| s as ScoreType),
221        })
222    }
223}
224
225#[derive(Debug, Clone)]
226pub struct CoreSearchRequestBatch {
227    pub searches: Vec<CoreSearchRequest>,
228}
qdrant_edge/shard/search.rs

qdrant_edge/shard/
search.rs
