hermes_core/query/vector/
dense.rs

1//! Dense vector query for similarity search (ANN)
2
3use crate::dsl::Field;
4use crate::segment::SegmentReader;
5use std::sync::{Arc, Mutex};
6
7use super::VectorResultScorer;
8use super::combiner::MultiValueCombiner;
9use crate::query::traits::{CountFuture, Query, Scorer, ScorerFuture};
10
11/// Maximum number of IVF clusters a single dense query may probe.
12///
13/// This guard bounds explicit overrides independently of the trained global
14/// leaf count. Automatic billion-scale fields normally use far fewer probes.
15pub const MAX_DENSE_NPROBE: usize = 65_536;
16
17/// Maximum exact-rerank candidate multiplier accepted by dense search.
18///
19/// Normal deployments use 2-5x. 32x leaves ample room for recall tuning while
20/// bounding the heap and raw-vector buffers derived from the requested top-k.
21pub const MAX_DENSE_RERANK_FACTOR: f32 = 32.0;
22
23/// Dense vector query for similarity search
24#[derive(Debug, Clone)]
25pub struct DenseVectorQuery {
26    /// Field containing the dense vectors
27    pub field: Field,
28    /// Query vector
29    pub vector: Vec<f32>,
30    /// Number of clusters to probe (for IVF indexes)
31    pub nprobe: usize,
32    /// Re-ranking factor (multiplied by k for candidate selection, e.g. 3.0)
33    pub rerank_factor: f32,
34    /// How to combine scores for multi-valued documents
35    pub combiner: MultiValueCombiner,
36    /// One global-quantizer route, shared by all segment scorers spawned for
37    /// this query. The cache is versioned, so a query reused after an index
38    /// generation change recomputes safely.
39    probe_cache: Arc<Mutex<Option<Arc<crate::structures::IvfPqQueryPlan>>>>,
40}
41
42impl std::fmt::Display for DenseVectorQuery {
43    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44        write!(
45            f,
46            "Dense({}, dim={}, nprobe={}, rerank={})",
47            self.field.0,
48            self.vector.len(),
49            self.nprobe,
50            self.rerank_factor
51        )
52    }
53}
54
55impl DenseVectorQuery {
56    /// Create a new dense vector query
57    pub fn new(field: Field, vector: Vec<f32>) -> Self {
58        Self {
59            field,
60            vector,
61            nprobe: 64,
62            rerank_factor: 3.0,
63            combiner: MultiValueCombiner::Max,
64            probe_cache: Arc::new(Mutex::new(None)),
65        }
66    }
67
68    /// Set the number of clusters to probe (for IVF indexes)
69    ///
70    /// Values are validated when the query is executed. See
71    /// [`MAX_DENSE_NPROBE`].
72    pub fn with_nprobe(mut self, nprobe: usize) -> Self {
73        self.nprobe = nprobe;
74        self
75    }
76
77    /// Set the re-ranking factor (e.g. 3.0 = fetch 3x candidates for reranking)
78    ///
79    /// Values are validated when the query is executed. See
80    /// [`MAX_DENSE_RERANK_FACTOR`].
81    pub fn with_rerank_factor(mut self, factor: f32) -> Self {
82        self.rerank_factor = factor;
83        self
84    }
85
86    /// Set the multi-value score combiner
87    pub fn with_combiner(mut self, combiner: MultiValueCombiner) -> Self {
88        self.combiner = combiner;
89        self
90    }
91}
92
93impl Query for DenseVectorQuery {
94    fn scorer<'a>(&self, reader: &'a SegmentReader, limit: usize) -> ScorerFuture<'a> {
95        let field = self.field;
96        let vector = self.vector.clone();
97        let nprobe = self.nprobe;
98        let rerank_factor = self.rerank_factor;
99        let combiner = self.combiner;
100        let probe_cache = Arc::clone(&self.probe_cache);
101        Box::pin(async move {
102            let results = reader
103                .search_dense_vector_with_probe_cache(
104                    field,
105                    &vector,
106                    limit,
107                    nprobe,
108                    rerank_factor,
109                    combiner,
110                    &probe_cache,
111                )
112                .await?;
113
114            Ok(Box::new(VectorResultScorer::new(results, field.0)) as Box<dyn Scorer>)
115        })
116    }
117
118    #[cfg(feature = "sync")]
119    fn scorer_sync<'a>(
120        &self,
121        reader: &'a SegmentReader,
122        limit: usize,
123    ) -> crate::Result<Box<dyn Scorer + 'a>> {
124        let results = reader.search_dense_vector_sync_with_probe_cache(
125            self.field,
126            &self.vector,
127            limit,
128            self.nprobe,
129            self.rerank_factor,
130            self.combiner,
131            &self.probe_cache,
132        )?;
133        Ok(Box::new(VectorResultScorer::new(results, self.field.0)) as Box<dyn Scorer>)
134    }
135
136    fn count_estimate<'a>(&self, _reader: &'a SegmentReader) -> CountFuture<'a> {
137        Box::pin(async move { Ok(u32::MAX) })
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    #[test]
146    fn test_dense_vector_query_builder() {
147        let query = DenseVectorQuery::new(Field(0), vec![1.0, 2.0, 3.0])
148            .with_nprobe(64)
149            .with_rerank_factor(5.0);
150
151        assert_eq!(query.field, Field(0));
152        assert_eq!(query.vector.len(), 3);
153        assert_eq!(query.nprobe, 64);
154        assert_eq!(query.rerank_factor, 5.0);
155    }
156}
hermes_core/query/vector/dense.rs

hermes_core/query/vector/
dense.rs
