sparse_vector/
mmap_index.rs

1//! Flat binary mmap format for sparse-vector posting lists.
2//!
3//! Layout (every struct is `#[repr(C)]` for a stable layout):
4//!
5//! ```text
6//! [FileHeader]                    16 bytes
7//! [DimHeader × num_dims]          16 bytes × N
8//! [PostingEntry × total_entries]  16 bytes × M
9//! ```
10//!
11//! Each entry carries a `max_next_weight` ceiling. Files written here store
12//! the inclusive ceiling of the wand module (`tail_max`: the maximum weight
13//! over the entry and everything after it); readers fold
14//! `max(weight, max_next_weight)`, so a file whose ceiling excludes the
15//! entry itself reads identically.
16
17use std::collections::HashMap;
18use std::path::Path;
19
20use memmap2::Mmap;
21
22use crate::index::{run_search, SparseVector};
23use crate::wand::{MmapCursor, Postings};
24
25const MAGIC: u32 = 0x53505253; // "SPRS"
26const FORMAT_VERSION: u32 = 1;
27
28#[repr(C)]
29struct FileHeader {
30    magic: u32,
31    version: u32,
32    num_dims: u32,
33    num_vectors: u32,
34}
35
36#[repr(C)]
37struct DimHeader {
38    offset: u64,
39    count: u32,
40    _pad: u32,
41}
42
43/// On-disk posting entry.
44#[repr(C)]
45#[derive(Clone, Copy, Debug)]
46pub struct PostingEntry {
47    pub record_id: u64,
48    pub weight: f32,
49    /// Ceiling over the rest of the list (see the module docs).
50    pub max_next_weight: f32,
51}
52
53/// Mmap'd posting data — read-only view of `sparse.mmap`.
54pub struct MmapPostingData {
55    mmap: Mmap,
56    num_dims: usize,
57    num_vectors: usize,
58}
59
60impl MmapPostingData {
61    pub fn open(path: &Path) -> Result<Self, String> {
62        let file = std::fs::File::open(path)
63            .map_err(|e| format!("cannot open {}: {e}", path.display()))?;
64        let mmap = unsafe { Mmap::map(&file) }
65            .map_err(|e| format!("cannot mmap {}: {e}", path.display()))?;
66
67        if mmap.len() < std::mem::size_of::<FileHeader>() {
68            return Err("sparse.mmap too small for header".into());
69        }
70        let header = unsafe { &*(mmap.as_ptr() as *const FileHeader) };
71        if header.magic != MAGIC {
72            return Err(format!("bad magic: {:#x}", header.magic));
73        }
74        if header.version != FORMAT_VERSION {
75            return Err(format!("unsupported version: {}", header.version));
76        }
77
78        Ok(Self {
79            mmap,
80            num_dims: header.num_dims as usize,
81            num_vectors: header.num_vectors as usize,
82        })
83    }
84
85    pub fn num_dims(&self) -> usize {
86        self.num_dims
87    }
88
89    pub fn num_vectors(&self) -> usize {
90        self.num_vectors
91    }
92
93    /// The entries of a remapped dimension, sorted by record id; empty for
94    /// an empty or unknown dimension.
95    pub fn entries(&self, dim_idx: usize) -> &[PostingEntry] {
96        if dim_idx >= self.num_dims {
97            return &[];
98        }
99        let dim_headers_offset = std::mem::size_of::<FileHeader>();
100        let dh_ptr = unsafe {
101            self.mmap
102                .as_ptr()
103                .add(dim_headers_offset + dim_idx * std::mem::size_of::<DimHeader>())
104        } as *const DimHeader;
105        let dh = unsafe { &*dh_ptr };
106
107        if dh.count == 0 {
108            return &[];
109        }
110
111        let entries_ptr =
112            unsafe { self.mmap.as_ptr().add(dh.offset as usize) } as *const PostingEntry;
113        unsafe { std::slice::from_raw_parts(entries_ptr, dh.count as usize) }
114    }
115
116    /// Cursor over a remapped dimension, `None` when it has no postings.
117    pub fn cursor(&self, dim_idx: usize) -> Option<MmapCursor<'_>> {
118        MmapCursor::open(self, dim_idx as u32)
119    }
120
121    /// Load a dimension's entries into an in-RAM list, recomputing the
122    /// ceilings from the weights.
123    pub fn load_postings(&self, dim_idx: usize) -> Postings {
124        let pairs: Vec<(u64, f32)> = self
125            .entries(dim_idx)
126            .iter()
127            .map(|e| (e.record_id, e.weight))
128            .collect();
129        Postings::from_sorted_pairs(&pairs)
130    }
131}
132
133// ---------------------------------------------------------------------------
134// Search using mmap data (no RAM postings needed)
135// ---------------------------------------------------------------------------
136
137/// Top-`limit` search straight from the mapping, without loading anything
138/// into RAM. `dim_map` translates the query's token ids into the file's
139/// dimension indices.
140pub fn search_mmap<F: Fn(u64) -> bool>(
141    mmap: &MmapPostingData,
142    dim_map: &HashMap<u32, usize>,
143    query: &SparseVector,
144    limit: usize,
145    filter: &F,
146) -> Vec<(u64, f32)> {
147    if query.is_empty() || mmap.num_vectors() == 0 {
148        return Vec::new();
149    }
150    run_search(query, dim_map, limit, filter, |dim| mmap.cursor(dim as usize))
151}
152
153/// [`search_mmap`] restricted to `allowed` ids (see
154/// [`run_search_allowed`](crate::index::run_search_allowed)).
155pub fn search_mmap_allowed(
156    mmap: &MmapPostingData,
157    dim_map: &HashMap<u32, usize>,
158    query: &SparseVector,
159    limit: usize,
160    allowed: &[u64],
161) -> Vec<(u64, f32)> {
162    if query.is_empty() || mmap.num_vectors() == 0 {
163        return Vec::new();
164    }
165    crate::index::run_search_allowed(query, dim_map, limit, allowed, |dim| mmap.cursor(dim as usize))
166}
167
168// ---------------------------------------------------------------------------
169// Write mmap format
170// ---------------------------------------------------------------------------
171
172/// Write the flat binary mmap format from in-RAM posting lists, one per
173/// remapped dimension.
174pub fn write_mmap_file(
175    path: &Path,
176    postings: &[Postings],
177    num_vectors: u32,
178) -> Result<(), String> {
179    use std::io::{BufWriter, Write};
180
181    let num_dims = postings.len() as u32;
182    let header_size = std::mem::size_of::<FileHeader>();
183    let dim_headers_size = num_dims as usize * std::mem::size_of::<DimHeader>();
184    let entries_start = header_size + dim_headers_size;
185
186    let file = std::fs::File::create(path)
187        .map_err(|e| format!("cannot create {}: {e}", path.display()))?;
188    let mut out = BufWriter::new(file);
189
190    let header = FileHeader {
191        magic: MAGIC,
192        version: FORMAT_VERSION,
193        num_dims,
194        num_vectors,
195    };
196    out.write_all(as_bytes(&header))
197        .map_err(|e| format!("write header: {e}"))?;
198
199    let mut current_offset = entries_start;
200    for p in postings {
201        let dh = DimHeader {
202            offset: current_offset as u64,
203            count: p.len() as u32,
204            _pad: 0,
205        };
206        out.write_all(as_bytes(&dh))
207            .map_err(|e| format!("write dim header: {e}"))?;
208        current_offset += p.len() * std::mem::size_of::<PostingEntry>();
209    }
210
211    for p in postings {
212        for x in p.as_slice() {
213            let entry = PostingEntry {
214                record_id: x.id,
215                weight: x.weight,
216                max_next_weight: x.tail_max,
217            };
218            out.write_all(as_bytes(&entry))
219                .map_err(|e| format!("write entry: {e}"))?;
220        }
221    }
222
223    out.flush().map_err(|e| format!("flush {}: {e}", path.display()))
224}
225
226/// Reinterpret a `#[repr(C)]` struct without padding as bytes.
227fn as_bytes<T: Sized>(val: &T) -> &[u8] {
228    unsafe { std::slice::from_raw_parts(val as *const T as *const u8, std::mem::size_of::<T>()) }
229}
sparse_vector/mmap_index.rs

sparse_vector/
mmap_index.rs
