urna_runtime/
mmap_file.rs

1//! `MmapUrnaFile`: owns the mmap, parses section metadata once, exposes
2//! search/inspect entry points. Search math lives in `super::search`;
3//! dtype-agnostic vector materialization lives in `super::materialize`.
4
5use std::path::Path;
6
7use memmap2::Mmap;
8use urna_format::layout::{
9    SECTION_BM25_INDEX, SECTION_CHUNK_IDS, SECTION_CHUNKS_ORIGINAL_SPANS, SECTION_EMBEDDINGS,
10    SECTION_EMBEDDINGS_FP, SECTION_GRAPH_ADJACENCY, SECTION_HNSW_INDEX,
11};
12use urna_format::reader::{UrnaView, validate_slab_values};
13use urna_format::sections::{
14    BlobRefRecord, OriginalSpan, decode_chunk_ids, decode_chunks_original_spans,
15};
16
17use crate::ann;
18use crate::bm25;
19use crate::dtype::DType;
20use crate::error::RuntimeError;
21use crate::graph::CsrIndex;
22use crate::materialize::PackedVectors;
23use crate::rerank::FpSlab;
24use crate::simd::{self, SimdBackend};
25
26pub struct MmapUrnaFile {
27    pub(crate) _mmap: Mmap,
28    pub(crate) embedding_dim: usize,
29    pub(crate) n_embeddings: usize,
30    pub(crate) dtype: DType,
31    /// Byte offset (within the mmap) of the embeddings section payload.
32    pub(crate) embeddings_offset: usize,
33    /// Total physical bytes of the embeddings section.
34    pub(crate) embeddings_size: usize,
35    /// Optional full-precision rerank slab (`embeddings_fp`, 0x09): when
36    /// present the exact rerank reads it instead of the stored dtype slab.
37    pub(crate) embeddings_fp: Option<FpSlab>,
38    pub(crate) chunk_ids: Vec<String>,
39    pub(crate) spans: Vec<OriginalSpan>,
40    pub(crate) embedding_model: String,
41    /// Manifest `model_hash`: the granular fingerprint of the model that
42    /// produced the corpus embeddings. Exposed so a caller embedding a query
43    /// can verify its embedder matches the corpus (the honesty gate).
44    pub(crate) model_hash: String,
45    pub(crate) file_hash: String,
46    pub(crate) content_hash: String,
47    /// Optional ANN index. Built from the HNSW section payload at open
48    /// time (eager: build cost is paid once, queries get fast path).
49    pub(crate) ann_index: Option<ann::HnswIndex>,
50    /// Optional BM25 index. Mostly tiny ints; deserialized eagerly.
51    pub(crate) bm25_index: Option<bm25::Bm25Index>,
52    /// Optional chunk-to-chunk graph (graph_adjacency 0x0C). Opened behind
53    /// the manifest `graph_present` capability, like ann/bm25. A candidate
54    /// generator only: its frontier feeds the exact rerank, never a score.
55    pub(crate) graph_index: Option<CsrIndex>,
56    /// Optional blob_refs (0x14) table, opened behind the additive
57    /// `blobs_present` capability. content-hash references to the source
58    /// media blobs (self-contained or catalog).
59    pub(crate) blob_refs: Option<Vec<BlobRefRecord>>,
60    /// Optional blob_data (0x17) offset table; blob bytes are sliced
61    /// lazily off the mmap by `blob_bytes` (impl in blobs.rs).
62    pub(crate) blob_data: Option<crate::blobs::OpenBlobData>,
63    /// Optional multimodal spaces (0x15 + bands), opened behind the
64    /// additive `supports_multimodal` capability. each space's band slab
65    /// is scored by the per-space exact search, never by the text path.
66    pub(crate) spaces: Option<Vec<crate::spaces::OpenSpace>>,
67    /// What the manifest says the search path is. The runtime honors
68    /// this at search time.
69    pub(crate) declared_index_type: String,
70    pub(crate) declared_score_type: String,
71}
72
73impl MmapUrnaFile {
74    pub fn open(path: &Path) -> Result<Self, RuntimeError> {
75        let file = std::fs::File::open(path)?;
76        // SAFETY: `Mmap::map` is unsafe because a concurrent writer could
77        // truncate or rewrite the file under the mapping. The runtime opens
78        // corpora read-only and never rewrites a `.urna` in place, so only an
79        // external live writer can violate this, which is outside the
80        // supported model of every mmap-based store (checksums detect
81        // corruption after the fact, they do not guard a live writer).
82        let mmap = unsafe { Mmap::map(&file)? };
83        let view = UrnaView::from_bytes(&mmap)?;
84        view.validate_embeddings_values()?;
85
86        let dim = view.header.embedding_dim as usize;
87        let n = view.header.n_embeddings as usize;
88        let dtype = DType::from_str(&view.manifest.dtype)?;
89
90        let entry = view.entry(SECTION_EMBEDDINGS)?;
91        let embeddings_offset = entry.offset as usize;
92        let embeddings_size = entry.size as usize;
93
94        // Optional full-precision rerank slab (0x09); see rerank::FpSlab.
95        let embeddings_fp = FpSlab::detect(&view.section_table, n, dim)?;
96        if embeddings_fp.is_some() {
97            // the rerank scores this slab, so it passes the same NaN gate.
98            let e = view.entry(SECTION_EMBEDDINGS_FP)?;
99            let fp = view.get_section_data(SECTION_EMBEDDINGS_FP)?;
100            validate_slab_values(e.encoding, fp, n, dim)?;
101        }
102
103        // Decoded chunk_ids / spans (handles zstd transparently).
104        let chunk_ids = decode_chunk_ids(&view.decoded_section(SECTION_CHUNK_IDS)?, n)?;
105        let mut spans =
106            decode_chunks_original_spans(&view.decoded_section(SECTION_CHUNKS_ORIGINAL_SPANS)?, n)?;
107
108        // Optional ANN section. Materialize f32 vectors from the
109        // embeddings section so the graph can compute distances at
110        // search time independent of the on-disk dtype.
111        let ann_index = if view
112            .section_table
113            .iter()
114            .any(|e| e.section_id == SECTION_HNSW_INDEX)
115        {
116            let bytes = view.decoded_section(SECTION_HNSW_INDEX)?;
117            let mut idx = ann::HnswIndex::from_bytes(&bytes, n, dim)?;
118            let emb_bytes = view.get_section_data(SECTION_EMBEDDINGS)?;
119            // Keep the vectors in their on-disk packing (no n*dim*4 f32
120            // expansion): the graph decodes one row at a time on demand.
121            let store = PackedVectors::from_section(&view.manifest.dtype, emb_bytes, n, dim)?;
122            idx.attach_store(store);
123            Some(idx)
124        } else {
125            None
126        };
127
128        let bm25_index = if view
129            .section_table
130            .iter()
131            .any(|e| e.section_id == SECTION_BM25_INDEX)
132        {
133            let bytes = view.decoded_section(SECTION_BM25_INDEX)?;
134            Some(bm25::Bm25Index::from_bytes(&bytes)?)
135        } else {
136            None
137        };
138
139        // Optional graph_adjacency section (0x0C). gated behind the additive
140        // `graph_present` capability (delivered via Option<CapabilitiesExt>),
141        // mirroring how ann/bm25 are opened. raw payload; the csr bitpacks its
142        // own integer columns with intpack internally.
143        let graph_present = view
144            .manifest
145            .capabilities_ext
146            .as_ref()
147            .and_then(|e| e.graph_present)
148            .unwrap_or(false);
149        let graph_index = if graph_present
150            && view
151                .section_table
152                .iter()
153                .any(|e| e.section_id == SECTION_GRAPH_ADJACENCY)
154        {
155            let bytes = view.decoded_section(SECTION_GRAPH_ADJACENCY)?;
156            Some(CsrIndex::from_bytes(&bytes, n)?)
157        } else {
158            None
159        };
160
161        // Optional blob pair (0x14/0x16): opens behind the additive
162        // `blobs_present` capability; the overlay rewrites blob-pointing
163        // spans in place. see blobs.rs.
164        let blob_refs = crate::blobs::open_blob_sections(&view, &mut spans)?;
165        let blob_data = match &blob_refs {
166            Some(refs) => crate::blobs::open_blob_data(&view, refs.len())?,
167            None => None,
168        };
169
170        // Optional multimodal spaces (0x15 + bands): opens behind the
171        // additive `supports_multimodal` capability. see spaces.rs.
172        let spaces = crate::spaces::open_space_sections(&view)?;
173        let embedding_model = view.manifest.embedding_model.clone();
174        let model_hash = view.manifest.model_hash.clone();
175        let declared_index_type = view.manifest.index_type.clone();
176        let declared_score_type = view.manifest.score_type.clone();
177        let file_hash = view.file_hash_hex();
178        let content_hash = view.content_hash_hex()?;
179        drop(view);
180
181        Ok(Self {
182            _mmap: mmap,
183            embedding_dim: dim,
184            n_embeddings: n,
185            dtype,
186            embeddings_offset,
187            embeddings_size,
188            embeddings_fp,
189            chunk_ids,
190            spans,
191            embedding_model,
192            model_hash,
193            file_hash,
194            content_hash,
195            ann_index,
196            bm25_index,
197            graph_index,
198            blob_refs,
199            blob_data,
200            spaces,
201            declared_index_type,
202            declared_score_type,
203        })
204    }
205
206    pub fn embedding_dim(&self) -> usize {
207        self.embedding_dim
208    }
209    pub fn n_embeddings(&self) -> usize {
210        self.n_embeddings
211    }
212    pub fn dtype(&self) -> DType {
213        self.dtype
214    }
215    pub fn file_hash(&self) -> &str {
216        &self.file_hash
217    }
218    pub fn content_hash(&self) -> &str {
219        &self.content_hash
220    }
221    pub fn simd_backend(&self) -> SimdBackend {
222        simd::detect_backend()
223    }
224    pub fn declared_index_type(&self) -> &str {
225        &self.declared_index_type
226    }
227    /// Manifest `model_hash` (the model fingerprint the corpus was built
228    /// with). Callers embed the query with their own model and compare.
229    pub fn model_hash(&self) -> &str {
230        &self.model_hash
231    }
232    pub fn declared_score_type(&self) -> &str {
233        &self.declared_score_type
234    }
235    pub fn has_ann(&self) -> bool {
236        self.ann_index.is_some()
237    }
238    /// the hnsw beam floor (the file's `ef_construction`, 400 for python
239    /// builds): `search_ann(_, _, ef)` searches `max(ef, k, floor)`
240    /// candidates, so an `ef` below the floor changes nothing. `None`
241    /// without an hnsw section.
242    pub fn ann_ef_floor(&self) -> Option<usize> {
243        self.ann_index.as_ref().map(|idx| idx.ef_search)
244    }
245    pub fn has_bm25(&self) -> bool {
246        self.bm25_index.is_some()
247    }
248    pub fn has_graph(&self) -> bool {
249        self.graph_index.is_some()
250    }
251    pub fn has_blobs(&self) -> bool {
252        self.blob_refs.is_some()
253    }
254    /// The blob_refs (0x14) table, when the file declares `blobs_present`:
255    /// content-hash references to the source media blobs, in table order
256    /// (the overlay's `blob_ref_index` addresses this slice).
257    pub fn blob_refs(&self) -> Option<&[BlobRefRecord]> {
258        self.blob_refs.as_deref()
259    }
260    pub fn has_spaces(&self) -> bool {
261        self.spaces.is_some()
262    }
263    /// Names of the multimodal spaces listed in the space_table (empty
264    /// when the file has no `supports_multimodal` capability).
265    pub fn space_names(&self) -> Vec<&str> {
266        self.spaces
267            .as_deref()
268            .unwrap_or(&[])
269            .iter()
270            .map(|s| s.entry.name.as_str())
271            .collect()
272    }
273    /// One opened space by name, with its band slab sliced off the mmap.
274    pub(crate) fn space(&self, name: &str) -> Option<(&crate::spaces::OpenSpace, &[u8])> {
275        let s = self
276            .spaces
277            .as_deref()?
278            .iter()
279            .find(|s| s.entry.name == name)?;
280        Some((s, &self._mmap[s.offset..s.offset + s.size]))
281    }
282
283    /// Re-run all reader-side validation. The file was already
284    /// validated at `open()` time, but callers can invoke this
285    /// explicitly to detect tampering after the fact (e.g. the mmap
286    /// pages got swapped under the runtime).
287    pub fn revalidate(&self) -> Result<(), RuntimeError> {
288        let view = UrnaView::from_bytes(&self._mmap)?;
289        view.validate_embeddings_values()?;
290        let _ = view.search_contract()?;
291        Ok(())
292    }
293
294    pub(crate) fn embeddings_bytes(&self) -> &[u8] {
295        &self._mmap[self.embeddings_offset..self.embeddings_offset + self.embeddings_size]
296    }
297
298    /// The full-precision rerank slab + its dtype, if an `embeddings_fp`
299    /// (0x09) section is present. The rerank handle prefers this over the
300    /// stored dtype slab so a sub-int8 corpus still returns a real cosine.
301    pub(crate) fn embeddings_fp_slab(&self) -> Option<(&[u8], DType)> {
302        self.embeddings_fp
303            .map(|fp| (&self._mmap[fp.offset..fp.offset + fp.size], fp.dtype))
304    }
305}
urna_runtime/mmap_file.rs

urna_runtime/
mmap_file.rs
