uqa_storage/diskann_index/format/
node.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7use uqa_core::{memory::BudgetedVec, DocId};
8
9use super::{field, invalid, zeros, DiskANNNodeLayout, DiskANNVectorVersion, NODE_HEADER_BYTES};
10use crate::diskann_index::metric::{checkpoint, norms};
11use crate::mvcc::{DatabaseId, StorageTransactionId};
12use crate::{read_control::StorageReadControl, StorageBackendResult};
13
14pub struct DiskANNNodeInput<'a> {
15    pub node_id: u64,
16    pub doc_id: DocId,
17    pub ordinal: u32,
18    pub version: DiskANNVectorVersion,
19    pub vector: &'a [f32],
20    pub neighbors: &'a [u64],
21}
22
23#[derive(Debug)]
24pub struct DiskANNNode {
25    node_id: u64,
26    doc_id: DocId,
27    ordinal: u32,
28    version: DiskANNVectorVersion,
29    raw_norm: f32,
30    vector: BudgetedVec<f32>,
31    neighbors: BudgetedVec<u64>,
32}
33
34impl DiskANNNode {
35    pub fn node_id(&self) -> u64 {
36        self.node_id
37    }
38    pub fn doc_id(&self) -> DocId {
39        self.doc_id
40    }
41    pub fn ordinal(&self) -> u32 {
42        self.ordinal
43    }
44    pub fn version(&self) -> DiskANNVectorVersion {
45        self.version
46    }
47    pub fn raw_norm(&self) -> f32 {
48        self.raw_norm
49    }
50    pub fn vector(&self) -> &[f32] {
51        &self.vector
52    }
53    pub fn neighbors(&self) -> &[u64] {
54        &self.neighbors
55    }
56}
57
58impl DiskANNNodeLayout {
59    pub fn encode_node(
60        self,
61        node: &DiskANNNodeInput<'_>,
62        control: &StorageReadControl,
63    ) -> StorageBackendResult<BudgetedVec<u8>> {
64        control.check()?;
65        self.node_address(node.node_id)?;
66        self.check_neighbors(node.node_id, node.neighbors, control)?;
67        let raw_norm = graph_norm(self.dimensions, node.vector, control)?;
68        let mut header = [0_u8; NODE_HEADER_BYTES];
69        header[0..8].copy_from_slice(&node.node_id.to_le_bytes());
70        header[8..16].copy_from_slice(&node.doc_id.to_le_bytes());
71        header[16..20].copy_from_slice(&node.ordinal.to_le_bytes());
72        header[20..24].copy_from_slice(&raw_norm.to_bits().to_le_bytes());
73        header[24..32].copy_from_slice(&(node.neighbors.len() as u64).to_le_bytes());
74        header[32..48].copy_from_slice(&node.version.writer.database().as_bytes());
75        header[48..56].copy_from_slice(&node.version.writer.allocation().to_le_bytes());
76        header[56..64].copy_from_slice(&node.version.revision.to_le_bytes());
77        let mut bytes = BudgetedVec::new(control.memory());
78        bytes.reserve(self.slot_bytes)?;
79        bytes.extend_from_slice(&header)?;
80        for (offset, value) in node.vector.iter().enumerate() {
81            checkpoint(offset, control)?;
82            bytes.extend_from_slice(&value.to_bits().to_le_bytes())?;
83        }
84        for slot in 0..self.max_degree {
85            checkpoint(slot, control)?;
86            bytes
87                .extend_from_slice(&node.neighbors.get(slot).copied().unwrap_or(0).to_le_bytes())?;
88        }
89        control.check()?;
90        Ok(bytes)
91    }
92
93    pub fn decode_node(
94        self,
95        expected_node: u64,
96        bytes: &[u8],
97        control: &StorageReadControl,
98    ) -> StorageBackendResult<DiskANNNode> {
99        control.check()?;
100        self.node_address(expected_node)?;
101        if bytes.len() != self.slot_bytes || u64::from_le_bytes(field(bytes, 0)?) != expected_node {
102            return Err(invalid("node slot length or identity differs"));
103        }
104        let degree = usize::try_from(u64::from_le_bytes(field(bytes, 24)?))
105            .map_err(|_| invalid("degree range"))?;
106        if degree > self.max_degree || degree as u64 >= self.node_count {
107            return Err(invalid("degree exceeds generation bounds"));
108        }
109        let writer = StorageTransactionId::new(
110            DatabaseId::from_bytes(field(bytes, 32)?),
111            u64::from_le_bytes(field(bytes, 48)?),
112        )
113        .map_err(|_| invalid("invalid origin writer allocation"))?;
114        let version = DiskANNVectorVersion::new(writer, u64::from_le_bytes(field(bytes, 56)?))?;
115        let vector_end = NODE_HEADER_BYTES + self.dimensions as usize * 4;
116        let mut vector = BudgetedVec::new(control.memory());
117        vector.reserve(self.dimensions as usize)?;
118        for (offset, component) in bytes[NODE_HEADER_BYTES..vector_end]
119            .as_chunks::<4>()
120            .0
121            .iter()
122            .enumerate()
123        {
124            checkpoint(offset, control)?;
125            vector.push(f32::from_bits(u32::from_le_bytes(field(component, 0)?)))?;
126        }
127        let raw_norm = graph_norm(self.dimensions, &vector, control)?;
128        if raw_norm.to_bits() != u32::from_le_bytes(field(bytes, 20)?) {
129            return Err(invalid("stored norm differs from canonical raw norm"));
130        }
131        let neighbor_end = vector_end + degree * 8;
132        zeros(&bytes[neighbor_end..], control)?;
133        let mut neighbors = BudgetedVec::new(control.memory());
134        neighbors.reserve(degree)?;
135        for (offset, neighbor) in bytes[vector_end..neighbor_end]
136            .as_chunks::<8>()
137            .0
138            .iter()
139            .enumerate()
140        {
141            checkpoint(offset, control)?;
142            neighbors.push(u64::from_le_bytes(field(neighbor, 0)?))?;
143        }
144        self.check_neighbors(expected_node, &neighbors, control)?;
145        control.check()?;
146        Ok(DiskANNNode {
147            node_id: expected_node,
148            doc_id: u64::from_le_bytes(field(bytes, 8)?),
149            ordinal: u32::from_le_bytes(field(bytes, 16)?),
150            version,
151            raw_norm,
152            vector,
153            neighbors,
154        })
155    }
156
157    fn check_neighbors(
158        self,
159        node: u64,
160        neighbors: &[u64],
161        control: &StorageReadControl,
162    ) -> StorageBackendResult<()> {
163        if neighbors.len() > self.max_degree || neighbors.len() as u64 >= self.node_count {
164            return Err(invalid("degree exceeds generation bounds"));
165        }
166        let mut previous = None;
167        for (offset, &neighbor) in neighbors.iter().enumerate() {
168            checkpoint(offset, control)?;
169            if neighbor >= self.node_count
170                || neighbor == node
171                || previous.is_some_and(|id| id >= neighbor)
172            {
173                return Err(invalid(
174                    "neighbors must be sorted, unique, in range and exclude self",
175                ));
176            }
177            previous = Some(neighbor);
178        }
179        Ok(())
180    }
181}
182
183fn graph_norm(
184    dimensions: u32,
185    raw: &[f32],
186    control: &StorageReadControl,
187) -> StorageBackendResult<f32> {
188    let (norm, _) = norms(dimensions, raw, control)?;
189    if !norm.is_finite() || norm == 0.0 {
190        return Err(invalid(
191            "non-navigable vector belongs in the exact side stream",
192        ));
193    }
194    Ok(norm)
195}
uqa_storage/diskann_index/format/node.rs

uqa_storage/diskann_index/format/
node.rs
