qdrant_edge/shard/query/
query_enum.rs

1use crate::segment::data_types::vectors::*;
2use crate::segment::types::VectorName;
3use crate::segment::vector_storage::query::*;
4use serde::Serialize;
5use crate::sparse::common::sparse_vector::SparseVector;
6
7/// Every kind of vector query that can be performed on segment level.
8#[derive(Clone, Debug, PartialEq, Hash, Serialize)]
9pub enum QueryEnum {
10    Nearest(NamedQuery<VectorInternal>),
11    RecommendBestScore(NamedQuery<RecoQuery<VectorInternal>>),
12    RecommendSumScores(NamedQuery<RecoQuery<VectorInternal>>),
13    Discover(NamedQuery<DiscoverQuery<VectorInternal>>),
14    Context(NamedQuery<ContextQuery<VectorInternal>>),
15    FeedbackNaive(NamedQuery<NaiveFeedbackQuery<VectorInternal>>),
16}
17
18impl QueryEnum {
19    pub fn get_vector_name(&self) -> &VectorName {
20        match self {
21            QueryEnum::Nearest(vector) => vector.get_name(),
22            QueryEnum::RecommendBestScore(reco_query) => reco_query.get_name(),
23            QueryEnum::RecommendSumScores(reco_query) => reco_query.get_name(),
24            QueryEnum::Discover(discover_query) => discover_query.get_name(),
25            QueryEnum::Context(context_query) => context_query.get_name(),
26            QueryEnum::FeedbackNaive(feedback_query) => feedback_query.get_name(),
27        }
28    }
29
30    /// Only when the distance is the scoring, this will return true.
31    pub fn is_distance_scored(&self) -> bool {
32        match self {
33            QueryEnum::Nearest(_) => true,
34            QueryEnum::RecommendBestScore(_)
35            | QueryEnum::RecommendSumScores(_)
36            | QueryEnum::Discover(_)
37            | QueryEnum::Context(_)
38            | QueryEnum::FeedbackNaive(_) => false,
39        }
40    }
41
42    pub fn iterate_sparse(&self, mut f: impl FnMut(&VectorName, &SparseVector)) {
43        match self {
44            QueryEnum::Nearest(named) => match &named.query {
45                VectorInternal::Sparse(sparse_vector) => f(named.get_name(), sparse_vector),
46                VectorInternal::Dense(_) | VectorInternal::MultiDense(_) => {}
47            },
48            QueryEnum::RecommendBestScore(reco_query)
49            | QueryEnum::RecommendSumScores(reco_query) => {
50                let name = reco_query.get_name();
51                for vector in reco_query.query.flat_iter() {
52                    match vector {
53                        VectorInternal::Sparse(sparse_vector) => f(name, sparse_vector),
54                        VectorInternal::Dense(_) | VectorInternal::MultiDense(_) => {}
55                    }
56                }
57            }
58            QueryEnum::Discover(discover_query) => {
59                let name = discover_query.get_name();
60                for vector in discover_query.query.flat_iter() {
61                    match vector {
62                        VectorInternal::Sparse(sparse_vector) => f(name, sparse_vector),
63                        VectorInternal::Dense(_) | VectorInternal::MultiDense(_) => {}
64                    }
65                }
66            }
67            QueryEnum::Context(context_query) => {
68                let name = context_query.get_name();
69                for vector in context_query.query.flat_iter() {
70                    match vector {
71                        VectorInternal::Sparse(sparse_vector) => f(name, sparse_vector),
72                        VectorInternal::Dense(_) | VectorInternal::MultiDense(_) => {}
73                    }
74                }
75            }
76            QueryEnum::FeedbackNaive(feedback_query) => {
77                let name = feedback_query.get_name();
78                for vector in feedback_query.query.flat_iter() {
79                    match vector {
80                        VectorInternal::Sparse(sparse_vector) => f(name, sparse_vector),
81                        VectorInternal::Dense(_) | VectorInternal::MultiDense(_) => {}
82                    }
83                }
84            }
85        }
86    }
87
88    /// Returns the estimated cost of using this query in terms of number of vectors.
89    /// The cost approximates how many similarity comparisons this query will make against one point.
90    pub fn search_cost(&self) -> usize {
91        match self {
92            QueryEnum::Nearest(named_query) => search_cost([&named_query.query]),
93            QueryEnum::RecommendBestScore(named_query) => {
94                search_cost(named_query.query.flat_iter())
95            }
96            QueryEnum::RecommendSumScores(named_query) => {
97                search_cost(named_query.query.flat_iter())
98            }
99            QueryEnum::Discover(named_query) => search_cost(named_query.query.flat_iter()),
100            QueryEnum::Context(named_query) => search_cost(named_query.query.flat_iter()),
101            QueryEnum::FeedbackNaive(named_query) => search_cost(named_query.query.flat_iter()),
102        }
103    }
104}
105
106fn search_cost<'a>(vectors: impl IntoIterator<Item = &'a VectorInternal>) -> usize {
107    vectors
108        .into_iter()
109        .map(VectorInternal::similarity_cost)
110        .sum()
111}
112
113impl AsRef<QueryEnum> for QueryEnum {
114    fn as_ref(&self) -> &QueryEnum {
115        self
116    }
117}
118
119impl From<DenseVector> for QueryEnum {
120    fn from(vector: DenseVector) -> Self {
121        QueryEnum::Nearest(NamedQuery {
122            query: VectorInternal::Dense(vector),
123            using: None,
124        })
125    }
126}
127
128impl From<NamedQuery<DiscoverQuery<VectorInternal>>> for QueryEnum {
129    fn from(query: NamedQuery<DiscoverQuery<VectorInternal>>) -> Self {
130        QueryEnum::Discover(query)
131    }
132}
133
134impl From<QueryEnum> for QueryVector {
135    fn from(query: QueryEnum) -> Self {
136        match query {
137            QueryEnum::Nearest(named) => QueryVector::Nearest(named.query),
138            QueryEnum::RecommendBestScore(named) => QueryVector::RecommendBestScore(named.query),
139            QueryEnum::RecommendSumScores(named) => QueryVector::RecommendSumScores(named.query),
140            QueryEnum::Discover(named) => QueryVector::Discover(named.query),
141            QueryEnum::Context(named) => QueryVector::Context(named.query),
142            QueryEnum::FeedbackNaive(named) => QueryVector::FeedbackNaive(named.query),
143        }
144    }
145}
qdrant_edge/shard/query/query_enum.rs

qdrant_edge/shard/query/
query_enum.rs
