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
55    #[test]
56    fn mmap_roundtrip() {
57        let mut b = HnswBuilder::new(4, VectorMetric::Cosine, HnswConfig::default());
58        b.insert(RowId::new(0), vec![1.0, 0.0, 0.0, 0.0]);
59        b.insert(RowId::new(1), vec![0.0, 1.0, 0.0, 0.0]);
60        let idx = b.build();
61        let bytes = HnswSerializer::to_bytes(&idx).unwrap();
62
63        let loaded = MmapLoader::from_bytes(&bytes).unwrap();
64        assert_eq!(loaded.node_count(), 2);
65        let r = loaded.search(&[1.0, 0.0, 0.0, 0.0], 1, 50);
66        assert_eq!(r[0].0, RowId::new(0));
67    }
68}
ailake_index/mmap_loader.rs

ailake_index/
mmap_loader.rs
