qdrant_edge/segment/vector_storage/query/
reco_query.rs

1use std::hash::Hash;
2
3use crate::common::math::scaled_fast_sigmoid;
4use crate::common::types::ScoreType;
5use itertools::Itertools;
6use serde::Serialize;
7
8use super::{Query, TransformInto};
9use crate::segment::common::operation_error::OperationResult;
10use crate::segment::data_types::vectors::{QueryVector, VectorInternal};
11
12#[derive(Debug, Clone, PartialEq, Serialize, Hash)]
13pub struct RecoQuery<T> {
14    pub positives: Vec<T>,
15    pub negatives: Vec<T>,
16}
17
18impl<T> RecoQuery<T> {
19    pub fn new(positives: Vec<T>, negatives: Vec<T>) -> Self {
20        Self {
21            positives,
22            negatives,
23        }
24    }
25
26    pub fn flat_iter(&self) -> impl Iterator<Item = &T> {
27        self.positives.iter().chain(self.negatives.iter())
28    }
29}
30
31impl<T, U> TransformInto<RecoQuery<U>, T, U> for RecoQuery<T> {
32    fn transform(self, f: &dyn Fn(T) -> OperationResult<U>) -> OperationResult<RecoQuery<U>> {
33        let positives = self.positives.into_iter().map(f).try_collect()?;
34        let negatives = self.negatives.into_iter().map(f).try_collect()?;
35        Ok(RecoQuery::new(positives, negatives))
36    }
37}
38
39#[derive(Debug, Clone, PartialEq)]
40pub struct RecoBestScoreQuery<T>(RecoQuery<T>);
41
42impl<T> From<RecoQuery<T>> for RecoBestScoreQuery<T> {
43    fn from(query: RecoQuery<T>) -> Self {
44        Self(query)
45    }
46}
47
48impl<T, U> TransformInto<RecoBestScoreQuery<U>, T, U> for RecoBestScoreQuery<T> {
49    fn transform(
50        self,
51        f: &dyn Fn(T) -> OperationResult<U>,
52    ) -> OperationResult<RecoBestScoreQuery<U>> {
53        Ok(RecoBestScoreQuery(self.0.transform(f)?))
54    }
55}
56
57impl From<RecoBestScoreQuery<VectorInternal>> for QueryVector {
58    fn from(query: RecoBestScoreQuery<VectorInternal>) -> Self {
59        QueryVector::RecommendBestScore(query.0)
60    }
61}
62
63impl<T> Query<T> for RecoBestScoreQuery<T> {
64    fn score_by(&self, similarity: impl Fn(&T) -> ScoreType) -> ScoreType {
65        // get similarities to all positives
66        let mut max_positive = ScoreType::NEG_INFINITY;
67        for vector in &self.0.positives {
68            let score = similarity(vector);
69            if score.total_cmp(&max_positive).is_gt() {
70                max_positive = score;
71            }
72        }
73
74        // and all negatives
75        let mut max_negative = ScoreType::NEG_INFINITY;
76        for vector in &self.0.negatives {
77            let score = similarity(vector);
78            if score.total_cmp(&max_negative).is_gt() {
79                max_negative = score;
80            }
81        }
82
83        if max_positive > max_negative {
84            scaled_fast_sigmoid(max_positive)
85        } else {
86            -scaled_fast_sigmoid(max_negative)
87        }
88    }
89}
90
91#[derive(Debug, Clone, PartialEq)]
92pub struct RecoSumScoresQuery<T>(RecoQuery<T>);
93
94impl<T> From<RecoQuery<T>> for RecoSumScoresQuery<T> {
95    fn from(query: RecoQuery<T>) -> Self {
96        Self(query)
97    }
98}
99
100impl<T, U> TransformInto<RecoSumScoresQuery<U>, T, U> for RecoSumScoresQuery<T> {
101    fn transform(
102        self,
103        f: &dyn Fn(T) -> OperationResult<U>,
104    ) -> OperationResult<RecoSumScoresQuery<U>> {
105        Ok(RecoSumScoresQuery(self.0.transform(f)?))
106    }
107}
108
109impl From<RecoSumScoresQuery<VectorInternal>> for QueryVector {
110    fn from(query: RecoSumScoresQuery<VectorInternal>) -> Self {
111        QueryVector::RecommendSumScores(query.0)
112    }
113}
114
115impl<T> Query<T> for RecoSumScoresQuery<T> {
116    fn score_by(&self, similarity: impl Fn(&T) -> ScoreType) -> ScoreType {
117        // Sum all positive vectors scores
118        let mut positive_score: ScoreType = 0.0;
119        for vector in &self.0.positives {
120            positive_score += similarity(vector);
121        }
122
123        // Sum all negative vectors scores
124        let mut negative_score: ScoreType = 0.0;
125        for vector in &self.0.negatives {
126            negative_score += similarity(vector);
127        }
128
129        // Subtract
130        positive_score - negative_score
131    }
132}
133
134#[cfg(test)]
135mod test {
136    use std::cmp::Ordering;
137
138    use crate::common::math::scaled_fast_sigmoid;
139    use crate::common::types::ScoreType;
140    use proptest::prelude::*;
141    use rstest::rstest;
142
143    use crate::segment::vector_storage::query::{Query, RecoBestScoreQuery, RecoQuery};
144
145    enum Chosen {
146        Positive,
147        Negative,
148    }
149
150    #[rstest]
151    #[case::higher_positive(vec![42], vec![4], Chosen::Positive, 42.0)]
152    #[case::higher_negative(vec![4], vec![42], Chosen::Negative, 42.0)]
153    #[case::negative_zero(vec![-1], vec![0], Chosen::Negative, 0.0)]
154    #[case::positive_zero(vec![0], vec![-1], Chosen::Positive, 0.0)]
155    #[case::both_under_zero(vec![-42], vec![-84], Chosen::Positive, -42.0)]
156    #[case::both_under_zero_but_negative_is_higher(vec![-84], vec![-42], Chosen::Negative, -42.0)]
157    #[case::multiple_with_negative_best(vec![1, 2, 3], vec![4, 5, 6], Chosen::Negative, 6.0)]
158    #[case::multiple_with_positive_best(vec![10, 2, 3], vec![4, 5, 6], Chosen::Positive, 10.0)]
159    fn score_query(
160        #[case] positives: Vec<isize>,
161        #[case] negatives: Vec<isize>,
162        #[case] chosen: Chosen,
163        #[case] expected: ScoreType,
164    ) {
165        use super::{RecoBestScoreQuery, RecoQuery};
166
167        let query = RecoBestScoreQuery::from(RecoQuery::new(positives, negatives));
168
169        let dummy_similarity = |x: &isize| *x as ScoreType;
170
171        let positive_transformation = scaled_fast_sigmoid;
172        let negative_transformation = |x| -scaled_fast_sigmoid(x);
173
174        let score = query.score_by(dummy_similarity);
175
176        match chosen {
177            Chosen::Positive => {
178                assert_eq!(score, positive_transformation(expected));
179            }
180            Chosen::Negative => {
181                assert_eq!(score, negative_transformation(expected));
182            }
183        }
184    }
185
186    fn ulps_eq(a: f32, b: f32, ulps: u32) -> bool {
187        if a.signum() != b.signum() {
188            return false;
189        }
190
191        let a = a.to_bits();
192        let b = b.to_bits();
193
194        a.abs_diff(b) <= ulps
195    }
196
197    /// Relaxes the comparison of floats to allow for a some difference in units of least precision
198    fn float_cmp(a: f32, b: f32) -> Ordering {
199        if ulps_eq(a, b, 80) {
200            Ordering::Equal
201        } else {
202            a.total_cmp(&b)
203        }
204    }
205
206    proptest! {
207        /// Checks that the negative-chosen scores invert the order of the candidates
208        #[test]
209        fn correct_negative_order(a in -100f32..=100f32, b in -100f32..=100f32) {
210            let dummy_similarity = |x: &f32| *x as ScoreType;
211
212            let ordering_before = float_cmp(dummy_similarity(&a), dummy_similarity(&b));
213
214            let query_a = RecoBestScoreQuery::from(RecoQuery::new(vec![], vec![a]));
215            let query_b = RecoBestScoreQuery::from(RecoQuery::new(vec![], vec![b]));
216
217            let score_a = query_a.score_by(dummy_similarity);
218            let score_b = query_b.score_by(dummy_similarity);
219
220            let ordering_after = float_cmp(score_a, score_b);
221
222            if ordering_before == std::cmp::Ordering::Equal {
223                assert_eq!(ordering_before, ordering_after);
224            } else {
225                assert_ne!(ordering_before, ordering_after)
226            }
227        }
228
229        /// Checks that the positive-chosen scores preserve the order of the candidates
230        #[test]
231        fn correct_positive_order(a in -100f32..=100f32, b in -100f32..=100f32) {
232            let dummy_similarity = |x: &f32| *x as ScoreType;
233
234            let ordering_before = float_cmp(dummy_similarity(&a), dummy_similarity(&b));
235
236            // Too similar scores can get compressed to the same value by the sigmoid function.
237            // This would make the test useless, so we skip those cases.
238            prop_assume!(ordering_before != Ordering::Equal);
239
240            let query_a = RecoBestScoreQuery::from(RecoQuery::new(vec![a], vec![]));
241            let query_b = RecoBestScoreQuery::from(RecoQuery::new(vec![b], vec![]));
242
243            let score_a = query_a.score_by(dummy_similarity);
244            let score_b = query_b.score_by(dummy_similarity);
245
246            let ordering_after = score_a.total_cmp(&score_b);
247
248            assert_eq!(ordering_before, ordering_after);
249        }
250
251        /// Guarantees that the point that was chosen from positive is always preferred on
252        /// the candidate list over a point that was chosen from negatives
253        #[test]
254        fn correct_positive_and_negative_order(p in -100f32..=100f32, n in -100f32..=100f32) {
255            let dummy_similarity = |x: &f32| *x as ScoreType;
256
257            let query_p = RecoBestScoreQuery::from(RecoQuery::new(vec![p], vec![]));
258            let query_n = RecoBestScoreQuery::from(RecoQuery::new(vec![], vec![n]));
259
260            let ordering = query_p.score_by(dummy_similarity).total_cmp(&query_n.score_by(dummy_similarity));
261
262            assert_ne!(ordering, std::cmp::Ordering::Less);
263        }
264    }
265}
qdrant_edge/segment/vector_storage/query/reco_query.rs

qdrant_edge/segment/vector_storage/query/
reco_query.rs
