writer: v2 B-trees with internal nodes (no 65 535-record limit)
Dense link and attribute indexes and the chunk index for several unlimited dimensions were single leaves, capping them at 65 535 records. btree_v2_write builds trees of any depth, with node capacities and pointer widths from libhdf5's H5B2__hdr_init arithmetic (now shared with the reader as btree_v2::node_info) and libhdf5's node sizes (512 dense, 2048 chunks). Indexes that fit the old one-leaf layout are written byte for byte as before (compared for 10..65 535 links, attrs and chunks, tracked and filtered). Tests: 100 000 links (short names; long names with creation order), 70 000 attributes, 200 000 chunks (and 80 000 deflated), read by h5py, h5dump and clawhdf5 and edited by h5py r+; h5rs check on the same shapes, asserting depths 2-3. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
osobh
@@ -0,0 +1,396 @@
|
||||
//! Writing version-2 B-trees: a header (`BTHD`) and its nodes, leaves
|
||||
//! (`BTLF`) and, for more records than one leaf holds, internal nodes
|
||||
//! (`BTIN`) to any depth.
|
||||
//!
|
||||
//! Node capacities come from [`crate::btree_v2::node_info`], the arithmetic
|
||||
//! libhdf5 uses (`H5B2__hdr_init`) and the reader decodes pointers with, so
|
||||
//! the pointer widths the writer encodes are the ones every reader expects.
|
||||
|
||||
#[cfg(not(feature = "std"))]
|
||||
use alloc::{format, vec, vec::Vec};
|
||||
|
||||
use crate::btree_v2::{NodeInfo, bytes_for_max_records, node_info};
|
||||
use crate::checksum::jenkins_lookup3;
|
||||
use crate::error::FormatError;
|
||||
|
||||
/// How a B-tree is laid out: its record type and node geometry, as the
|
||||
/// header records them.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub(crate) struct BTreeV2Params {
|
||||
/// Record type (5: link names, 6: link creation order, 8: attribute
|
||||
/// names, 9: attribute creation order, 10/11: chunks).
|
||||
pub(crate) tree_type: u8,
|
||||
/// Bytes per node.
|
||||
pub(crate) node_size: u32,
|
||||
/// Bytes per record.
|
||||
pub(crate) record_size: u16,
|
||||
/// Split and merge percentages. The writer fills nodes itself; these
|
||||
/// only tell libhdf5 when to split and merge as it modifies the tree.
|
||||
pub(crate) split_percent: u8,
|
||||
pub(crate) merge_percent: u8,
|
||||
}
|
||||
|
||||
/// Size of a B-tree v2 header.
|
||||
pub(crate) fn header_size(offset_size: u8, length_size: u8) -> usize {
|
||||
4 + 1 + 1 + 4 + 2 + 2 + 1 + 1 + offset_size as usize + 2 + length_size as usize + 4
|
||||
}
|
||||
|
||||
/// Deepest tree the writer builds. Even at the smallest fan-out libhdf5's
|
||||
/// arithmetic allows, a few levels hold more records than any file could.
|
||||
const MAX_WRITE_DEPTH: u16 = 32;
|
||||
|
||||
/// Write a B-tree v2 holding `records` (`record_size` bytes each,
|
||||
/// concatenated, already in the tree's key order) at `addr`: the header,
|
||||
/// then its nodes, each `node_size` bytes. No records gives a header with
|
||||
/// an undefined root.
|
||||
///
|
||||
/// The tree is as shallow as the node size allows: a single leaf when the
|
||||
/// records fit one, otherwise internal nodes above leaves. Records are
|
||||
/// spread evenly over each node's children, so every node but the root is
|
||||
/// at least about half full (above libhdf5's merge threshold, which is below
|
||||
/// half), and each node holds at most its depth's maximum.
|
||||
pub(crate) fn build_btree_v2(
|
||||
p: BTreeV2Params,
|
||||
records: &[u8],
|
||||
addr: u64,
|
||||
offset_size: u8,
|
||||
length_size: u8,
|
||||
) -> Result<Vec<u8>, FormatError> {
|
||||
let rs = usize::from(p.record_size);
|
||||
if rs == 0 || !records.len().is_multiple_of(rs) {
|
||||
return Err(FormatError::SerializationError(format!(
|
||||
"B-tree v2 records are {} bytes, not a multiple of the record size {rs}",
|
||||
records.len()
|
||||
)));
|
||||
}
|
||||
let n = (records.len() / rs) as u64;
|
||||
let hdr_len = header_size(offset_size, length_size);
|
||||
|
||||
// The shallowest depth whose subtree can hold every record.
|
||||
let mut info = node_info(p.node_size, p.record_size, offset_size, 0);
|
||||
let max_leaf = info[0].max_nrec;
|
||||
if max_leaf == 0 || max_leaf > u64::from(u16::MAX) {
|
||||
return Err(FormatError::SerializationError(format!(
|
||||
"a {}-byte B-tree v2 node holds {max_leaf} {}-byte records; \
|
||||
a node holds 1 to 65535",
|
||||
p.node_size, p.record_size
|
||||
)));
|
||||
}
|
||||
let mut depth = 0u16;
|
||||
while info[usize::from(depth)].cum_max_nrec < n {
|
||||
depth += 1;
|
||||
if depth > MAX_WRITE_DEPTH {
|
||||
return Err(FormatError::SerializationError(format!(
|
||||
"{n} records do not fit a B-tree v2 of {}-byte nodes",
|
||||
p.node_size
|
||||
)));
|
||||
}
|
||||
info = node_info(p.node_size, p.record_size, offset_size, depth);
|
||||
let max = info[usize::from(depth)].max_nrec;
|
||||
if max == 0 || max > u64::from(u16::MAX) {
|
||||
return Err(FormatError::SerializationError(format!(
|
||||
"a {}-byte B-tree v2 internal node holds {max} records; \
|
||||
a node holds 1 to 65535",
|
||||
p.node_size
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
let mut w = TreeWriter {
|
||||
p,
|
||||
records,
|
||||
info: &info,
|
||||
nrec_width: bytes_for_max_records(max_leaf),
|
||||
offset_size,
|
||||
first_node: addr + hdr_len as u64,
|
||||
nodes: Vec::new(),
|
||||
};
|
||||
let root = (n > 0).then(|| w.node(depth, 0, n as usize)).transpose()?;
|
||||
|
||||
let mut out = Vec::with_capacity(hdr_len + w.nodes.len() * p.node_size as usize);
|
||||
out.extend_from_slice(b"BTHD");
|
||||
out.push(0); // version
|
||||
out.push(p.tree_type);
|
||||
out.extend_from_slice(&p.node_size.to_le_bytes());
|
||||
out.extend_from_slice(&p.record_size.to_le_bytes());
|
||||
out.extend_from_slice(&depth.to_le_bytes());
|
||||
out.push(p.split_percent);
|
||||
out.push(p.merge_percent);
|
||||
match root {
|
||||
Some(r) => push_uint(&mut out, r.addr, offset_size as usize),
|
||||
None => out.extend(core::iter::repeat_n(0xFF, offset_size as usize)),
|
||||
}
|
||||
let root_nrec = root.map_or(0, |r| r.nrec);
|
||||
out.extend_from_slice(&(root_nrec as u16).to_le_bytes());
|
||||
push_uint(&mut out, n, length_size as usize);
|
||||
let sum = jenkins_lookup3(&out);
|
||||
out.extend_from_slice(&sum.to_le_bytes());
|
||||
debug_assert_eq!(out.len(), hdr_len);
|
||||
for node in &w.nodes {
|
||||
out.extend_from_slice(node);
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// A written node, as its parent points at it.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct NodeRef {
|
||||
addr: u64,
|
||||
/// Records in the node itself.
|
||||
nrec: u64,
|
||||
/// Records in the subtree it roots.
|
||||
all_nrec: u64,
|
||||
}
|
||||
|
||||
struct TreeWriter<'a> {
|
||||
p: BTreeV2Params,
|
||||
records: &'a [u8],
|
||||
info: &'a [NodeInfo],
|
||||
/// Width of a child's record count: what a leaf's maximum needs.
|
||||
nrec_width: usize,
|
||||
offset_size: u8,
|
||||
/// Address of the first node (right after the header).
|
||||
first_node: u64,
|
||||
/// Nodes in file order (children before their parent).
|
||||
nodes: Vec<Vec<u8>>,
|
||||
}
|
||||
|
||||
impl TreeWriter<'_> {
|
||||
fn record(&self, i: usize) -> &[u8] {
|
||||
let rs = usize::from(self.p.record_size);
|
||||
&self.records[i * rs..(i + 1) * rs]
|
||||
}
|
||||
|
||||
fn push_node(&mut self, mut node: Vec<u8>) -> u64 {
|
||||
// The checksum covers the node up to it, not the padding after.
|
||||
let sum = jenkins_lookup3(&node);
|
||||
node.extend_from_slice(&sum.to_le_bytes());
|
||||
debug_assert!(node.len() <= self.p.node_size as usize);
|
||||
node.resize(self.p.node_size as usize, 0);
|
||||
let addr = self.first_node + self.nodes.len() as u64 * u64::from(self.p.node_size);
|
||||
self.nodes.push(node);
|
||||
addr
|
||||
}
|
||||
|
||||
/// Write the subtree of `depth` holding records `first..first + n`.
|
||||
fn node(&mut self, depth: u16, first: usize, n: usize) -> Result<NodeRef, FormatError> {
|
||||
let rs = usize::from(self.p.record_size);
|
||||
let mut node = Vec::with_capacity(self.p.node_size as usize);
|
||||
if depth == 0 {
|
||||
debug_assert!(n as u64 <= self.info[0].max_nrec);
|
||||
node.extend_from_slice(b"BTLF");
|
||||
node.push(0); // version
|
||||
node.push(self.p.tree_type);
|
||||
node.extend_from_slice(&self.records[first * rs..(first + n) * rs]);
|
||||
let addr = self.push_node(node);
|
||||
return Ok(NodeRef {
|
||||
addr,
|
||||
nrec: n as u64,
|
||||
all_nrec: n as u64,
|
||||
});
|
||||
}
|
||||
|
||||
// As few children as hold the records, at least two, with the
|
||||
// records spread evenly: `k` children and `k - 1` records between
|
||||
// them.
|
||||
let below = self.info[usize::from(depth) - 1].cum_max_nrec;
|
||||
let k = (n as u64 + 1).div_ceil(below + 1).max(2);
|
||||
let max = self.info[usize::from(depth)].max_nrec;
|
||||
if k - 1 > max || (n as u64) < k - 1 + k {
|
||||
return Err(FormatError::SerializationError(format!(
|
||||
"cannot spread {n} B-tree v2 records over {k} children at depth {depth}"
|
||||
)));
|
||||
}
|
||||
let k = k as usize;
|
||||
let in_children = n - (k - 1);
|
||||
let (base, extra) = (in_children / k, in_children % k);
|
||||
|
||||
let mut children = Vec::with_capacity(k);
|
||||
let mut separators = Vec::with_capacity(k - 1);
|
||||
let mut next = first;
|
||||
for c in 0..k {
|
||||
let m = base + usize::from(c < extra);
|
||||
children.push(self.node(depth - 1, next, m)?);
|
||||
next += m;
|
||||
if c + 1 < k {
|
||||
separators.push(next);
|
||||
next += 1;
|
||||
}
|
||||
}
|
||||
debug_assert_eq!(next, first + n);
|
||||
|
||||
node.extend_from_slice(b"BTIN");
|
||||
node.push(0); // version
|
||||
node.push(self.p.tree_type);
|
||||
for &s in &separators {
|
||||
node.extend_from_slice(self.record(s));
|
||||
}
|
||||
let total_width = if depth > 1 {
|
||||
self.info[usize::from(depth) - 1].cum_max_nrec_size
|
||||
} else {
|
||||
0
|
||||
};
|
||||
for c in &children {
|
||||
push_uint(&mut node, c.addr, self.offset_size as usize);
|
||||
push_uint(&mut node, c.nrec, self.nrec_width);
|
||||
if depth > 1 {
|
||||
push_uint(&mut node, c.all_nrec, total_width);
|
||||
}
|
||||
}
|
||||
let addr = self.push_node(node);
|
||||
Ok(NodeRef {
|
||||
addr,
|
||||
nrec: (k - 1) as u64,
|
||||
all_nrec: n as u64,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Append `v` as a `width`-byte little-endian integer.
|
||||
fn push_uint(buf: &mut Vec<u8>, v: u64, width: usize) {
|
||||
let bytes = v.to_le_bytes();
|
||||
buf.extend_from_slice(&bytes[..width.min(8)]);
|
||||
buf.extend(vec![0u8; width.saturating_sub(8)]);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::btree_v2::{BTreeV2Header, collect_btree_v2_records};
|
||||
|
||||
fn params(node_size: u32, record_size: u16) -> BTreeV2Params {
|
||||
BTreeV2Params {
|
||||
tree_type: 5,
|
||||
node_size,
|
||||
record_size,
|
||||
split_percent: 100,
|
||||
merge_percent: 40,
|
||||
}
|
||||
}
|
||||
|
||||
/// `n` 11-byte records: a big-endian counter, so byte order is key order.
|
||||
fn records(n: usize, rs: usize) -> Vec<u8> {
|
||||
let mut out = Vec::with_capacity(n * rs);
|
||||
for i in 0..n {
|
||||
let mut r = vec![0u8; rs];
|
||||
r[..8].copy_from_slice(&(i as u64).to_be_bytes());
|
||||
out.extend_from_slice(&r);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn roundtrip(node_size: u32, rs: u16, n: usize, os: u8, ls: u8) -> BTreeV2Header {
|
||||
let recs = records(n, usize::from(rs));
|
||||
let base = 4096u64;
|
||||
let tree = build_btree_v2(params(node_size, rs), &recs, base, os, ls).unwrap();
|
||||
let mut file = vec![0u8; base as usize];
|
||||
file.extend_from_slice(&tree);
|
||||
let hdr = BTreeV2Header::parse(&file, base as usize, os, ls).unwrap();
|
||||
assert_eq!(hdr.total_records, n as u64);
|
||||
let got = collect_btree_v2_records(&file, &hdr, os, ls).unwrap();
|
||||
assert_eq!(got.len(), n);
|
||||
let flat: Vec<u8> = got.into_iter().flat_map(|r| r.data).collect();
|
||||
assert_eq!(flat, recs, "node {node_size} rs {rs} n {n}");
|
||||
hdr
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn one_leaf_then_deeper_trees_read_back_in_order() {
|
||||
// 512-byte nodes of 11-byte records: 45 per leaf, 1149 at depth 1,
|
||||
// 26 449 at depth 2.
|
||||
let info = node_info(512, 11, 8, 3);
|
||||
assert_eq!(
|
||||
info.iter().map(|i| i.cum_max_nrec).collect::<Vec<_>>(),
|
||||
[45, 1149, 26_449, 608_349]
|
||||
);
|
||||
for (n, depth) in [
|
||||
(0, 0),
|
||||
(1, 0),
|
||||
(45, 0),
|
||||
(46, 1),
|
||||
(1149, 1),
|
||||
(1150, 2),
|
||||
(26_449, 2),
|
||||
(26_450, 3),
|
||||
(100_000, 3),
|
||||
] {
|
||||
let hdr = roundtrip(512, 11, n, 8, 8);
|
||||
assert_eq!(hdr.depth, depth, "{n} records");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pointer_widths_follow_the_offset_and_length_sizes() {
|
||||
for (os, ls) in [(4, 4), (8, 4), (4, 8), (2, 2)] {
|
||||
roundtrip(512, 11, 5000, os, ls);
|
||||
}
|
||||
// Wide counts: a leaf of 2048 bytes / 9-byte records (226, one byte)
|
||||
// and deeper subtree totals of three bytes.
|
||||
roundtrip(2048, 9, 300_000, 8, 8);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn every_node_is_within_its_capacity_and_above_the_merge_threshold() {
|
||||
let rs = 17u16;
|
||||
let n = 70_000usize;
|
||||
let info = node_info(512, rs, 8, 3);
|
||||
let recs = records(n, usize::from(rs));
|
||||
let tree = build_btree_v2(params(512, rs), &recs, 0, 8, 8).unwrap();
|
||||
let hdr_len = header_size(8, 8);
|
||||
let nodes = (tree.len() - hdr_len) / 512;
|
||||
for i in 0..nodes {
|
||||
let node = &tree[hdr_len + i * 512..hdr_len + (i + 1) * 512];
|
||||
let sig = &node[..4];
|
||||
if sig == b"BTLF" {
|
||||
continue; // counts checked through the parents below
|
||||
}
|
||||
assert_eq!(sig, b"BTIN");
|
||||
}
|
||||
// Walk from the header: each child's count within [40%, 100%].
|
||||
let hdr = BTreeV2Header::parse(&tree, 0, 8, 8).unwrap();
|
||||
assert_eq!(hdr.depth, 3);
|
||||
assert!(u64::from(hdr.num_records_in_root) <= info[3].max_nrec);
|
||||
fn walk(tree: &[u8], addr: usize, nrec: usize, depth: usize, info: &[NodeInfo], rs: usize) {
|
||||
if depth == 0 {
|
||||
return;
|
||||
}
|
||||
let nrec_w = bytes_for_max_records(info[0].max_nrec);
|
||||
let tot_w = if depth > 1 {
|
||||
info[depth - 1].cum_max_nrec_size
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let mut pos = addr + 6 + nrec * rs;
|
||||
for _ in 0..=nrec {
|
||||
let a = u64::from_le_bytes(tree[pos..pos + 8].try_into().unwrap()) as usize;
|
||||
pos += 8;
|
||||
let mut c = 0usize;
|
||||
for b in 0..nrec_w {
|
||||
c |= usize::from(tree[pos + b]) << (8 * b);
|
||||
}
|
||||
pos += nrec_w + tot_w;
|
||||
let max = info[depth - 1].max_nrec as usize;
|
||||
assert!(c <= max && c * 100 > max * 40, "{c} of {max}");
|
||||
walk(tree, a, c, depth - 1, info, rs);
|
||||
}
|
||||
}
|
||||
walk(
|
||||
&tree,
|
||||
hdr.root_node_address as usize,
|
||||
usize::from(hdr.num_records_in_root),
|
||||
3,
|
||||
&info,
|
||||
usize::from(rs),
|
||||
);
|
||||
assert!(nodes > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_node_too_small_or_too_big_is_an_error() {
|
||||
assert!(build_btree_v2(params(16, 11), &records(1, 11), 0, 8, 8).is_err());
|
||||
// A leaf with room for more than 65 535 records.
|
||||
assert!(build_btree_v2(params(1 << 20, 11), &records(1, 11), 0, 8, 8).is_err());
|
||||
// Records that are not whole.
|
||||
assert!(build_btree_v2(params(512, 11), &[0u8; 12], 0, 8, 8).is_err());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user