ailake_index/
mmap_loader.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2use std::io::Write;
3
4use ailake_core::{AilakeError, AilakeResult};
5use memmap2::Mmap;
6use tracing::debug;
7
8use crate::hnsw::HnswIndex;
9use crate::serialize::HnswSerializer;
10
11pub struct MmapLoader;
12
13impl MmapLoader {
14    /// Write `bytes` to a temporary file, mmap it, and deserialize the HNSW index.
15    /// Using mmap lets the OS lazily page in only the graph nodes touched during search —
16    /// critical for large indexes (>1 GB) where loading the full file would waste RAM.
17    pub fn from_bytes(bytes: &[u8]) -> AilakeResult<HnswIndex> {
18        debug!(
19            "ailake: loading HNSW index via mmap ({} bytes)",
20            bytes.len()
21        );
22        let mut tmp = tempfile::tempfile().map_err(|e| {
23            AilakeError::Store(format!("failed to create tempfile for HNSW mmap: {e}"))
24        })?;
25        tmp.write_all(bytes).map_err(|e| {
26            AilakeError::Store(format!(
27                "failed to write {} bytes to HNSW tempfile: {e}",
28                bytes.len()
29            ))
30        })?;
31        // SAFETY: the backing file is not modified after mmap is created.
32        // The mmap is dropped before the function returns (index owns its data).
33        let mmap = unsafe { Mmap::map(&tmp) }.map_err(|e| {
34            AilakeError::Store(format!(
35                "mmap failed for HNSW tempfile ({} bytes): {e}",
36                bytes.len()
37            ))
38        })?;
39        let idx = HnswSerializer::from_bytes(&mmap)?;
40        debug!(
41            "ailake: HNSW index loaded — {} nodes via mmap",
42            idx.node_count()
43        );
44        Ok(idx)
45    }
46}
47
48#[cfg(test)]
49mod tests {
50    use super::*;
51    use crate::hnsw::{HnswBuilder, HnswConfig};
52    use crate::serialize::HnswSerializer;
53    use ailake_core::{RowId, VectorMetric};
54    use proptest::prelude::*;
55    use rand::Rng;
56
57    #[test]
58    fn mmap_roundtrip() {
59        let mut b = HnswBuilder::new(4, VectorMetric::Cosine, HnswConfig::default());
60        b.insert(RowId::new(0), vec![1.0, 0.0, 0.0, 0.0]);
61        b.insert(RowId::new(1), vec![0.0, 1.0, 0.0, 0.0]);
62        let idx = b.build();
63        let bytes = HnswSerializer::to_bytes(&idx).unwrap();
64
65        let loaded = MmapLoader::from_bytes(&bytes).unwrap();
66        assert_eq!(loaded.node_count(), 2);
67        let r = loaded.search(&[1.0, 0.0, 0.0, 0.0], 1, 50);
68        assert_eq!(r[0].0, RowId::new(0));
69    }
70
71    // ── Fuzz-style: property tests for mmap loader ────────────────────────
72
73    fn arb_query(dim: usize) -> impl Strategy<Value = Vec<f32>> {
74        let val = (-10.0f32..10.0).prop_filter("no NaN/Inf", |x| x.is_finite());
75        proptest::collection::vec(val, dim)
76    }
77
78    proptest! {
79        #[test]
80        fn prop_mmap_roundtrip_search(
81            n_nodes in 2usize..10,
82            m in 4usize..8,
83            query in arb_query(4),
84        ) {
85            let dim = 4u32;
86            let m = m.max(2);
87            let mut rng = rand::thread_rng();
88
89            let mut b = HnswBuilder::new(dim, VectorMetric::Cosine, HnswConfig { m, ef_construction: 50, max_elements: n_nodes.max(10) });
90            let mut inserted: Vec<(RowId, Vec<f32>)> = Vec::new();
91            for i in 0..n_nodes {
92                let v: Vec<f32> = (0..dim as usize).map(|_| rng.gen::<f32>() * 2.0 - 1.0).collect();
93                let rid = RowId::new(i as u64);
94                b.insert(rid, v.clone());
95                inserted.push((rid, v));
96            }
97            let idx = b.build();
98            let bytes = HnswSerializer::to_bytes(&idx)
99                .expect("serialization should succeed");
100
101            let loaded = MmapLoader::from_bytes(&bytes)
102                .expect("mmap load should succeed");
103            prop_assert_eq!(loaded.node_count(), n_nodes as u64);
104
105            let results = loaded.search(&query, 1, 50);
106            if !results.is_empty() {
107                prop_assert!(results[0].1 >= 0.0, "distance must be non-negative");
108            }
109        }
110
111        #[test]
112        fn prop_mmap_empty_index(
113            dim in 2u32..8,
114            m in 2usize..8,
115        ) {
116            let b = HnswBuilder::new(dim, VectorMetric::Cosine, HnswConfig { m, ef_construction: 50, max_elements: 10 });
117            let idx = b.build();
118            let bytes = HnswSerializer::to_bytes(&idx)
119                .expect("empty index should serialize");
120
121            let loaded = MmapLoader::from_bytes(&bytes)
122                .expect("empty index should load via mmap");
123            prop_assert_eq!(loaded.node_count(), 0);
124        }
125    }
126}
ailake_index/mmap_loader.rs

ailake_index/
mmap_loader.rs
