qdrant_edge/segment/vector_storage/query/
context_query.rs

1use std::hash::Hash;
2use std::iter::{self, Chain, Once};
3
4use crate::common::math::fast_sigmoid;
5use crate::common::types::ScoreType;
6use itertools::Itertools;
7use serde::Serialize;
8
9use super::{Query, TransformInto};
10use crate::segment::common::operation_error::OperationResult;
11use crate::segment::data_types::vectors::{QueryVector, VectorInternal};
12
13#[derive(Debug, Clone, PartialEq, Serialize, Hash)]
14pub struct ContextPair<T> {
15    pub positive: T,
16    pub negative: T,
17}
18
19impl<T> ContextPair<T> {
20    pub fn iter(&self) -> impl Iterator<Item = &T> {
21        iter::once(&self.positive).chain(iter::once(&self.negative))
22    }
23
24    pub fn transform<U>(
25        self,
26        f: &dyn Fn(T) -> OperationResult<U>,
27    ) -> OperationResult<ContextPair<U>> {
28        Ok(ContextPair {
29            positive: f(self.positive)?,
30            negative: f(self.negative)?,
31        })
32    }
33
34    /// In the first stage of discovery search, the objective is to get the best entry point
35    /// for the search. This is done by using a smooth loss function instead of hard ranking
36    /// to approach the best zone, once the best zone is reached, score will be same for all
37    /// points inside that zone.
38    /// e.g.:
39    /// ```text
40    ///                   │
41    ///                   │
42    ///                   │    +0
43    ///                   │             +0
44    ///                   │
45    ///         n         │         p
46    ///                   │
47    ///   ─►          ─►  │
48    ///  -0.4        -0.1 │   +0
49    ///                   │
50    /// ```
51    /// Simple 2D model:
52    /// <https://www.desmos.com/calculator/lbxycyh2hs>
53    pub fn loss_by(&self, similarity: impl Fn(&T) -> ScoreType) -> ScoreType {
54        const MARGIN: ScoreType = ScoreType::EPSILON;
55
56        let positive = similarity(&self.positive);
57        let negative = similarity(&self.negative);
58
59        let difference = positive - negative - MARGIN;
60
61        fast_sigmoid(ScoreType::min(difference, 0.0))
62    }
63}
64
65impl<T> IntoIterator for ContextPair<T> {
66    type Item = T;
67
68    type IntoIter = Chain<Once<T>, Once<T>>;
69
70    fn into_iter(self) -> Self::IntoIter {
71        iter::once(self.positive).chain(iter::once(self.negative))
72    }
73}
74
75#[cfg(test)]
76impl<T> From<(T, T)> for ContextPair<T> {
77    fn from(pair: (T, T)) -> Self {
78        Self {
79            positive: pair.0,
80            negative: pair.1,
81        }
82    }
83}
84
85#[derive(Debug, Clone, PartialEq, Serialize, Hash)]
86pub struct ContextQuery<T> {
87    pub pairs: Vec<ContextPair<T>>,
88}
89
90impl<T> ContextQuery<T> {
91    pub fn new(pairs: Vec<ContextPair<T>>) -> Self {
92        Self { pairs }
93    }
94
95    pub fn flat_iter(&self) -> impl Iterator<Item = &T> {
96        self.pairs.iter().flat_map(|pair| pair.iter())
97    }
98}
99
100impl<T, U> TransformInto<ContextQuery<U>, T, U> for ContextQuery<T> {
101    fn transform(self, f: &dyn Fn(T) -> OperationResult<U>) -> OperationResult<ContextQuery<U>> {
102        Ok(ContextQuery::new(
103            self.pairs
104                .into_iter()
105                .map(|pair| pair.transform(f))
106                .try_collect()?,
107        ))
108    }
109}
110
111impl<T> Query<T> for ContextQuery<T> {
112    fn score_by(&self, similarity: impl Fn(&T) -> ScoreType) -> ScoreType {
113        let mut sum = 0.0;
114        for pair in &self.pairs {
115            sum += pair.loss_by(&similarity);
116        }
117        sum
118    }
119}
120
121impl<T> From<Vec<ContextPair<T>>> for ContextQuery<T> {
122    fn from(pairs: Vec<ContextPair<T>>) -> Self {
123        ContextQuery::new(pairs)
124    }
125}
126
127impl From<ContextQuery<VectorInternal>> for QueryVector {
128    fn from(query: ContextQuery<VectorInternal>) -> Self {
129        QueryVector::Context(query)
130    }
131}
132
133#[cfg(test)]
134mod test {
135    use crate::common::types::ScoreType;
136    use proptest::prelude::*;
137
138    use super::*;
139
140    fn dummy_similarity(x: &f32) -> ScoreType {
141        *x as ScoreType
142    }
143
144    /// Possible similarities
145    fn sim() -> impl Strategy<Value = f32> {
146        (-100.0..=100.0).prop_map(|x| x as f32)
147    }
148
149    proptest! {
150        #![proptest_config(ProptestConfig::with_cases(1000))]
151
152        /// Checks that the loss is between 0 and -1
153        #[test]
154        fn loss_is_not_more_than_1_per_pair((p, n) in (sim(), sim())) {
155            let query = ContextQuery::new(vec![ContextPair::from((p, n))]);
156
157            let score = query.score_by(dummy_similarity);
158            assert!(score <= 0.0, "similarity: {score}");
159            assert!(score > -1.0, "similarity: {score}");
160        }
161    }
162}
qdrant_edge/segment/vector_storage/query/context_query.rs

qdrant_edge/segment/vector_storage/query/
context_query.rs
