precinct 0.2.0

Approximate nearest-neighbor search over region embeddings
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

//! Simple binary I/O for region embeddings.
//!
//! Format (little-endian):
//! - Header: `n: u32`, `dim: u32`
//! - For each of n regions: `id: u32`, `min: [f32; dim]`, `max: [f32; dim]`
//!
//! This is intentionally minimal -- no compression, no versioning.
//! For production use cases, prefer serde with a format of your choice.

use crate::{AxisBox, Region};
use std::io::{self, Read, Write};

/// Write a collection of boxes to a binary stream.
pub fn write_boxes<W: Write>(
    writer: &mut W,
    ids: &[u32],
    boxes: &[AxisBox],
) -> io::Result<()> {
    assert_eq!(ids.len(), boxes.len());
    if boxes.is_empty() {
        writer.write_all(&0u32.to_le_bytes())?;
        writer.write_all(&0u32.to_le_bytes())?;
        return Ok(());
    }

    let n = boxes.len() as u32;
    let dim = boxes[0].dim() as u32;
    writer.write_all(&n.to_le_bytes())?;
    writer.write_all(&dim.to_le_bytes())?;

    for (id, b) in ids.iter().zip(boxes.iter()) {
        writer.write_all(&id.to_le_bytes())?;
        for &v in b.min() {
            writer.write_all(&v.to_le_bytes())?;
        }
        for &v in b.max() {
            writer.write_all(&v.to_le_bytes())?;
        }
    }

    Ok(())
}

/// Read boxes from a binary stream.
///
/// Returns `(ids, boxes)`.
pub fn read_boxes<R: Read>(reader: &mut R) -> io::Result<(Vec<u32>, Vec<AxisBox>)> {
    let mut buf4 = [0u8; 4];

    reader.read_exact(&mut buf4)?;
    let n = u32::from_le_bytes(buf4) as usize;
    reader.read_exact(&mut buf4)?;
    let dim = u32::from_le_bytes(buf4) as usize;

    let mut ids = Vec::with_capacity(n);
    let mut boxes = Vec::with_capacity(n);

    let mut float_buf = vec![0u8; dim * 4];

    for _ in 0..n {
        reader.read_exact(&mut buf4)?;
        ids.push(u32::from_le_bytes(buf4));

        reader.read_exact(&mut float_buf)?;
        let min: Vec<f32> = float_buf
            .chunks_exact(4)
            .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
            .collect();

        reader.read_exact(&mut float_buf)?;
        let max: Vec<f32> = float_buf
            .chunks_exact(4)
            .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
            .collect();

        boxes.push(AxisBox::new(min, max));
    }

    Ok((ids, boxes))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn roundtrip() {
        let ids = vec![0, 1, 2];
        let boxes = vec![
            AxisBox::new(vec![0.0, 0.0], vec![1.0, 1.0]),
            AxisBox::new(vec![2.0, 3.0], vec![4.0, 5.0]),
            AxisBox::new(vec![-1.0, -1.0], vec![0.0, 0.0]),
        ];

        let mut buf = Vec::new();
        write_boxes(&mut buf, &ids, &boxes).unwrap();

        let (read_ids, read_boxes) = read_boxes(&mut buf.as_slice()).unwrap();
        assert_eq!(ids, read_ids);
        assert_eq!(read_boxes.len(), 3);

        for (orig, read) in boxes.iter().zip(read_boxes.iter()) {
            assert_eq!(orig.min(), read.min());
            assert_eq!(orig.max(), read.max());
        }
    }

    #[test]
    fn empty() {
        let mut buf = Vec::new();
        write_boxes(&mut buf, &[], &[]).unwrap();
        let (ids, boxes) = read_boxes(&mut buf.as_slice()).unwrap();
        assert!(ids.is_empty());
        assert!(boxes.is_empty());
    }
}