flow_knn/
io.rs

1//! Versioned on-disk encoding for [`KnnGraph`].
2//!
3//! Dense packed layout (little-endian):
4//! - magic `b"FKNN"` (4)
5//! - version `u32` (currently 1)
6//! - `n` `u64`, `k` `u64`
7//! - metric discriminant `u8` (0=Euclidean, 1=EuclideanSq, 2=Cosine, 3=Manhattan)
8//! - provenance length `u32` + UTF-8 bytes (0 = None)
9//! - `n * k` `u32` indices (row-major)
10//! - `n * k` `f32` distances (row-major)
11
12use std::fs::File;
13use std::io::{Read, Write};
14use std::path::Path;
15
16use crate::config::DistanceMetric;
17use crate::error::KnnError;
18use crate::graph::{KnnGraph, NeighborList};
19
20const MAGIC: &[u8; 4] = b"FKNN";
21const VERSION: u32 = 1;
22
23fn metric_to_u8(metric: DistanceMetric) -> u8 {
24    match metric {
25        DistanceMetric::Euclidean => 0,
26        DistanceMetric::EuclideanSq => 1,
27        DistanceMetric::Cosine => 2,
28        DistanceMetric::Manhattan => 3,
29    }
30}
31
32fn metric_from_u8(v: u8) -> Result<DistanceMetric, KnnError> {
33    match v {
34        0 => Ok(DistanceMetric::Euclidean),
35        1 => Ok(DistanceMetric::EuclideanSq),
36        2 => Ok(DistanceMetric::Cosine),
37        3 => Ok(DistanceMetric::Manhattan),
38        _ => Err(KnnError::Io(format!(
39            "unknown distance metric discriminant {v}"
40        ))),
41    }
42}
43
44fn write_u32<W: Write>(w: &mut W, v: u32) -> Result<(), KnnError> {
45    w.write_all(&v.to_le_bytes())
46        .map_err(|e| KnnError::Io(e.to_string()))
47}
48
49fn write_u64<W: Write>(w: &mut W, v: u64) -> Result<(), KnnError> {
50    w.write_all(&v.to_le_bytes())
51        .map_err(|e| KnnError::Io(e.to_string()))
52}
53
54fn read_exact_arr<R: Read, const N: usize>(r: &mut R) -> Result<[u8; N], KnnError> {
55    let mut buf = [0u8; N];
56    r.read_exact(&mut buf)
57        .map_err(|e| KnnError::Io(e.to_string()))?;
58    Ok(buf)
59}
60
61fn read_u32<R: Read>(r: &mut R) -> Result<u32, KnnError> {
62    Ok(u32::from_le_bytes(read_exact_arr(r)?))
63}
64
65fn read_u64<R: Read>(r: &mut R) -> Result<u64, KnnError> {
66    Ok(u64::from_le_bytes(read_exact_arr(r)?))
67}
68
69/// Write a [`KnnGraph`] to `path` (creates/overwrites the file).
70pub fn write_knn_graph(path: &Path, graph: &KnnGraph) -> Result<(), KnnError> {
71    if graph.neighbors.len() != graph.n {
72        return Err(KnnError::GraphSizeMismatch {
73            graph_n: graph.n,
74            neighbors_len: graph.neighbors.len(),
75            data_n: graph.n,
76        });
77    }
78    for (i, nbr) in graph.neighbors.iter().enumerate() {
79        if nbr.indices.len() != graph.k || nbr.distances.len() != graph.k {
80            return Err(KnnError::Io(format!(
81                "neighbor list {i} has indices={} distances={} but graph.k={}",
82                nbr.indices.len(),
83                nbr.distances.len(),
84                graph.k
85            )));
86        }
87    }
88
89    let mut file = File::create(path).map_err(|e| KnnError::Io(e.to_string()))?;
90    file.write_all(MAGIC)
91        .map_err(|e| KnnError::Io(e.to_string()))?;
92    write_u32(&mut file, VERSION)?;
93    write_u64(&mut file, graph.n as u64)?;
94    write_u64(&mut file, graph.k as u64)?;
95    file.write_all(&[metric_to_u8(graph.metric)])
96        .map_err(|e| KnnError::Io(e.to_string()))?;
97
98    let prov = graph.provenance.as_deref().unwrap_or("");
99    let prov_bytes = prov.as_bytes();
100    if prov_bytes.len() > u32::MAX as usize {
101        return Err(KnnError::Io(
102            "provenance string exceeds u32 length".to_string(),
103        ));
104    }
105    write_u32(&mut file, prov_bytes.len() as u32)?;
106    file.write_all(prov_bytes)
107        .map_err(|e| KnnError::Io(e.to_string()))?;
108
109    let total = graph
110        .n
111        .checked_mul(graph.k)
112        .ok_or_else(|| KnnError::Io("n*k overflow".to_string()))?;
113    let byte_len = total
114        .checked_mul(4)
115        .ok_or_else(|| KnnError::Io("n*k*4 overflow".to_string()))?;
116
117    // Stage packed LE payloads, then two bulk write_all calls (mirrors read path).
118    let mut idx_bytes = Vec::with_capacity(byte_len);
119    let mut dist_bytes = Vec::with_capacity(byte_len);
120    if cfg!(target_endian = "little") {
121        let mut indices = Vec::with_capacity(total);
122        let mut distances = Vec::with_capacity(total);
123        for nbr in &graph.neighbors {
124            indices.extend_from_slice(&nbr.indices);
125            distances.extend_from_slice(&nbr.distances);
126        }
127        idx_bytes.extend_from_slice(bytemuck::cast_slice::<u32, u8>(&indices));
128        dist_bytes.extend_from_slice(bytemuck::cast_slice::<f32, u8>(&distances));
129    } else {
130        for nbr in &graph.neighbors {
131            for &idx in &nbr.indices {
132                idx_bytes.extend_from_slice(&idx.to_le_bytes());
133            }
134        }
135        for nbr in &graph.neighbors {
136            for &dist in &nbr.distances {
137                dist_bytes.extend_from_slice(&dist.to_le_bytes());
138            }
139        }
140    }
141    file.write_all(&idx_bytes)
142        .map_err(|e| KnnError::Io(e.to_string()))?;
143    file.write_all(&dist_bytes)
144        .map_err(|e| KnnError::Io(e.to_string()))?;
145    file.sync_all().map_err(|e| KnnError::Io(e.to_string()))?;
146    Ok(())
147}
148
149/// Read a [`KnnGraph`] previously written by [`write_knn_graph`].
150pub fn read_knn_graph(path: &Path) -> Result<KnnGraph, KnnError> {
151    let mut file = File::open(path).map_err(|e| KnnError::Io(e.to_string()))?;
152    let magic = read_exact_arr::<_, 4>(&mut file)?;
153    if &magic != MAGIC {
154        return Err(KnnError::Io(format!(
155            "bad knn.bin magic: expected FKNN, got {:?}",
156            String::from_utf8_lossy(&magic)
157        )));
158    }
159    let version = read_u32(&mut file)?;
160    if version != VERSION {
161        return Err(KnnError::Io(format!(
162            "unsupported knn.bin version {version} (expected {VERSION})"
163        )));
164    }
165    let n = read_u64(&mut file)? as usize;
166    let k = read_u64(&mut file)? as usize;
167    let metric = metric_from_u8(read_exact_arr::<_, 1>(&mut file)?[0])?;
168    let prov_len = read_u32(&mut file)? as usize;
169    let provenance = if prov_len == 0 {
170        None
171    } else {
172        let mut buf = vec![0u8; prov_len];
173        file.read_exact(&mut buf)
174            .map_err(|e| KnnError::Io(e.to_string()))?;
175        let s = String::from_utf8(buf).map_err(|e| KnnError::Io(e.to_string()))?;
176        Some(s)
177    };
178
179    let total = n
180        .checked_mul(k)
181        .ok_or_else(|| KnnError::Io("n*k overflow".to_string()))?;
182    let byte_len = total
183        .checked_mul(4)
184        .ok_or_else(|| KnnError::Io("n*k*4 overflow".to_string()))?;
185
186    // Bulk-read directly into typed buffers (avoids u8 staging + second copy on LE).
187    let (indices, distances) = if cfg!(target_endian = "little") {
188        let mut indices = vec![0u32; total];
189        file.read_exact(bytemuck::cast_slice_mut(&mut indices))
190            .map_err(|e| KnnError::Io(e.to_string()))?;
191        let mut distances = vec![0f32; total];
192        file.read_exact(bytemuck::cast_slice_mut(&mut distances))
193            .map_err(|e| KnnError::Io(e.to_string()))?;
194        (indices, distances)
195    } else {
196        let mut idx_bytes = vec![0u8; byte_len];
197        file.read_exact(&mut idx_bytes)
198            .map_err(|e| KnnError::Io(e.to_string()))?;
199        let mut dist_bytes = vec![0u8; byte_len];
200        file.read_exact(&mut dist_bytes)
201            .map_err(|e| KnnError::Io(e.to_string()))?;
202        let indices: Vec<u32> = idx_bytes
203            .chunks_exact(4)
204            .map(|c| u32::from_le_bytes(c.try_into().unwrap()))
205            .collect();
206        let distances: Vec<f32> = dist_bytes
207            .chunks_exact(4)
208            .map(|c| f32::from_le_bytes(c.try_into().unwrap()))
209            .collect();
210        (indices, distances)
211    };
212
213    let mut neighbors = Vec::with_capacity(n);
214    for i in 0..n {
215        let start = i * k;
216        let end = start + k;
217        neighbors.push(NeighborList {
218            indices: indices[start..end].to_vec(),
219            distances: distances[start..end].to_vec(),
220        });
221    }
222
223    Ok(KnnGraph {
224        neighbors,
225        n,
226        k,
227        metric,
228        provenance,
229    })
230}
231
232#[cfg(test)]
233mod tests {
234    use super::*;
235    use crate::{KnnMethod, compute_knn};
236    use std::env::temp_dir;
237
238    fn make_grid(n: usize) -> Vec<f32> {
239        (0..n).flat_map(|i| [i as f32, 0.0]).collect()
240    }
241
242    #[test]
243    fn round_trip_preserves_graph() {
244        let data = make_grid(12);
245        let graph = compute_knn(
246            &data,
247            12,
248            2,
249            4,
250            &KnnMethod::Exact,
251            DistanceMetric::Euclidean,
252        )
253        .unwrap();
254        let path = temp_dir().join(format!(
255            "flow-knn-roundtrip-{}.bin",
256            std::process::id()
257        ));
258        write_knn_graph(&path, &graph).unwrap();
259        let loaded = read_knn_graph(&path).unwrap();
260        let _ = std::fs::remove_file(&path);
261
262        assert_eq!(loaded.n, graph.n);
263        assert_eq!(loaded.k, graph.k);
264        assert_eq!(loaded.metric, graph.metric);
265        assert_eq!(loaded.provenance, graph.provenance);
266        for (a, b) in loaded.neighbors.iter().zip(graph.neighbors.iter()) {
267            assert_eq!(a.indices, b.indices);
268            assert_eq!(a.distances, b.distances);
269        }
270    }
271
272    #[test]
273    fn rejects_bad_magic() {
274        let path = temp_dir().join(format!("flow-knn-bad-magic-{}.bin", std::process::id()));
275        std::fs::write(&path, b"XXXX\0\0\0\x01").unwrap();
276        let err = read_knn_graph(&path).unwrap_err();
277        let _ = std::fs::remove_file(&path);
278        assert!(matches!(err, KnnError::Io(_)));
279        assert!(err.to_string().contains("magic"));
280    }
281
282    #[test]
283    fn rejects_bad_version() {
284        let path = temp_dir().join(format!("flow-knn-bad-ver-{}.bin", std::process::id()));
285        let mut bytes = Vec::new();
286        bytes.extend_from_slice(MAGIC);
287        bytes.extend_from_slice(&99u32.to_le_bytes());
288        std::fs::write(&path, &bytes).unwrap();
289        let err = read_knn_graph(&path).unwrap_err();
290        let _ = std::fs::remove_file(&path);
291        assert!(err.to_string().contains("version"));
292    }
293
294    #[test]
295    fn round_trip_preserves_metric_and_provenance() {
296        let data = make_grid(8);
297        let mut graph = compute_knn(
298            &data,
299            8,
300            2,
301            2,
302            &KnnMethod::Exact,
303            DistanceMetric::Manhattan,
304        )
305        .unwrap();
306        graph.provenance = Some("unit-test".to_string());
307        let path = temp_dir().join(format!("flow-knn-metric-{}.bin", std::process::id()));
308        write_knn_graph(&path, &graph).unwrap();
309        let loaded = read_knn_graph(&path).unwrap();
310        let _ = std::fs::remove_file(&path);
311        assert_eq!(loaded.metric, DistanceMetric::Manhattan);
312        assert_eq!(loaded.provenance.as_deref(), Some("unit-test"));
313    }
314}
flow_knn/io.rs

flow_knn/
io.rs
