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(|p| p.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(
150                value
151                    .with_vectors
152                    .map(|with_vectors| with_vectors.into())
153                    .unwrap_or_default(),
154            ),
155            score_threshold: value.score_threshold,
156        })
157    }
158}
159
160#[cfg(feature = "api")]
161fn try_context_pair_from_grpc(
162    pair: api::grpc::qdrant::ContextPair,
163) -> Result<ContextPair<VectorInternal>, tonic::Status> {
164    let api::grpc::qdrant::ContextPair { positive, negative } = pair;
165    match (positive, negative) {
166        (Some(positive), Some(negative)) => Ok(ContextPair {
167            positive: positive.try_into()?,
168            negative: negative.try_into()?,
169        }),
170        _ => Err(tonic::Status::invalid_argument(
171            "All context pairs must have both positive and negative parts",
172        )),
173    }
174}
175
176#[cfg(feature = "api")]
177impl TryFrom<api::grpc::qdrant::SearchPoints> for CoreSearchRequest {
178    type Error = tonic::Status;
179
180    fn try_from(value: api::grpc::qdrant::SearchPoints) -> Result<Self, Self::Error> {
181        use crate::sparse::common::sparse_vector::validate_sparse_vector_impl;
182
183        let api::grpc::qdrant::SearchPoints {
184            collection_name: _,
185            vector,
186            filter,
187            limit,
188            with_payload,
189            params,
190            score_threshold,
191            offset,
192            vector_name,
193            with_vectors,
194            read_consistency: _,
195            timeout: _,
196            shard_key_selector: _,
197            sparse_indices,
198        } = value;
199
200        if let Some(sparse_indices) = &sparse_indices {
201            let api::grpc::qdrant::SparseIndices { data } = sparse_indices;
202            validate_sparse_vector_impl(data, &vector).map_err(|e| {
203                tonic::Status::invalid_argument(format!(
204                    "Sparse indices does not match sparse vector conditions: {e}"
205                ))
206            })?;
207        }
208
209        let vector_internal =
210            VectorInternal::from_vector_and_indices(vector, sparse_indices.map(|v| v.data));
211
212        let vector_struct =
213            api::grpc::conversions::into_named_vector_struct(vector_name, vector_internal)?;
214
215        Ok(Self {
216            query: QueryEnum::Nearest(NamedQuery::from(vector_struct)),
217            filter: filter.map(Filter::try_from).transpose()?,
218            params: params.map(SearchParams::from),
219            limit: limit as usize,
220            offset: offset.map(|v| v as usize).unwrap_or_default(),
221            with_payload: with_payload
222                .map(WithPayloadInterface::try_from)
223                .transpose()?,
224            with_vector: with_vectors.map(WithVector::from),
225            score_threshold: score_threshold.map(|s| s as ScoreType),
226        })
227    }
228}
229
230#[derive(Debug, Clone)]
231pub struct CoreSearchRequestBatch {
232    pub searches: Vec<CoreSearchRequest>,
233}
qdrant_edge/shard/search.rs

qdrant_edge/shard/
search.rs
