uqa_storage/diskann_index/format/
layout.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7use super::{invalid, NODE_HEADER_BYTES, PAGE_PAYLOAD_BYTES};
8use crate::StorageBackendResult;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub struct DiskANNNodeLayout {
12    pub(super) dimensions: u32,
13    pub(super) max_degree: usize,
14    pub(super) node_count: u64,
15    pub(super) slot_bytes: usize,
16    nodes_per_page: u32,
17    fragments: u32,
18    page_count: u64,
19}
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub struct DiskANNNodeAddress {
23    pub first_page: u64,
24    pub slot: u32,
25    pub fragments: u32,
26}
27
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub struct DiskANNPageShape {
30    pub first_node: u64,
31    pub slots: u32,
32    pub fragment_index: u32,
33    pub fragments: u32,
34    pub payload_bytes: u32,
35}
36
37impl DiskANNNodeLayout {
38    pub fn new(dimensions: u32, max_degree: usize, node_count: u64) -> StorageBackendResult<Self> {
39        if dimensions == 0 || max_degree < 2 {
40            return Err(invalid(
41                "require positive dimensions and maximum degree at least two",
42            ));
43        }
44        let dimensions_usize =
45            usize::try_from(dimensions).map_err(|_| invalid("dimension range"))?;
46        let slot_bytes = max_degree
47            .checked_mul(8)
48            .and_then(|neighbors| {
49                dimensions_usize
50                    .checked_mul(4)
51                    .and_then(|vectors| vectors.checked_add(neighbors))
52            })
53            .and_then(|body| body.checked_add(NODE_HEADER_BYTES))
54            .ok_or_else(|| invalid("node slot size overflow"))?;
55        std::alloc::Layout::array::<u8>(slot_bytes)
56            .map_err(|_| invalid("node slot exceeds addressable memory"))?;
57        let fragments = u32::try_from(slot_bytes.div_ceil(PAGE_PAYLOAD_BYTES))
58            .map_err(|_| invalid("fragment count exceeds format"))?;
59        let nodes_per_page = (PAGE_PAYLOAD_BYTES / slot_bytes) as u32;
60        let page_count = if nodes_per_page == 0 {
61            node_count
62                .checked_mul(u64::from(fragments))
63                .ok_or_else(|| invalid("page count overflow"))?
64        } else {
65            node_count.div_ceil(u64::from(nodes_per_page))
66        };
67        Ok(Self {
68            dimensions,
69            max_degree,
70            node_count,
71            slot_bytes,
72            nodes_per_page,
73            fragments,
74            page_count,
75        })
76    }
77
78    pub fn dimensions(self) -> u32 {
79        self.dimensions
80    }
81    pub fn max_degree(self) -> usize {
82        self.max_degree
83    }
84    pub fn node_count(self) -> u64 {
85        self.node_count
86    }
87    pub fn slot_bytes(self) -> usize {
88        self.slot_bytes
89    }
90    pub fn page_count(self) -> u64 {
91        self.page_count
92    }
93
94    pub fn node_address(self, node: u64) -> StorageBackendResult<DiskANNNodeAddress> {
95        if node >= self.node_count {
96            return Err(invalid("node ID outside generation"));
97        }
98        if self.nodes_per_page > 0 {
99            let count = u64::from(self.nodes_per_page);
100            Ok(DiskANNNodeAddress {
101                first_page: node / count,
102                slot: (node % count) as u32,
103                fragments: 1,
104            })
105        } else {
106            Ok(DiskANNNodeAddress {
107                first_page: node * u64::from(self.fragments),
108                slot: 0,
109                fragments: self.fragments,
110            })
111        }
112    }
113
114    pub fn page_shape(self, page: u64) -> StorageBackendResult<DiskANNPageShape> {
115        if page >= self.page_count {
116            return Err(invalid("page ID outside generation"));
117        }
118        if self.nodes_per_page > 0 {
119            let first_node = page * u64::from(self.nodes_per_page);
120            let slots = (self.node_count - first_node).min(u64::from(self.nodes_per_page)) as u32;
121            Ok(DiskANNPageShape {
122                first_node,
123                slots,
124                fragment_index: 0,
125                fragments: 1,
126                payload_bytes: (slots as usize * self.slot_bytes) as u32,
127            })
128        } else {
129            let fragments = u64::from(self.fragments);
130            let fragment_index = (page % fragments) as u32;
131            let offset = fragment_index as usize * PAGE_PAYLOAD_BYTES;
132            Ok(DiskANNPageShape {
133                first_node: page / fragments,
134                slots: 1,
135                fragment_index,
136                fragments: self.fragments,
137                payload_bytes: (self.slot_bytes - offset).min(PAGE_PAYLOAD_BYTES) as u32,
138            })
139        }
140    }
141}
uqa_storage/diskann_index/format/layout.rs

uqa_storage/diskann_index/format/
layout.rs
