urna_runtime/
space_search.rs

1//! Per-space exact search over the multimodal vector bands (0x20-0x2F).
2//!
3//! This is the ONLY path that reads a non-text band, and the text paths
4//! (exact/ann/graph/hybrid in `search.rs`) are the only paths that read
5//! the canonical 0x04 slab: the two never mix, so a text query can never
6//! be scored against the vision band by accident. every score is the real
7//! cosine recomputed by the shared `RerankSource` (the same kernel the
8//! text exact path uses), never an ann proxy.
9
10use crate::dtype::DType;
11use crate::error::RuntimeError;
12use crate::mmap_file::MmapUrnaFile;
13use crate::rerank::RerankSource;
14use crate::search::sort_scores_desc;
15use crate::{RerankSourceKind, SearchExplain, SearchResult};
16
17impl MmapUrnaFile {
18    /// Exact flat search over one named multimodal space (e.g. "vision").
19    /// the query must be embedded with the model the space's
20    /// `model_hash` fingerprints and have the space's dim; both are
21    /// checked up front (the per-space honesty gate), so a text-tower
22    /// query fails loudly instead of silently scoring the vision band.
23    /// recall is 1.0: this is the exact ground truth over that space.
24    pub fn search_space(
25        &self,
26        name: &str,
27        query: &[f32],
28        k: i32,
29        expected_model_hash: Option<&str>,
30    ) -> Result<SearchResult, RuntimeError> {
31        let t0 = std::time::Instant::now();
32        let (space, band) = self
33            .space(name)
34            .ok_or_else(|| RuntimeError::SpaceNotFound(name.to_string()))?;
35        let entry = &space.entry;
36
37        if let Some(expected) = expected_model_hash {
38            if expected != entry.model_hash {
39                return Err(RuntimeError::SpaceModelMismatch {
40                    space: name.to_string(),
41                    expected: expected.to_string(),
42                    actual: entry.model_hash.clone(),
43                });
44            }
45        }
46        if k <= 0 {
47            return Err(RuntimeError::InvalidK(k));
48        }
49        if query.is_empty() {
50            return Err(RuntimeError::EmptyQuery);
51        }
52        if query.len() != entry.dim as usize {
53            return Err(RuntimeError::DimensionMismatch {
54                expected: entry.dim as usize,
55                got: query.len(),
56            });
57        }
58        for &v in query {
59            if v.is_nan() || v.is_infinite() {
60                return Err(RuntimeError::InvalidQueryValue);
61            }
62        }
63        let mut qnorm = query.to_vec();
64        let norm: f32 = qnorm.iter().map(|x| x * x).sum::<f32>().sqrt();
65        if norm == 0.0 {
66            return Err(RuntimeError::ZeroNormQuery);
67        }
68        for x in &mut qnorm {
69            *x /= norm;
70        }
71
72        let dtype = DType::from_str(entry.dtype_str())?;
73        let n = entry.n_vectors as usize;
74        let mut src = RerankSource::new(dtype, band, n, entry.dim as usize)?;
75        let mut scores: Vec<(usize, f32)> = Vec::with_capacity(n);
76        for i in 0..n {
77            scores.push((i, src.score(&qnorm, i)));
78        }
79        sort_scores_desc(&mut scores);
80        let k_usize = k as usize;
81        let truncated = k_usize < n;
82        let top = &scores[..k_usize.min(n)];
83        let hits = self.materialize_hits(top, "space", false);
84        Ok(SearchResult {
85            hits: hits.clone(),
86            query_time_ms: t0.elapsed().as_secs_f64() * 1000.0,
87            index_type: "space",
88            recall: 1.0,
89            truncated,
90            k_requested: k,
91            k_returned: hits.len(),
92            explain: SearchExplain::exact(n, RerankSourceKind::from_dtype(dtype)),
93        })
94    }
95}
urna_runtime/space_search.rs

urna_runtime/
space_search.rs
