1#[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, pub badge_rows: u64, }
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 pub fn badge_bytes(&self) -> u64 {
32 self.badge_rows * self.row_bytes
33 }
34
35 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 pub fn byte_offset(&self, rowid: u64) -> u64 {
46 rowid * self.row_bytes
47 }
48}
49
50pub fn aligned_badge_rows(row_bytes: u64, min_bytes: u64) -> u64 {
52 (min_bytes / row_bytes.max(1)).max(1)
53}
54
55#[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 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 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 assert_eq!(l.badge_rows, 12);
85 }
86}