rullst-orm 12.1.2

Fast, un-opinionated, database-agnostic ORM and query builder for Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

use std::{cmp::Ordering, collections::BTreeMap};

use tokio::sync::RwLock;

use super::{
    VectorCollectionName, VectorDimensions, VectorMatch, VectorPoint, VectorQueryLimit,
    validate_payload, validate_vector,
};
use crate::polyglot::PolyglotError;

#[derive(Default)]
pub(super) struct MockQdrant {
    collections: RwLock<BTreeMap<String, MockCollection>>,
}

struct MockCollection {
    dimensions: VectorDimensions,
    points: BTreeMap<u64, VectorPoint>,
}

impl MockQdrant {
    pub(super) async fn create_collection(
        &self,
        collection: &VectorCollectionName,
        dimensions: VectorDimensions,
    ) -> Result<(), PolyglotError> {
        let mut collections = self.collections.write().await;
        if collections.contains_key(collection.as_str()) {
            return Err(PolyglotError::Conflict);
        }
        collections.insert(
            collection.as_str().to_owned(),
            MockCollection {
                dimensions,
                points: BTreeMap::new(),
            },
        );
        Ok(())
    }

    pub(super) async fn upsert(
        &self,
        collection: &VectorCollectionName,
        point: VectorPoint,
    ) -> Result<(), PolyglotError> {
        let mut collections = self.collections.write().await;
        let target = collections
            .get_mut(collection.as_str())
            .ok_or(PolyglotError::NotFound)?;
        validate_vector(point.vector(), Some(target.dimensions))?;
        target.points.insert(point.id(), point);
        Ok(())
    }

    pub(super) async fn delete(
        &self,
        collection: &VectorCollectionName,
        id: u64,
    ) -> Result<(), PolyglotError> {
        let mut collections = self.collections.write().await;
        let target = collections
            .get_mut(collection.as_str())
            .ok_or(PolyglotError::NotFound)?;
        target.points.remove(&id);
        Ok(())
    }

    pub(super) async fn search(
        &self,
        collection: &VectorCollectionName,
        query: &[f32],
        limit: VectorQueryLimit,
    ) -> Result<Vec<VectorMatch>, PolyglotError> {
        let collections = self.collections.read().await;
        let target = collections
            .get(collection.as_str())
            .ok_or(PolyglotError::NotFound)?;
        validate_vector(query, Some(target.dimensions))?;
        let mut matches = target
            .points
            .values()
            .map(|point| {
                VectorMatch::new(
                    point.id(),
                    cosine_similarity(query, point.vector()),
                    point.payload().clone(),
                )
            })
            .collect::<Vec<_>>();
        matches.sort_by(|left, right| {
            right
                .score()
                .partial_cmp(&left.score())
                .unwrap_or(Ordering::Equal)
                .then_with(|| left.id().cmp(&right.id()))
        });
        matches.truncate(usize::from(limit.get()));
        for result in &matches {
            validate_payload(result.payload())?;
        }
        Ok(matches)
    }
}

fn cosine_similarity(left: &[f32], right: &[f32]) -> f32 {
    let dot = left
        .iter()
        .zip(right)
        .map(|(left, right)| f64::from(*left) * f64::from(*right))
        .sum::<f64>();
    let left_norm = left
        .iter()
        .map(|value| f64::from(*value).powi(2))
        .sum::<f64>()
        .sqrt();
    let right_norm = right
        .iter()
        .map(|value| f64::from(*value).powi(2))
        .sum::<f64>()
        .sqrt();
    (dot / (left_norm * right_norm)) as f32
}