qdrant_edge/sparse/common/
sparse_vector.rs

1use std::borrow::Cow;
2use std::hash::Hash;
3
4use crate::blobstore::Blob;
5use crate::common::types::ScoreType;
6use itertools::Itertools;
7use ordered_float::OrderedFloat;
8use schemars::JsonSchema;
9use serde::{Deserialize, Serialize};
10use validator::{Validate, ValidationError, ValidationErrors};
11
12use crate::sparse::common::types::{DimId, DimOffset, DimWeight};
13
14/// Sparse vector structure
15#[derive(Debug, PartialEq, Clone, Default, Serialize, Deserialize, JsonSchema)]
16#[serde(rename_all = "snake_case")]
17pub struct SparseVector {
18    /// Indices must be unique
19    pub indices: Vec<DimId>,
20    /// Values and indices must be the same length
21    pub values: Vec<DimWeight>,
22}
23
24impl Hash for SparseVector {
25    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
26        let Self { indices, values } = self;
27        indices.hash(state);
28        for &value in values {
29            OrderedFloat(value).hash(state);
30        }
31    }
32}
33
34/// Same as `SparseVector` but with `DimOffset` indices.
35/// Meaning that is uses internal segment-specific indices.
36#[derive(Debug, PartialEq, Clone, Default, Serialize, Deserialize)]
37pub struct RemappedSparseVector {
38    /// indices must be unique
39    pub indices: Vec<DimOffset>,
40    /// values and indices must be the same length
41    pub values: Vec<DimWeight>,
42}
43
44/// Sort two arrays by the first array.
45pub fn double_sort<T: Ord + Copy, V: Copy>(indices: &mut [T], values: &mut [V]) {
46    // Check if the indices are already sorted
47    if indices.array_windows().all(|[a, b]| a < b) {
48        return;
49    }
50
51    let mut indexed_values: Vec<(T, V)> = indices
52        .iter()
53        .zip(values.iter())
54        .map(|(&i, &v)| (i, v))
55        .collect();
56
57    // Sort the vector of tuples by indices
58    indexed_values.sort_unstable_by_key(|&(i, _)| i);
59
60    for (i, (index, value)) in indexed_values.into_iter().enumerate() {
61        indices[i] = index;
62        values[i] = value;
63    }
64}
65
66pub fn score_vectors<T: Ord + Eq>(
67    self_indices: &[T],
68    self_values: &[DimWeight],
69    other_indices: &[T],
70    other_values: &[DimWeight],
71) -> Option<ScoreType> {
72    let mut score = 0.0;
73    // track whether there is any overlap
74    let mut overlap = false;
75    let mut i = 0;
76    let mut j = 0;
77    while i < self_indices.len() && j < other_indices.len() {
78        match self_indices[i].cmp(&other_indices[j]) {
79            std::cmp::Ordering::Less => i += 1,
80            std::cmp::Ordering::Greater => j += 1,
81            std::cmp::Ordering::Equal => {
82                overlap = true;
83                score += self_values[i] * other_values[j];
84                i += 1;
85                j += 1;
86            }
87        }
88    }
89    overlap.then_some(score)
90}
91
92impl RemappedSparseVector {
93    pub fn new(indices: Vec<DimId>, values: Vec<DimWeight>) -> Result<Self, ValidationErrors> {
94        let vector = Self { indices, values };
95        vector.validate()?;
96        Ok(vector)
97    }
98
99    pub fn sort_by_indices(&mut self) {
100        double_sort(&mut self.indices, &mut self.values);
101    }
102
103    /// Check if this vector is sorted by indices.
104    pub fn is_sorted(&self) -> bool {
105        self.indices.array_windows().all(|[a, b]| a < b)
106    }
107
108    /// Score this vector against another vector using dot product.
109    /// Warning: Expects both vectors to be sorted by indices.
110    ///
111    /// Return None if the vectors do not overlap.
112    pub fn score(&self, other: &RemappedSparseVector) -> Option<ScoreType> {
113        debug_assert!(self.is_sorted());
114        debug_assert!(other.is_sorted());
115        score_vectors(&self.indices, &self.values, &other.indices, &other.values)
116    }
117
118    /// Returns the number of elements in the vector.
119    pub fn len(&self) -> usize {
120        self.indices.len()
121    }
122
123    pub fn is_empty(&self) -> bool {
124        self.len() == 0
125    }
126}
127
128impl SparseVector {
129    pub fn new(indices: Vec<DimId>, values: Vec<DimWeight>) -> Result<Self, ValidationErrors> {
130        let vector = SparseVector { indices, values };
131        vector.validate()?;
132        Ok(vector)
133    }
134
135    #[cfg(feature = "testing")]
136    pub fn new_unchecked(indices: Vec<DimId>, values: Vec<DimWeight>) -> Self {
137        SparseVector { indices, values }
138    }
139
140    /// Sort this vector by indices.
141    ///
142    /// Sorting is required for scoring and overlap checks.
143    pub fn sort_by_indices(&mut self) {
144        double_sort(&mut self.indices, &mut self.values);
145    }
146
147    /// Check if this vector is sorted by indices.
148    pub fn is_sorted(&self) -> bool {
149        self.indices.windows(2).all(|w| w[0] < w[1])
150    }
151
152    /// Check if this vector is empty.
153    pub fn is_empty(&self) -> bool {
154        self.indices.is_empty() && self.values.is_empty()
155    }
156
157    /// Returns the number of elements in the vector.
158    pub fn len(&self) -> usize {
159        self.indices.len()
160    }
161
162    /// Score this vector against another vector using dot product.
163    /// Warning: Expects both vectors to be sorted by indices.
164    ///
165    /// Return None if the vectors do not overlap.
166    pub fn score(&self, other: &SparseVector) -> Option<ScoreType> {
167        debug_assert!(self.is_sorted());
168        debug_assert!(other.is_sorted());
169        score_vectors(&self.indices, &self.values, &other.indices, &other.values)
170    }
171
172    /// Construct a new vector that is the result of performing all indices-wise operations.
173    /// Automatically sort input vectors if necessary.
174    pub fn combine_aggregate(
175        &self,
176        other: &SparseVector,
177        op: impl Fn(DimWeight, DimWeight) -> DimWeight,
178    ) -> Self {
179        // Copy and sort `self` vector if not already sorted
180        let this: Cow<SparseVector> = if !self.is_sorted() {
181            let mut this = self.clone();
182            this.sort_by_indices();
183            Cow::Owned(this)
184        } else {
185            Cow::Borrowed(self)
186        };
187        assert!(this.is_sorted());
188
189        // Copy and sort `other` vector if not already sorted
190        let cow_other: Cow<SparseVector> = if !other.is_sorted() {
191            let mut other = other.clone();
192            other.sort_by_indices();
193            Cow::Owned(other)
194        } else {
195            Cow::Borrowed(other)
196        };
197        let other = &cow_other;
198        assert!(other.is_sorted());
199
200        let mut result = SparseVector::default();
201        let mut i = 0;
202        let mut j = 0;
203        while i < this.indices.len() && j < other.indices.len() {
204            match this.indices[i].cmp(&other.indices[j]) {
205                std::cmp::Ordering::Less => {
206                    result.indices.push(this.indices[i]);
207                    result.values.push(op(this.values[i], 0.0));
208                    i += 1;
209                }
210                std::cmp::Ordering::Greater => {
211                    result.indices.push(other.indices[j]);
212                    result.values.push(op(0.0, other.values[j]));
213                    j += 1;
214                }
215                std::cmp::Ordering::Equal => {
216                    result.indices.push(this.indices[i]);
217                    result.values.push(op(this.values[i], other.values[j]));
218                    i += 1;
219                    j += 1;
220                }
221            }
222        }
223        while i < this.indices.len() {
224            result.indices.push(this.indices[i]);
225            result.values.push(op(this.values[i], 0.0));
226            i += 1;
227        }
228        while j < other.indices.len() {
229            result.indices.push(other.indices[j]);
230            result.values.push(op(0.0, other.values[j]));
231            j += 1;
232        }
233        debug_assert!(result.is_sorted());
234        debug_assert!(result.validate().is_ok());
235        result
236    }
237
238    /// Create [RemappedSparseVector] from this vector in a naive way. Only suitable for testing.
239    #[cfg(feature = "testing")]
240    pub fn into_remapped(self) -> RemappedSparseVector {
241        RemappedSparseVector {
242            indices: self.indices,
243            values: self.values,
244        }
245    }
246}
247
248impl TryFrom<Vec<(u32, f32)>> for RemappedSparseVector {
249    type Error = ValidationErrors;
250
251    fn try_from(tuples: Vec<(u32, f32)>) -> Result<Self, Self::Error> {
252        let (indices, values): (Vec<_>, Vec<_>) = tuples.into_iter().unzip();
253        RemappedSparseVector::new(indices, values)
254    }
255}
256
257impl TryFrom<Vec<(u32, f32)>> for SparseVector {
258    type Error = ValidationErrors;
259
260    fn try_from(tuples: Vec<(u32, f32)>) -> Result<Self, Self::Error> {
261        let (indices, values): (Vec<_>, Vec<_>) = tuples.into_iter().unzip();
262        SparseVector::new(indices, values)
263    }
264}
265
266impl Blob for SparseVector {
267    fn to_bytes(&self) -> Vec<u8> {
268        bincode::serialize(&self).expect("Sparse vector serialization should not fail")
269    }
270
271    fn from_bytes(data: &[u8]) -> Self {
272        bincode::deserialize(data).expect("Sparse vector deserialization should not fail")
273    }
274}
275
276#[cfg(test)]
277impl<const N: usize> From<[(u32, f32); N]> for SparseVector {
278    fn from(value: [(u32, f32); N]) -> Self {
279        value.to_vec().try_into().unwrap()
280    }
281}
282
283#[cfg(test)]
284impl<const N: usize> From<[(u32, f32); N]> for RemappedSparseVector {
285    fn from(value: [(u32, f32); N]) -> Self {
286        value.to_vec().try_into().unwrap()
287    }
288}
289
290impl Validate for SparseVector {
291    fn validate(&self) -> Result<(), ValidationErrors> {
292        validate_sparse_vector_impl(&self.indices, &self.values)
293    }
294}
295
296impl Validate for RemappedSparseVector {
297    fn validate(&self) -> Result<(), ValidationErrors> {
298        validate_sparse_vector_impl(&self.indices, &self.values)
299    }
300}
301
302pub fn validate_sparse_vector_impl<T: Clone + Eq + Hash>(
303    indices: &[T],
304    values: &[DimWeight],
305) -> Result<(), ValidationErrors> {
306    let mut errors = ValidationErrors::default();
307
308    if indices.len() != values.len() {
309        errors.add(
310            "values",
311            ValidationError::new("must be the same length as indices"),
312        );
313    }
314    if indices.iter().unique().count() != indices.len() {
315        errors.add("indices", ValidationError::new("must be unique"));
316    }
317
318    if errors.is_empty() {
319        Ok(())
320    } else {
321        Err(errors)
322    }
323}
324
325#[cfg(test)]
326mod tests {
327    use super::*;
328
329    #[test]
330    fn test_score_aligned_same_size() {
331        let v1 = RemappedSparseVector::new(vec![1, 2, 3], vec![1.0, 2.0, 3.0]).unwrap();
332        let v2 = RemappedSparseVector::new(vec![1, 2, 3], vec![1.0, 2.0, 3.0]).unwrap();
333        assert_eq!(v1.score(&v2), Some(14.0));
334    }
335
336    #[test]
337    fn test_score_not_aligned_same_size() {
338        let v1 = RemappedSparseVector::new(vec![1, 2, 3], vec![1.0, 2.0, 3.0]).unwrap();
339        let v2 = RemappedSparseVector::new(vec![2, 3, 4], vec![2.0, 3.0, 4.0]).unwrap();
340        assert_eq!(v1.score(&v2), Some(13.0));
341    }
342
343    #[test]
344    fn test_score_aligned_different_size() {
345        let v1 = RemappedSparseVector::new(vec![1, 2, 3], vec![1.0, 2.0, 3.0]).unwrap();
346        let v2 = RemappedSparseVector::new(vec![1, 2, 3, 4], vec![1.0, 2.0, 3.0, 4.0]).unwrap();
347        assert_eq!(v1.score(&v2), Some(14.0));
348    }
349
350    #[test]
351    fn test_score_not_aligned_different_size() {
352        let v1 = RemappedSparseVector::new(vec![1, 2, 3], vec![1.0, 2.0, 3.0]).unwrap();
353        let v2 = RemappedSparseVector::new(vec![2, 3, 4, 5], vec![2.0, 3.0, 4.0, 5.0]).unwrap();
354        assert_eq!(v1.score(&v2), Some(13.0));
355    }
356
357    #[test]
358    fn test_score_no_overlap() {
359        let v1 = RemappedSparseVector::new(vec![1, 2, 3], vec![1.0, 2.0, 3.0]).unwrap();
360        let v2 = RemappedSparseVector::new(vec![4, 5, 6], vec![2.0, 3.0, 4.0]).unwrap();
361        assert!(v1.score(&v2).is_none());
362    }
363
364    #[test]
365    fn validation_test() {
366        let fully_empty = SparseVector::new(vec![], vec![]);
367        assert!(fully_empty.is_ok());
368        assert!(fully_empty.unwrap().is_empty());
369
370        let different_length = SparseVector::new(vec![1, 2, 3], vec![1.0, 2.0]);
371        assert!(different_length.is_err());
372
373        let not_sorted = SparseVector::new(vec![1, 3, 2], vec![1.0, 2.0, 3.0]);
374        assert!(not_sorted.is_ok());
375
376        let not_unique = SparseVector::new(vec![1, 2, 3, 2], vec![1.0, 2.0, 3.0, 4.0]);
377        assert!(not_unique.is_err());
378    }
379
380    #[test]
381    fn sorting_test() {
382        let mut not_sorted = SparseVector::new(vec![1, 3, 2], vec![1.0, 2.0, 3.0]).unwrap();
383        assert!(!not_sorted.is_sorted());
384        not_sorted.sort_by_indices();
385        assert!(not_sorted.is_sorted());
386    }
387
388    #[test]
389    fn combine_aggregate_test() {
390        // Test with missing index
391        let a = SparseVector::new(vec![1, 2, 3], vec![0.1, 0.2, 0.3]).unwrap();
392        let b = SparseVector::new(vec![2, 3, 4], vec![2.0, 3.0, 4.0]).unwrap();
393        let sum = a.combine_aggregate(&b, |x, y| x + 2.0 * y);
394        assert_eq!(sum.indices, vec![1, 2, 3, 4]);
395        assert_eq!(sum.values, vec![0.1, 4.2, 6.3, 8.0]);
396
397        // reverse arguments
398        let sum = b.combine_aggregate(&a, |x, y| x + 2.0 * y);
399        assert_eq!(sum.indices, vec![1, 2, 3, 4]);
400        assert_eq!(sum.values, vec![0.2, 2.4, 3.6, 4.0]);
401
402        // Test with non-sorted input
403        let a = SparseVector::new(vec![1, 2, 3], vec![0.1, 0.2, 0.3]).unwrap();
404        let b = SparseVector::new(vec![4, 2, 3], vec![4.0, 2.0, 3.0]).unwrap();
405        let sum = a.combine_aggregate(&b, |x, y| x + 2.0 * y);
406        assert_eq!(sum.indices, vec![1, 2, 3, 4]);
407        assert_eq!(sum.values, vec![0.1, 4.2, 6.3, 8.0]);
408    }
409}
qdrant_edge/sparse/common/sparse_vector.rs

qdrant_edge/sparse/common/
sparse_vector.rs
