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/// Dense vector query for similarity search
11#[derive(Debug, Clone)]
12pub struct DenseVectorQuery {
13    /// Field containing the dense vectors
14    pub field: Field,
15    /// Query vector
16    pub vector: Vec<f32>,
17    /// Number of clusters to probe (for IVF indexes)
18    pub nprobe: usize,
19    /// Re-ranking factor (multiplied by k for candidate selection, e.g. 3.0)
20    pub rerank_factor: f32,
21    /// How to combine scores for multi-valued documents
22    pub combiner: MultiValueCombiner,
23}
24
25impl std::fmt::Display for DenseVectorQuery {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        write!(
28            f,
29            "Dense({}, dim={}, nprobe={}, rerank={})",
30            self.field.0,
31            self.vector.len(),
32            self.nprobe,
33            self.rerank_factor
34        )
35    }
36}
37
38impl DenseVectorQuery {
39    /// Create a new dense vector query
40    pub fn new(field: Field, vector: Vec<f32>) -> Self {
41        Self {
42            field,
43            vector,
44            nprobe: 32,
45            rerank_factor: 3.0,
46            combiner: MultiValueCombiner::Max,
47        }
48    }
49
50    /// Set the number of clusters to probe (for IVF indexes)
51    pub fn with_nprobe(mut self, nprobe: usize) -> Self {
52        self.nprobe = nprobe;
53        self
54    }
55
56    /// Set the re-ranking factor (e.g. 3.0 = fetch 3x candidates for reranking)
57    pub fn with_rerank_factor(mut self, factor: f32) -> Self {
58        self.rerank_factor = factor;
59        self
60    }
61
62    /// Set the multi-value score combiner
63    pub fn with_combiner(mut self, combiner: MultiValueCombiner) -> Self {
64        self.combiner = combiner;
65        self
66    }
67}
68
69impl Query for DenseVectorQuery {
70    fn scorer<'a>(&self, reader: &'a SegmentReader, limit: usize) -> ScorerFuture<'a> {
71        let field = self.field;
72        let vector = self.vector.clone();
73        let nprobe = self.nprobe;
74        let rerank_factor = self.rerank_factor;
75        let combiner = self.combiner;
76        Box::pin(async move {
77            let results = reader
78                .search_dense_vector(field, &vector, limit, nprobe, rerank_factor, combiner)
79                .await?;
80
81            Ok(Box::new(VectorResultScorer::new(results, field.0)) as Box<dyn Scorer>)
82        })
83    }
84
85    #[cfg(feature = "sync")]
86    fn scorer_sync<'a>(
87        &self,
88        reader: &'a SegmentReader,
89        limit: usize,
90    ) -> crate::Result<Box<dyn Scorer + 'a>> {
91        let results = reader.search_dense_vector_sync(
92            self.field,
93            &self.vector,
94            limit,
95            self.nprobe,
96            self.rerank_factor,
97            self.combiner,
98        )?;
99        Ok(Box::new(VectorResultScorer::new(results, self.field.0)) as Box<dyn Scorer>)
100    }
101
102    fn count_estimate<'a>(&self, _reader: &'a SegmentReader) -> CountFuture<'a> {
103        Box::pin(async move { Ok(u32::MAX) })
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    #[test]
112    fn test_dense_vector_query_builder() {
113        let query = DenseVectorQuery::new(Field(0), vec![1.0, 2.0, 3.0])
114            .with_nprobe(64)
115            .with_rerank_factor(5.0);
116
117        assert_eq!(query.field, Field(0));
118        assert_eq!(query.vector.len(), 3);
119        assert_eq!(query.nprobe, 64);
120        assert_eq!(query.rerank_factor, 5.0);
121    }
122}
hermes_core/query/vector/dense.rs

hermes_core/query/vector/
dense.rs
