fix(format): index datasets with several unlimited dims by B-tree v2

A dataset with more than one unlimited dimension got an Extensible Array
index, which libhdf5 refuses ("already found unlimited dimension"), so
the whole file failed to open in h5py and h5dump. The previous commit
turned that into a write error; this one writes what the library itself
uses there: a version-2 B-tree chunk index (record type 10/11), as a
single leaf of the library's 2048-byte node size, or a larger leaf when
the records do not fit. The root's record count is 16-bit, so more than
65535 chunks is still refused rather than written wrong.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
osobh
2026-09-25 21:14:47 -05:00
co-authored by Claude Opus 5.5
parent 540fa08907
commit 1dba7b465a
2 changed files with 268 additions and 99 deletions
+188 -38
View File
@@ -444,6 +444,27 @@ fn serialize_v4_fixed_array(
element_size: u32,
max_bits: u8,
) -> Vec<u8> {
let mut buf = layout_v4_chunked_prefix(chunk_dims, element_size);
// chunk index type = 3 (Fixed Array)
buf.push(3);
// max_dblk_page_nelmts_bits — must match FAHD max_nelmts_bits
buf.push(max_bits);
// Fixed Array header address
match offset_size {
4 => buf.extend_from_slice(&(fixed_array_address as u32).to_le_bytes()),
8 => buf.extend_from_slice(&fixed_array_address.to_le_bytes()),
_ => {}
}
buf
}
/// The part of a v4 chunked layout message before the chunk index type:
/// version, class, flags and the chunk dimensions (plus the element size).
fn layout_v4_chunked_prefix(chunk_dims: &[u32], element_size: u32) -> Vec<u8> {
let mut buf = Vec::new();
buf.push(4); // version
buf.push(2); // class = chunked
@@ -483,20 +504,6 @@ fn serialize_v4_fixed_array(
4 => buf.extend_from_slice(&element_size.to_le_bytes()),
_ => {}
}
// chunk index type = 3 (Fixed Array)
buf.push(3);
// max_dblk_page_nelmts_bits — must match FAHD max_nelmts_bits
buf.push(max_bits);
// Fixed Array header address
match offset_size {
4 => buf.extend_from_slice(&(fixed_array_address as u32).to_le_bytes()),
8 => buf.extend_from_slice(&fixed_array_address.to_le_bytes()),
_ => {}
}
buf
}
@@ -733,7 +740,8 @@ pub fn build_chunked_data_from_precompressed(
data_buf.resize(aligned_idx, 0u8);
}
let layout_message = if let ChunkIndexPlan::ExtensibleArray(grid) = &index {
let layout_message = match &index {
ChunkIndexPlan::ExtensibleArray(grid) => {
let ea_address = base_address + data_buf.len() as u64;
let slots = index_slots(grid, &pre.shape, &pre.chunk_dims, &written_chunks, None)?;
let ea_bytes = ea_writer::build_extensible_array_at(
@@ -750,7 +758,8 @@ pub fn build_chunked_data_from_precompressed(
offset_size,
element_size as u32,
)
} else if matches!(index, ChunkIndexPlan::SingleChunk) {
}
ChunkIndexPlan::SingleChunk => {
let chunk_addr = written_chunks[0].address;
let filtered_size = if pre.has_filters {
Some(written_chunks[0].compressed_size)
@@ -766,7 +775,8 @@ pub fn build_chunked_data_from_precompressed(
offset_size,
element_size as u32,
)
} else if let ChunkIndexPlan::FixedArray(grid, nslots) = &index {
}
ChunkIndexPlan::FixedArray(grid, nslots) => {
let fa_address = base_address + data_buf.len() as u64;
let slots = index_slots(
grid,
@@ -790,8 +800,31 @@ pub fn build_chunked_data_from_precompressed(
element_size as u32,
FA_PAGE_BITS,
)
} else {
unreachable!("every chunk index plan is handled above")
}
ChunkIndexPlan::BTreeV2 => {
let bt_address = base_address + data_buf.len() as u64;
let records: Vec<(Vec<u64>, &WrittenChunk)> = written_chunks
.iter()
.enumerate()
.map(|(i, c)| (scaled_coords(&pre.shape, &pre.chunk_dims, i), c))
.collect();
let (bt_bytes, node_size) = build_btree_v2_chunk_index_at(
pre.shape.len(),
&records,
offset_size,
length_size,
pre.has_filters,
bt_address,
)?;
data_buf.extend_from_slice(&bt_bytes);
serialize_v4_btree_v2(
&chunk_dims_u32,
bt_address,
offset_size,
element_size as u32,
node_size,
)
}
};
Ok(ChunkedDataResult {
@@ -807,14 +840,15 @@ pub fn build_chunked_data_from_precompressed(
const MAX_FIXED_ARRAY_SLOTS: u64 = 1 << 26;
/// Which chunk index a dataset gets, following the library's choice in
/// `H5D__layout_set_latest_indexing`: Extensible Array for exactly one
/// unlimited dimension, Fixed Array for a finite maxshape, Single Chunk when
/// the whole maximum extent is one chunk.
/// `H5D__layout_set_latest_indexing`: version-2 B-tree for more than one
/// unlimited dimension, Extensible Array for exactly one, Fixed Array for a
/// finite maxshape, Single Chunk when the whole maximum extent is one chunk.
enum ChunkIndexPlan {
SingleChunk,
/// The grid and the number of array elements (chunks of the max extent).
FixedArray(ChunkGrid, usize),
ExtensibleArray(ChunkGrid),
BTreeV2,
}
impl ChunkIndexPlan {
@@ -860,10 +894,7 @@ impl ChunkIndexPlan {
Some(max),
chunk_dims,
)?)),
_ => Err(bad(
"more than one unlimited dimension needs a B-tree v2 chunk index, \
which the writer does not support",
)),
_ => Ok(Self::BTreeV2),
}
}
}
@@ -879,20 +910,9 @@ fn index_slots(
chunks: &[WrittenChunk],
len: Option<usize>,
) -> Result<Vec<Option<WrittenChunk>>, FormatError> {
let rank = shape.len();
let cur: Vec<u64> = shape
.iter()
.zip(chunk_dims)
.map(|(&s, &c)| s.div_ceil(c))
.collect();
let mut placed: Vec<(usize, &WrittenChunk)> = Vec::with_capacity(chunks.len());
let mut scaled = vec![0u64; rank];
for (i, chunk) in chunks.iter().enumerate() {
let mut rem = i as u64;
for d in (0..rank).rev() {
scaled[d] = rem % cur[d];
rem /= cur[d];
}
let scaled = scaled_coords(shape, chunk_dims, i);
let idx = usize::try_from(grid.linear_index(&scaled))
.map_err(|_| FormatError::Overflow("chunk index slot".into()))?;
placed.push((idx, chunk));
@@ -907,6 +927,136 @@ fn index_slots(
Ok(slots)
}
/// Scaled coordinates (`offset / chunk_dim`) of the `i`-th chunk in the
/// row-major order `split_into_chunks` produces over the current extent.
fn scaled_coords(shape: &[u64], chunk_dims: &[u64], i: usize) -> Vec<u64> {
let rank = shape.len();
let mut scaled = vec![0u64; rank];
let mut rem = i as u64;
for d in (0..rank).rev() {
let n = shape[d].div_ceil(chunk_dims[d]);
scaled[d] = rem % n;
rem /= n;
}
scaled
}
/// Node size the library gives a chunk index B-tree (`H5D_BT2_NODE_SIZE`),
/// with its split and merge percentages.
const BT2_NODE_SIZE: u32 = 2048;
const BT2_SPLIT_PERCENT: u8 = 100;
const BT2_MERGE_PERCENT: u8 = 40;
/// B-tree v2 record types for chunk indexes (`H5B2_CDSET_ID`,
/// `H5B2_CDSET_FILT_ID`).
const BT2_CHUNK_UNFILTERED: u8 = 10;
const BT2_CHUNK_FILTERED: u8 = 11;
/// Build a version-2 B-tree chunk index (the library's index for datasets
/// with more than one unlimited dimension) at a known absolute address.
///
/// `records` are `(scaled coordinates, chunk)` in lexicographic order of the
/// coordinates, which is the order the library's comparator
/// (`H5VM_vector_cmp_u`) keeps them in. The tree is a single leaf: the
/// library's 2048-byte node when the records fit, otherwise a leaf node
/// sized to hold them all (the root's record count is 16-bit, so at most
/// 65535 chunks). Returns the bytes and the node size the layout message
/// must record.
fn build_btree_v2_chunk_index_at(
rank: usize,
records: &[(Vec<u64>, &WrittenChunk)],
offset_size: u8,
length_size: u8,
has_filters: bool,
base_address: u64,
) -> Result<(Vec<u8>, u32), FormatError> {
let os = offset_size as usize;
let nrec = u16::try_from(records.len()).map_err(|_| {
FormatError::ChunkedReadError(
"more than 65535 chunks with more than one unlimited dimension: \
use larger chunks"
.into(),
)
})?;
let chunk_size_bytes = has_filters.then(|| {
let slots: Vec<Option<WrittenChunk>> =
records.iter().map(|(_, c)| Some((*c).clone())).collect();
filtered_chunk_size_len(&slots)
});
let record_size = os + chunk_size_bytes.map_or(0, |n| n + 4) + 8 * rank;
// Leaf: signature, version, type, records, checksum.
let leaf_len = 4 + 1 + 1 + records.len() * record_size + 4;
let node_size = u32::try_from(leaf_len)
.map_err(|_| FormatError::Overflow("B-tree v2 leaf size".into()))?
.max(BT2_NODE_SIZE);
let tree_type = if has_filters {
BT2_CHUNK_FILTERED
} else {
BT2_CHUNK_UNFILTERED
};
let hdr_len = 4 + 1 + 1 + 4 + 2 + 2 + 1 + 1 + os + 2 + length_size as usize + 4;
let leaf_address = base_address + hdr_len as u64;
let mut out = Vec::with_capacity(hdr_len + node_size as usize);
out.extend_from_slice(b"BTHD");
out.push(0); // version
out.push(tree_type);
out.extend_from_slice(&node_size.to_le_bytes());
out.extend_from_slice(&(record_size as u16).to_le_bytes());
out.extend_from_slice(&0u16.to_le_bytes()); // depth
out.push(BT2_SPLIT_PERCENT);
out.push(BT2_MERGE_PERCENT);
if records.is_empty() {
out.extend(core::iter::repeat_n(0xFF, os));
} else {
push_addr(&mut out, leaf_address, offset_size);
}
out.extend_from_slice(&nrec.to_le_bytes());
match length_size {
4 => out.extend_from_slice(&(records.len() as u32).to_le_bytes()),
_ => out.extend_from_slice(&(records.len() as u64).to_le_bytes()),
}
let sum = jenkins_lookup3(&out);
out.extend_from_slice(&sum.to_le_bytes());
debug_assert_eq!(out.len(), hdr_len);
if records.is_empty() {
return Ok((out, node_size));
}
let leaf_start = out.len();
out.extend_from_slice(b"BTLF");
out.push(0); // version
out.push(tree_type);
for (scaled, chunk) in records {
push_index_element(&mut out, Some(chunk), offset_size, chunk_size_bytes);
for &c in scaled {
out.extend_from_slice(&c.to_le_bytes());
}
}
let sum = jenkins_lookup3(&out[leaf_start..]);
out.extend_from_slice(&sum.to_le_bytes());
// The library reads whole nodes; pad the leaf out to the node size.
out.resize(leaf_start + node_size as usize, 0);
Ok((out, node_size))
}
/// Serialize a v4 layout message for a version-2 B-tree chunk index.
fn serialize_v4_btree_v2(
chunk_dims: &[u32],
btree_address: u64,
offset_size: u8,
element_size: u32,
node_size: u32,
) -> Vec<u8> {
let mut buf = layout_v4_chunked_prefix(chunk_dims, element_size);
buf.push(5); // chunk index type = 5 (version-2 B-tree)
buf.extend_from_slice(&node_size.to_le_bytes());
buf.push(BT2_SPLIT_PERCENT);
buf.push(BT2_MERGE_PERCENT);
push_addr(&mut buf, btree_address, offset_size);
buf
}
/// Build chunked data with absolute addresses.
/// If `maxshape` has unlimited dims, uses Extensible Array index.
pub fn build_chunked_data_at(
+26 -7
View File
@@ -515,16 +515,35 @@ fn we_write_maxshape_larger_than_shape() {
check_we_write(&cases);
}
/// More than one unlimited dimension needs a B-tree v2 chunk index; the
/// writer must not produce a file libhdf5 cannot open.
/// More than one unlimited dimension needs a version-2 B-tree chunk index,
/// as the library uses; an Extensible Array for `(None, None)` made libhdf5
/// refuse the whole file ("already found unlimited dimension").
#[test]
fn two_unlimited_dims_are_refused() {
fn we_write_btree_v2_for_several_unlimited_dims() {
const U: u64 = u64::MAX;
let mut cases = vec![
wcase("unl_unl", &[20, 30], &[5, 5], Some(&[U, U])),
wcase("unl_fin_unl", &[6, 7, 8], &[4, 3, 5], Some(&[U, 9, U])),
// More records than the library's 2048-byte node holds (84 here).
wcase("unl_unl_2400", &[40, 60], &[1, 1], Some(&[U, U])),
wcase("unl_unl_empty", &[0, 0], &[4, 4], Some(&[U, U])),
];
let mut filtered = wcase("unl_unl_deflate", &[6, 7, 8], &[4, 3, 5], Some(&[U, U, U]));
filtered.deflate = true;
cases.push(filtered);
check_we_write(&cases);
}
/// A single-leaf B-tree has a 16-bit record count; beyond it the writer
/// refuses rather than writing a tree libhdf5 would misread.
#[test]
fn btree_v2_index_past_one_leaf_is_refused() {
let mut b = FileBuilder::new();
b.create_dataset("d")
.with_i32_data(&(0..600).collect::<Vec<i32>>())
.with_shape(&[20, 30])
.with_chunks(&[5, 5])
.with_i32_data(&vec![0i32; 70_000])
.with_shape(&[70_000, 1])
.with_chunks(&[1, 1])
.with_maxshape(&[u64::MAX, u64::MAX]);
let dir = tempfile::tempdir().unwrap();
assert!(b.write(dir.path().join("unl_unl.h5")).is_err());
assert!(b.write(dir.path().join("too_many.h5")).is_err());
}