hermes_core/query/vector/
dense.rs

1//! Dense vector query for similarity search (ANN)
2
3use crate::dsl::Field;
4use crate::segment::SegmentReader;
5
6use super::VectorResultScorer;
7use super::combiner::MultiValueCombiner;
8use crate::query::traits::{CountFuture, Query, Scorer, ScorerFuture};
9
10/// Maximum number of IVF clusters a single dense query may probe.
11///
12/// Automatic IVF construction caps the cluster count at 4096. Keeping the
13/// query override within the same bound prevents malformed requests from
14/// turning into unbounded work while still permitting a full-cluster scan.
15pub const MAX_DENSE_NPROBE: usize = 4096;
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}
37
38impl std::fmt::Display for DenseVectorQuery {
39    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40        write!(
41            f,
42            "Dense({}, dim={}, nprobe={}, rerank={})",
43            self.field.0,
44            self.vector.len(),
45            self.nprobe,
46            self.rerank_factor
47        )
48    }
49}
50
51impl DenseVectorQuery {
52    /// Create a new dense vector query
53    pub fn new(field: Field, vector: Vec<f32>) -> Self {
54        Self {
55            field,
56            vector,
57            nprobe: 32,
58            rerank_factor: 3.0,
59            combiner: MultiValueCombiner::Max,
60        }
61    }
62
63    /// Set the number of clusters to probe (for IVF indexes)
64    ///
65    /// Values are validated when the query is executed. See
66    /// [`MAX_DENSE_NPROBE`].
67    pub fn with_nprobe(mut self, nprobe: usize) -> Self {
68        self.nprobe = nprobe;
69        self
70    }
71
72    /// Set the re-ranking factor (e.g. 3.0 = fetch 3x candidates for reranking)
73    ///
74    /// Values are validated when the query is executed. See
75    /// [`MAX_DENSE_RERANK_FACTOR`].
76    pub fn with_rerank_factor(mut self, factor: f32) -> Self {
77        self.rerank_factor = factor;
78        self
79    }
80
81    /// Set the multi-value score combiner
82    pub fn with_combiner(mut self, combiner: MultiValueCombiner) -> Self {
83        self.combiner = combiner;
84        self
85    }
86}
87
88impl Query for DenseVectorQuery {
89    fn scorer<'a>(&self, reader: &'a SegmentReader, limit: usize) -> ScorerFuture<'a> {
90        let field = self.field;
91        let vector = self.vector.clone();
92        let nprobe = self.nprobe;
93        let rerank_factor = self.rerank_factor;
94        let combiner = self.combiner;
95        Box::pin(async move {
96            let results = reader
97                .search_dense_vector(field, &vector, limit, nprobe, rerank_factor, combiner)
98                .await?;
99
100            Ok(Box::new(VectorResultScorer::new(results, field.0)) as Box<dyn Scorer>)
101        })
102    }
103
104    #[cfg(feature = "sync")]
105    fn scorer_sync<'a>(
106        &self,
107        reader: &'a SegmentReader,
108        limit: usize,
109    ) -> crate::Result<Box<dyn Scorer + 'a>> {
110        let results = reader.search_dense_vector_sync(
111            self.field,
112            &self.vector,
113            limit,
114            self.nprobe,
115            self.rerank_factor,
116            self.combiner,
117        )?;
118        Ok(Box::new(VectorResultScorer::new(results, self.field.0)) as Box<dyn Scorer>)
119    }
120
121    fn count_estimate<'a>(&self, _reader: &'a SegmentReader) -> CountFuture<'a> {
122        Box::pin(async move { Ok(u32::MAX) })
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    #[test]
131    fn test_dense_vector_query_builder() {
132        let query = DenseVectorQuery::new(Field(0), vec![1.0, 2.0, 3.0])
133            .with_nprobe(64)
134            .with_rerank_factor(5.0);
135
136        assert_eq!(query.field, Field(0));
137        assert_eq!(query.vector.len(), 3);
138        assert_eq!(query.nprobe, 64);
139        assert_eq!(query.rerank_factor, 5.0);
140    }
141}
hermes_core/query/vector/dense.rs

hermes_core/query/vector/
dense.rs
