urna_runtime/graph/
csr.rs

1//! `CsrIndex`: the runtime reader over the graph_adjacency (0x0C) csr
2//! payload. parses the offsets + neighbor columns off the section bytes
3//! once at open time, then exposes `neighbors(node)` as a single contiguous
4//! slice walked with zero allocation, mirroring `ann::HnswIndex::from_bytes`.
5//! typed errors, never panics on a truncated mmap; the graph only ever
6//! GENERATES candidates, it never scores.
7
8use urna_format::error::UrnaError;
9use urna_format::sections::graph::parse_csr_parts;
10
11use crate::error::RuntimeError;
12
13/// flat csr adjacency over chunk ordinals. `offsets[node..node+1]` bounds
14/// node's run in `neighbors` (and `edge_types`), so one node's neighbors are
15/// a contiguous `&[u32]` slice.
16pub struct CsrIndex {
17    n_nodes: usize,
18    /// row pointers, len `n_nodes + 1`.
19    offsets: Vec<u32>,
20    /// absolute (decoded) neighbor ids, concatenated per node.
21    neighbors: Vec<u32>,
22    /// edge type per neighbor, aligned 1:1 with `neighbors`.
23    edge_types: Vec<u8>,
24}
25
26impl CsrIndex {
27    /// parse a graph_adjacency csr payload. `n_embeddings` is the corpus
28    /// chunk count; the csr's `n_nodes` must match it so neighbor ids index
29    /// the same ordinal space as the embeddings.
30    pub fn from_bytes(bytes: &[u8], n_embeddings: usize) -> Result<Self, RuntimeError> {
31        let parts = parse_csr_parts(bytes).map_err(RuntimeError::Format)?;
32        if parts.n_nodes != n_embeddings {
33            return Err(RuntimeError::Format(UrnaError::MalformedSectionPayload {
34                section_id: urna_format::layout::SECTION_GRAPH_ADJACENCY,
35                reason: format!(
36                    "graph node count {} != n_embeddings {}",
37                    parts.n_nodes, n_embeddings
38                ),
39            }));
40        }
41        // narrow the validated u64 columns to u32 (ids index a u32 ordinal
42        // space; parse_csr_parts already range-checks dst < n_nodes and the
43        // monotone offsets, so these casts cannot lose information).
44        let offsets: Vec<u32> = parts.offsets.iter().map(|&o| o as u32).collect();
45        Ok(Self {
46            n_nodes: parts.n_nodes,
47            offsets,
48            neighbors: parts.neighbors,
49            edge_types: parts.edge_types,
50        })
51    }
52
53    pub fn n_nodes(&self) -> usize {
54        self.n_nodes
55    }
56
57    /// node `node`'s out-neighbors as a contiguous slice. empty (not an
58    /// error) for an out-of-range node so the bfs can probe freely.
59    #[inline]
60    pub fn neighbors(&self, node: usize) -> &[u32] {
61        if node + 1 >= self.offsets.len() {
62            return &[];
63        }
64        let start = self.offsets[node] as usize;
65        let end = self.offsets[node + 1] as usize;
66        &self.neighbors[start..end]
67    }
68
69    /// edge type of node's `i`-th out-edge (0-based within the node's run).
70    #[inline]
71    pub fn edge_type(&self, node: usize, i: usize) -> Option<u8> {
72        if node + 1 >= self.offsets.len() {
73            return None;
74        }
75        let start = self.offsets[node] as usize;
76        let end = self.offsets[node + 1] as usize;
77        let idx = start + i;
78        if idx < end {
79            Some(self.edge_types[idx])
80        } else {
81            None
82        }
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use urna_format::sections::graph::{
90        EDGE_TYPE_NEXT_CHUNK, EDGE_TYPE_SEMANTIC, Edge, encode_graph_adjacency,
91    };
92
93    fn edge(src: u32, dst: u32, t: u8) -> Edge {
94        Edge {
95            src,
96            dst,
97            edge_type: t,
98        }
99    }
100
101    #[test]
102    fn neighbors_are_contiguous_and_correct() {
103        // node 0 -> {1,2}, node 1 -> {2}, node 2 -> {}, node 3 -> {0}.
104        let edges = vec![
105            edge(0, 2, EDGE_TYPE_SEMANTIC),
106            edge(0, 1, EDGE_TYPE_NEXT_CHUNK),
107            edge(1, 2, EDGE_TYPE_NEXT_CHUNK),
108            edge(3, 0, EDGE_TYPE_NEXT_CHUNK),
109        ];
110        let payload = encode_graph_adjacency(4, &edges).unwrap();
111        let idx = CsrIndex::from_bytes(&payload, 4).unwrap();
112        assert_eq!(idx.n_nodes(), 4);
113        // canonical sort is (src, edge_type, dst): node 0's edges become
114        // [NEXT_CHUNK->1, SEMANTIC->2].
115        assert_eq!(idx.neighbors(0), &[1, 2]);
116        assert_eq!(idx.neighbors(1), &[2]);
117        assert_eq!(idx.neighbors(2), &[] as &[u32]);
118        assert_eq!(idx.neighbors(3), &[0]);
119        assert_eq!(idx.edge_type(0, 0), Some(EDGE_TYPE_NEXT_CHUNK));
120        assert_eq!(idx.edge_type(0, 1), Some(EDGE_TYPE_SEMANTIC));
121        assert_eq!(idx.edge_type(0, 2), None);
122        // out-of-range node is empty, not a panic.
123        assert_eq!(idx.neighbors(99), &[] as &[u32]);
124    }
125
126    #[test]
127    fn rejects_node_count_mismatch() {
128        let edges = vec![edge(0, 1, EDGE_TYPE_NEXT_CHUNK)];
129        let payload = encode_graph_adjacency(2, &edges).unwrap();
130        assert!(CsrIndex::from_bytes(&payload, 3).is_err());
131    }
132
133    #[test]
134    fn truncated_payload_errors_never_panics() {
135        let edges = vec![
136            edge(0, 1, EDGE_TYPE_NEXT_CHUNK),
137            edge(1, 2, EDGE_TYPE_SEMANTIC),
138            edge(2, 0, EDGE_TYPE_SEMANTIC),
139        ];
140        let payload = encode_graph_adjacency(3, &edges).unwrap();
141        for cut in 0..payload.len() {
142            // every prefix must error or parse, never panic.
143            let _ = CsrIndex::from_bytes(&payload[..cut], 3);
144        }
145    }
146
147    #[test]
148    fn empty_graph_parses() {
149        let payload = encode_graph_adjacency(0, &[]).unwrap();
150        let idx = CsrIndex::from_bytes(&payload, 0).unwrap();
151        assert_eq!(idx.n_nodes(), 0);
152        assert_eq!(idx.neighbors(0), &[] as &[u32]);
153    }
154}
urna_runtime/graph/csr.rs

urna_runtime/graph/
csr.rs
