Skip to main content

engramdb_core/
layout.rs

1//! 布局:行宽、badge、分片划分(与真实 checkpoint 结构一致的参数化小实现)。
2
3/// 物理布局描述:分片化的定长行存储(如 PLE: 128 shards × [2_500_012, 160] FP8)。
4#[derive(Debug, Clone, Copy, PartialEq)]
5pub struct Layout {
6    pub shards: u64,
7    pub rows_per_shard: u64,
8    pub width: u64,
9    pub row_bytes: u64,  // width * elem_bytes
10    pub badge_rows: u64, // 每个 badge 的行数(用户可调,对齐 4KB 优先)
11}
12
13impl Layout {
14    pub fn new(shards: u64, rows_per_shard: u64, width: u64, elem_bytes: u64) -> Self {
15        let row_bytes = width * elem_bytes;
16        let badge_rows = aligned_badge_rows(row_bytes, 4096);
17        Self {
18            shards,
19            rows_per_shard,
20            width,
21            row_bytes,
22            badge_rows,
23        }
24    }
25
26    pub fn total_rows(&self) -> u64 {
27        self.shards * self.rows_per_shard
28    }
29
30    /// badge 字节数(4KB 对齐推荐)
31    pub fn badge_bytes(&self) -> u64 {
32        self.badge_rows * self.row_bytes
33    }
34
35    /// rowid → (shard_id, badge_id, in_badge_row)
36    pub fn locate(&self, rowid: u64) -> (u64, u64, u64) {
37        let shard = rowid / self.rows_per_shard;
38        let in_shard = rowid % self.rows_per_shard;
39        let badge = in_shard / self.badge_rows;
40        let in_badge = in_shard % self.badge_rows;
41        (shard, badge, in_badge)
42    }
43
44    /// 一维张量模式:总表 = 单列大数组(当 shards == 1 时)
45    pub fn byte_offset(&self, rowid: u64) -> u64 {
46        rowid * self.row_bytes
47    }
48}
49
50/// 行宽 -> 尽量使 badge ≥ 4KB 且行数整除性佳的 badge_rows。
51pub fn aligned_badge_rows(row_bytes: u64, min_bytes: u64) -> u64 {
52    (min_bytes / row_bytes.max(1)).max(1)
53}
54
55/// 行范围(为批量连续读服务)
56#[derive(Debug, Clone, Copy)]
57pub struct RowExtent {
58    pub rows: u64,
59    pub row_bytes: u64,
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65
66    #[test]
67    fn locate_matches_struct() {
68        // 真实 PLE 规格的 1/16 mock:8 shards × 156_251 行 × 160 宽
69        let l = Layout::new(8, 156_251, 160, 1);
70        assert_eq!(l.total_rows(), 8 * 156_251);
71        assert_eq!(l.row_bytes, 160);
72        // 4KB / 160B = 25.6 → 25 行/badge
73        assert_eq!(l.badge_rows, 25);
74        assert_eq!(l.badge_bytes(), 25 * 160);
75        let (s, b, r) = l.locate(156_253);
76        assert_eq!((s, b, r), (1, 0, 2));
77    }
78
79    #[test]
80    fn row_bytes_bf16() {
81        let l = Layout::new(8, 156_251, 160, 2);
82        assert_eq!(l.row_bytes, 320);
83        // 4096/320=12.8 → 12 行
84        assert_eq!(l.badge_rows, 12);
85    }
86}