Skip to main content

erigon_seg/writer/
bt_writer.rs

1//! Writer for `.bt` B-tree index files, in both on-disk layouts.
2//!
3//! The index records the `.kv` byte offset of every key as an Elias-Fano array. Two
4//! layouts are produced (selectable via [`BtLayout`]):
5//!
6//! * [`BtLayout::Legacy`] — just the serialized Elias-Fano (`[EF]`). The reader and
7//!   erigon both treat trailing B-tree nodes as optional.
8//! * [`BtLayout::Footer`] — `[0x01][nodes][pad→4096][EF][pad→8][footer][anchor]`, where
9//!   `nodes` holds the key at every `M`-th position (`keyLen:u16-BE | key`) for
10//!   co-located binary search, and the trailing footer/anchor carry `keys_count`, `M`,
11//!   `ef_offset`, and the `erigon\0\0` magic. Port of erigon `BtIndexWriter`.
12
13use std::fs::File;
14use std::io::Write;
15use std::path::Path;
16
17use super::ef_builder::EfBuilder;
18use crate::error::{Error, Result};
19use crate::seg::Seg;
20
21/// Default B-tree fanout (`DefaultBtreeM`), the number of keys per leaf.
22pub const DEFAULT_BTREE_M: u64 = 256;
23
24const BT_EF_ALIGN: usize = 4096;
25const BT_FOOTER_ALIGN: usize = 8;
26const BT_VERSION: u16 = 1;
27const BT_METADATA_LEN: u32 = 24;
28const FOOTER_MAGIC: [u8; 8] = *b"erigon\x00\x00";
29const FIRST_BYTE_FOOTER: u8 = 0x01;
30
31/// Which `.bt` on-disk layout to emit.
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub enum BtLayout {
34    /// Just the Elias-Fano offset array (smallest; binary search only).
35    Legacy,
36    /// Footer layout with di-nodes for co-located reads (erigon's current default).
37    Footer,
38}
39
40/// Options for building a `.bt` index.
41#[derive(Debug, Clone, Copy)]
42pub struct BtOptions {
43    /// Which layout to write.
44    pub layout: BtLayout,
45    /// B-tree fanout `M` (footer layout only).
46    pub m: u64,
47}
48
49impl Default for BtOptions {
50    fn default() -> BtOptions {
51        BtOptions {
52            layout: BtLayout::Footer,
53            m: DEFAULT_BTREE_M,
54        }
55    }
56}
57
58/// Build a `.bt` index for a `.kv` file, writing it to `bt_path`.
59pub fn build_bt(
60    kv_path: impl AsRef<Path>,
61    bt_path: impl AsRef<Path>,
62    opts: BtOptions,
63) -> Result<()> {
64    let seg = Seg::open(kv_path)?;
65    build_bt_from_seg(&seg, bt_path, opts)
66}
67
68/// Build a `.bt` index from an already-open [`Seg`].
69pub fn build_bt_from_seg(seg: &Seg, bt_path: impl AsRef<Path>, opts: BtOptions) -> Result<()> {
70    let bt_path = bt_path.as_ref();
71    let key_count = seg.words_count() / 2;
72
73    // An empty domain yields an empty (0-byte) index, which the reader treats as 0 keys.
74    if key_count == 0 {
75        File::create(bt_path).map_err(|e| Error::io(bt_path, e))?;
76        return Ok(());
77    }
78
79    let max_offset = seg.len() as u64;
80    let m = opts.m.max(1);
81    let mut ef = EfBuilder::new(key_count, max_offset);
82
83    // The footer layout streams the di-nodes (keys at every M-th position) ahead of the EF.
84    let mut nodes: Vec<u8> = Vec::new();
85    let footer = opts.layout == BtLayout::Footer;
86    if footer {
87        nodes.push(FIRST_BYTE_FOOTER);
88    }
89
90    let mut g = seg.getter();
91    for di in 0..key_count {
92        let off = g.offset();
93        if footer && di % m == 0 {
94            let key = g.next(); // need the bytes for this node
95            let klen = u16::try_from(key.len()).map_err(|_| {
96                Error::format("key longer than 65535 bytes (unsupported in .bt node)")
97            })?;
98            nodes.extend_from_slice(&klen.to_be_bytes());
99            nodes.extend_from_slice(&key);
100        } else {
101            g.skip(); // key
102        }
103        g.skip(); // value
104        ef.add_offset(off);
105    }
106    ef.build();
107
108    match opts.layout {
109        BtLayout::Legacy => {
110            let mut out = Vec::with_capacity(ef.serialized_len());
111            ef.write_to(&mut out);
112            write_all(bt_path, &out)
113        }
114        BtLayout::Footer => {
115            let mut out = nodes;
116            pad_to(&mut out, BT_EF_ALIGN);
117            let ef_offset = out.len() as u64;
118            ef.write_to(&mut out);
119            pad_to(&mut out, BT_FOOTER_ALIGN);
120            // Footer payload: keys_count | M | ef_offset.
121            out.extend_from_slice(&key_count.to_be_bytes());
122            out.extend_from_slice(&m.to_be_bytes());
123            out.extend_from_slice(&ef_offset.to_be_bytes());
124            // Anchor: footer_len(u32) | flags(u16) | format_version(u16) | magic(u64).
125            out.extend_from_slice(&BT_METADATA_LEN.to_be_bytes());
126            out.extend_from_slice(&0u16.to_be_bytes()); // flags
127            out.extend_from_slice(&BT_VERSION.to_be_bytes());
128            out.extend_from_slice(&FOOTER_MAGIC);
129            write_all(bt_path, &out)
130        }
131    }
132}
133
134fn pad_to(out: &mut Vec<u8>, align: usize) {
135    let rem = out.len() % align;
136    if rem != 0 {
137        out.resize(out.len() + (align - rem), 0);
138    }
139}
140
141fn write_all(path: &Path, bytes: &[u8]) -> Result<()> {
142    let mut f = File::create(path).map_err(|e| Error::io(path, e))?;
143    f.write_all(bytes).map_err(|e| Error::io(path, e))?;
144    f.flush().map_err(|e| Error::io(path, e))
145}