Skip to main content

qdrant_edge/shard/query/
mod.rs

1#[cfg(feature = "api")]
2mod conversions;
3pub mod formula;
4pub mod mmr;
5pub mod planned_query;
6pub mod query_enum;
7pub mod scroll;
8mod validation;
9
10pub mod query_context;
11#[cfg(test)]
12mod tests;
13
14use crate::common::types::ScoreType;
15use ordered_float::OrderedFloat;
16use crate::segment::data_types::order_by::OrderBy;
17use crate::segment::data_types::vectors::VectorInternal;
18use crate::segment::index::query_optimization::rescore_formula::parsed_formula::ParsedFormula;
19use crate::segment::types::*;
20use serde::Serialize;
21
22use self::query_enum::*;
23use crate::shard::search::CoreSearchRequest;
24
25/// Internal response type for a universal query request.
26///
27/// Capable of returning multiple intermediate results if needed, like the case of RRF (Reciprocal Rank Fusion)
28pub type ShardQueryResponse = Vec<Vec<ScoredPoint>>;
29
30/// Internal representation of a universal query request.
31///
32/// Direct translation of the user-facing request, but with all point ids substituted with their corresponding vectors.
33///
34/// For the case of formula queries, it collects conditions and variables too.
35#[derive(Clone, Debug, Hash, Serialize)]
36pub struct ShardQueryRequest {
37    pub prefetches: Vec<ShardPrefetch>,
38    #[serde(skip_serializing_if = "Option::is_none")]
39    pub query: Option<ScoringQuery>,
40    #[serde(skip_serializing_if = "Option::is_none")]
41    pub filter: Option<Filter>,
42    #[serde(skip_serializing_if = "Option::is_none")]
43    pub score_threshold: Option<OrderedFloat<ScoreType>>,
44    pub limit: usize,
45    pub offset: usize,
46    /// Search params for when there is no prefetch
47    #[serde(skip_serializing_if = "Option::is_none")]
48    pub params: Option<SearchParams>,
49    pub with_vector: WithVector,
50    pub with_payload: WithPayloadInterface,
51}
52
53impl ShardQueryRequest {
54    pub fn prefetches_depth(&self) -> usize {
55        self.prefetches
56            .iter()
57            .map(ShardPrefetch::depth)
58            .max()
59            .unwrap_or(0)
60    }
61
62    pub fn filter_refs(&self) -> Vec<Option<&Filter>> {
63        let mut filters = vec![];
64        filters.push(self.filter.as_ref());
65
66        for prefetch in &self.prefetches {
67            filters.extend(prefetch.filter_refs())
68        }
69
70        filters
71    }
72}
73
74#[derive(Clone, Debug, Hash, Serialize)]
75pub struct ShardPrefetch {
76    pub prefetches: Vec<ShardPrefetch>,
77    #[serde(skip_serializing_if = "Option::is_none")]
78    pub query: Option<ScoringQuery>,
79    pub limit: usize,
80    #[serde(skip_serializing_if = "Option::is_none")]
81    pub params: Option<SearchParams>,
82    #[serde(skip_serializing_if = "Option::is_none")]
83    pub filter: Option<Filter>,
84    #[serde(skip_serializing_if = "Option::is_none")]
85    pub score_threshold: Option<OrderedFloat<ScoreType>>,
86}
87
88impl ShardPrefetch {
89    pub fn depth(&self) -> usize {
90        let mut depth = 1;
91        for prefetch in &self.prefetches {
92            depth = depth.max(prefetch.depth() + 1);
93        }
94        depth
95    }
96
97    fn filter_refs(&self) -> Vec<Option<&Filter>> {
98        let mut filters = vec![];
99
100        filters.push(self.filter.as_ref());
101
102        for prefetch in &self.prefetches {
103            filters.extend(prefetch.filter_refs())
104        }
105
106        filters
107    }
108}
109
110/// Same as `Query`, but with the resolved vector references.
111#[derive(Clone, Debug, PartialEq, Hash, Serialize)]
112pub enum ScoringQuery {
113    /// Score points against some vector(s)
114    Vector(QueryEnum),
115
116    /// Reciprocal rank fusion
117    Fusion(FusionInternal),
118
119    /// Order by a payload field
120    OrderBy(OrderBy),
121
122    /// Score boosting via an arbitrary formula
123    Formula(ParsedFormula),
124
125    /// Sample points
126    Sample(SampleInternal),
127
128    /// Maximal Marginal Relevance
129    ///
130    /// This one behaves a little differently than the other scorings, since it is two parts.
131    /// It will create one nearest neighbor search in segment space and then try to resolve MMR algorithm higher up.
132    ///
133    /// E.g. If it is the root query of a request:
134    ///   1. Performs search all the way down to segments.
135    ///   2. MMR gets calculated once results reach collection level.
136    Mmr(MmrInternal),
137}
138
139impl ScoringQuery {
140    /// Whether the query needs the prefetches results from all shards to compute the final score
141    ///
142    /// If false, there is a single list of scored points which contain the final score.
143    pub fn needs_intermediate_results(&self) -> bool {
144        match self {
145            Self::Fusion(fusion) => match fusion {
146                // We need the ranking information of each prefetch
147                FusionInternal::Rrf { k: _, weights: _ } => true,
148                // We need the score distribution information of each prefetch
149                FusionInternal::Dbsf => true,
150            },
151            // MMR is a nearest neighbors search before computing diversity at collection level
152            Self::Mmr(_) => false,
153            Self::Vector(_) | Self::OrderBy(_) | Self::Formula(_) | Self::Sample(_) => false,
154        }
155    }
156
157    /// Get the vector name if it is scored against a vector
158    pub fn get_vector_name(&self) -> Option<&VectorName> {
159        match self {
160            Self::Vector(query) => Some(query.get_vector_name()),
161            Self::Mmr(mmr) => Some(&mmr.using),
162            _ => None,
163        }
164    }
165}
166
167#[derive(Clone, Debug, PartialEq, Hash, Serialize)]
168pub enum FusionInternal {
169    /// Reciprocal Rank Fusion with optional weights per prefetch
170    Rrf {
171        k: usize,
172        /// Weights for each prefetch source. Higher weight = more influence on final ranking.
173        /// If None, all sources are weighted equally.
174        weights: Option<Vec<ordered_float::OrderedFloat<f32>>>,
175    },
176    /// Distribution-based score fusion
177    Dbsf,
178}
179
180#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, Serialize)]
181pub enum SampleInternal {
182    Random,
183}
184
185/// Maximal Marginal Relevance configuration
186#[derive(Clone, Debug, PartialEq, Hash, Serialize)]
187pub struct MmrInternal {
188    /// Query vector, used to get the relevance of each point.
189    pub vector: VectorInternal,
190    /// Vector name to use for similarity computation, defaults to empty string (default vector)
191    pub using: VectorNameBuf,
192    /// Lambda parameter controlling diversity vs relevance trade-off (0.0 = full diversity, 1.0 = full relevance)
193    pub lambda: OrderedFloat<f32>,
194    /// Maximum number of candidates to preselect using nearest neighbors.
195    pub candidates_limit: usize,
196}
197
198impl From<CoreSearchRequest> for ShardQueryRequest {
199    fn from(value: CoreSearchRequest) -> Self {
200        let CoreSearchRequest {
201            query,
202            filter,
203            score_threshold,
204            limit,
205            offset,
206            params,
207            with_vector,
208            with_payload,
209        } = value;
210
211        Self {
212            prefetches: vec![],
213            query: Some(ScoringQuery::Vector(query)),
214            filter,
215            score_threshold: score_threshold.map(OrderedFloat),
216            limit,
217            offset,
218            params,
219            with_vector: with_vector.unwrap_or_default(),
220            with_payload: with_payload.unwrap_or_default(),
221        }
222    }
223}