qdrant_edge/edge/read_view/ops/
matrix.rs

1use ahash::AHashSet;
2use crate::segment::common::operation_error::{OperationError, OperationResult};
3use crate::segment::data_types::vectors::NamedQuery;
4use crate::segment::types::{
5    Condition, Filter, HasIdCondition, HasVectorCondition, PointIdType, ScoredPoint,
6    WithPayloadInterface, WithVector,
7};
8use crate::shard::query::query_enum::QueryEnum;
9use crate::shard::query::{SampleInternal, ScoringQuery, ShardQueryRequest};
10
11use crate::edge::read_view::{EdgeReadView, ReadSegmentHandle};
12use crate::edge::requests::SearchMatrixRequest;
13
14#[derive(Debug, Default)]
15pub struct SearchMatrixResponse {
16    pub sample_ids: Vec<PointIdType>,
17    pub nearests: Vec<Vec<ScoredPoint>>,
18}
19
20impl<H: ReadSegmentHandle> EdgeReadView<H> {
21    pub(crate) fn search_matrix(
22        &self,
23        request: SearchMatrixRequest,
24    ) -> OperationResult<SearchMatrixResponse> {
25        let SearchMatrixRequest {
26            sample_size,
27            limit_per_sample,
28            filter,
29            using,
30        } = request;
31        if sample_size == 0 || limit_per_sample == 0 {
32            return Ok(SearchMatrixResponse::default());
33        }
34
35        // Only sample points that actually carry the `using` vector.
36        let has_vector = Filter::new_must(Condition::HasVector(HasVectorCondition::from(
37            using.clone(),
38        )));
39        let sample_filter = Some(filter.map(|f| f.merge(&has_vector)).unwrap_or(has_vector));
40
41        let sampling = ShardQueryRequest {
42            prefetches: vec![],
43            query: Some(ScoringQuery::Sample(SampleInternal::Random)),
44            filter: sample_filter,
45            score_threshold: None,
46            limit: sample_size,
47            offset: 0,
48            params: None,
49            with_vector: WithVector::Selector(vec![using.clone()]),
50            with_payload: WithPayloadInterface::Bool(false),
51        };
52        let mut sampled = self.query(sampling)?;
53        if sampled.len() < 2 {
54            return Ok(SearchMatrixResponse::default());
55        }
56        sampled.truncate(sample_size);
57        sampled.sort_unstable_by_key(|p| p.id);
58        let sample_ids: Vec<PointIdType> = sampled.iter().map(|p| p.id).collect();
59
60        // Restrict each nearest search to the sampled set.
61        let id_filter = Filter::new_must(Condition::HasId(HasIdCondition::from(
62            sample_ids.iter().copied().collect::<AHashSet<_>>(),
63        )));
64
65        let mut nearests = Vec::with_capacity(sampled.len());
66        for point in &sampled {
67            let vector = point
68                .vector
69                .as_ref()
70                .and_then(|v| v.get(&using))
71                .map(|v| v.to_owned())
72                .ok_or_else(|| {
73                    OperationError::service_error("sampled point is missing its vector")
74                })?;
75            let nearest = ShardQueryRequest {
76                prefetches: vec![],
77                query: Some(ScoringQuery::Vector(QueryEnum::Nearest(NamedQuery::new(
78                    vector,
79                    using.clone(),
80                )))),
81                filter: Some(id_filter.clone()),
82                score_threshold: None,
83                limit: limit_per_sample.saturating_add(1), // +1 to drop the point itself afterwards
84                offset: 0,
85                params: None,
86                with_vector: WithVector::Bool(false),
87                with_payload: WithPayloadInterface::Bool(false),
88            };
89            let mut scores = self.query(nearest)?;
90            if let Some(pos) = scores.iter().position(|p| p.id == point.id) {
91                scores.remove(pos);
92            } else if scores.len() == limit_per_sample.saturating_add(1) {
93                scores.pop();
94            }
95            nearests.push(scores);
96        }
97
98        Ok(SearchMatrixResponse {
99            sample_ids,
100            nearests,
101        })
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108    use crate::edge::EdgeShardRead;
109    use crate::edge::test_helpers::{VECTOR_NAME, point, test_config, upsert};
110
111    #[test]
112    fn matrix_returns_sample_and_nearests() {
113        let dir = tempfile::tempdir().unwrap();
114        let shard = crate::edge::EdgeShard::new(dir.path(), test_config()).unwrap();
115        upsert(&shard, (1..=4).map(point).collect());
116
117        let resp = shard
118            .search_matrix(SearchMatrixRequest {
119                sample_size: 4,
120                limit_per_sample: 2,
121                filter: None,
122                using: VECTOR_NAME.to_string(),
123            })
124            .unwrap();
125
126        assert_eq!(resp.sample_ids.len(), 4);
127        assert_eq!(resp.nearests.len(), 4);
128        for (id, row) in resp.sample_ids.iter().zip(&resp.nearests) {
129            assert!(row.iter().all(|p| &p.id != id));
130            assert!(row.len() <= 2);
131        }
132    }
133
134    #[test]
135    fn matrix_empty_when_fewer_than_two_points() {
136        let dir = tempfile::tempdir().unwrap();
137        let shard = crate::edge::EdgeShard::new(dir.path(), test_config()).unwrap();
138        upsert(&shard, vec![point(1)]);
139
140        let resp = shard
141            .search_matrix(SearchMatrixRequest {
142                sample_size: 10,
143                limit_per_sample: 3,
144                filter: None,
145                using: VECTOR_NAME.to_string(),
146            })
147            .unwrap();
148
149        assert!(resp.sample_ids.is_empty());
150    }
151}
qdrant_edge/edge/read_view/ops/matrix.rs

qdrant_edge/edge/read_view/ops/
matrix.rs
