fix(format): write chunk indexes over the max extent, swizzled for EA

The writer indexed chunks by their position in the current shape, the
same mistake the reader had. With a finite maxshape larger than the shape
the Fixed Array was sized for the shape, so libhdf5 looked up chunks past
its end ("addr overflow"); with the unlimited dimension anywhere but first,
e.g. maxshape (20, None), libhdf5 swizzles that dimension to the slowest
position and read our Extensible Array scrambled. Two unlimited dimensions
produced a file libhdf5 refused to open ("already found unlimited
dimension").

Chunks are now placed with the shared chunk_grid linearisation: Fixed
Array slots cover every chunk of the maximum extent (unwritten ones
undefined), Extensible Array indexes are swizzled, Single Chunk is only
used when the maximum extent is one chunk, and a maxshape that is smaller
than the shape, has more than one unlimited dimension, or would need an
absurd Fixed Array is an error instead of a bad file.
build_chunked_data_from_precompressed now returns a Result.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
osobh
2026-09-25 21:12:47 -05:00
co-authored by Claude Opus 5.5
parent 44f5f8b5c5
commit 540fa08907
4 changed files with 243 additions and 16 deletions
+124 -13
View File
@@ -8,6 +8,7 @@ use alloc::{vec, vec::Vec};
use crate::checksum::jenkins_lookup3;
use crate::chunk_cache::{CACHE_LINE_SIZE, align_to_cache_line};
use crate::chunk_grid::ChunkGrid;
use crate::ea_writer;
use crate::error::FormatError;
use crate::filter_pipeline::{
@@ -699,7 +700,8 @@ pub fn build_chunked_data_from_precompressed(
pre: &PrecompressedChunks,
base_address: u64,
maxshape: Option<&[u64]>,
) -> ChunkedDataResult {
) -> Result<ChunkedDataResult, FormatError> {
let index = ChunkIndexPlan::new(&pre.shape, maxshape, &pre.chunk_dims)?;
let offset_size: u8 = 8;
let length_size: u8 = 8;
let num_chunks = pre.chunks.len();
@@ -725,16 +727,15 @@ pub fn build_chunked_data_from_precompressed(
}
let chunk_dims_u32: Vec<u32> = pre.chunk_dims.iter().map(|&d| d as u32).collect();
let use_extensible = maxshape.is_some_and(|ms| ms.contains(&u64::MAX));
let aligned_idx = align_to_cache_line(data_buf.len());
if aligned_idx > data_buf.len() {
data_buf.resize(aligned_idx, 0u8);
}
let layout_message = if use_extensible {
let layout_message = if let ChunkIndexPlan::ExtensibleArray(grid) = &index {
let ea_address = base_address + data_buf.len() as u64;
let slots: Vec<Option<WrittenChunk>> = written_chunks.iter().cloned().map(Some).collect();
let slots = index_slots(grid, &pre.shape, &pre.chunk_dims, &written_chunks, None)?;
let ea_bytes = ea_writer::build_extensible_array_at(
&slots,
offset_size,
@@ -749,7 +750,7 @@ pub fn build_chunked_data_from_precompressed(
offset_size,
element_size as u32,
)
} else if num_chunks == 1 {
} else if matches!(index, ChunkIndexPlan::SingleChunk) {
let chunk_addr = written_chunks[0].address;
let filtered_size = if pre.has_filters {
Some(written_chunks[0].compressed_size)
@@ -765,9 +766,15 @@ pub fn build_chunked_data_from_precompressed(
offset_size,
element_size as u32,
)
} else {
} else if let ChunkIndexPlan::FixedArray(grid, nslots) = &index {
let fa_address = base_address + data_buf.len() as u64;
let slots: Vec<Option<WrittenChunk>> = written_chunks.iter().cloned().map(Some).collect();
let slots = index_slots(
grid,
&pre.shape,
&pre.chunk_dims,
&written_chunks,
Some(*nslots),
)?;
let fa_bytes = build_fixed_array_at(
&slots,
offset_size,
@@ -783,15 +790,123 @@ pub fn build_chunked_data_from_precompressed(
element_size as u32,
FA_PAGE_BITS,
)
} else {
unreachable!("every chunk index plan is handled above")
};
ChunkedDataResult {
Ok(ChunkedDataResult {
data_bytes: data_buf,
layout_message,
pipeline_message: pre.pipeline_message.clone(),
})
}
/// Most slots a Fixed Array index may have before we refuse to build it: its
/// data block holds one element per chunk of the *maximum* extent, so a huge
/// finite maxshape with small chunks would otherwise exhaust memory.
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.
enum ChunkIndexPlan {
SingleChunk,
/// The grid and the number of array elements (chunks of the max extent).
FixedArray(ChunkGrid, usize),
ExtensibleArray(ChunkGrid),
}
impl ChunkIndexPlan {
fn new(
shape: &[u64],
maxshape: Option<&[u64]>,
chunk_dims: &[u64],
) -> Result<Self, FormatError> {
let bad = |what: &str| FormatError::ChunkedReadError(format!("maxshape: {what}"));
if let Some(ms) = maxshape {
if ms.len() != shape.len() {
return Err(bad("rank differs from the shape"));
}
if ms.iter().zip(shape).any(|(&m, &s)| m < s) {
return Err(bad("smaller than the shape"));
}
}
let max = maxshape.unwrap_or(shape);
let nunlim = max.iter().filter(|&&d| d == u64::MAX).count();
match nunlim {
0 => {
let nslots = max
.iter()
.zip(chunk_dims)
.try_fold(1u64, |acc, (&m, &c)| acc.checked_mul(m.div_ceil(c.max(1))))
.filter(|&n| n <= MAX_FIXED_ARRAY_SLOTS)
.ok_or_else(|| {
bad("too many chunks for a Fixed Array index; \
use larger chunks or an unlimited dimension")
})?;
// A Single Chunk index needs that one chunk to exist; an
// empty dataset gets an all-unallocated Fixed Array instead.
let empty = shape.contains(&0);
if nslots == 1 && !empty {
Ok(Self::SingleChunk)
} else {
let grid = ChunkGrid::fixed_array(shape, Some(max), chunk_dims)?;
Ok(Self::FixedArray(grid, nslots as usize))
}
}
1 => Ok(Self::ExtensibleArray(ChunkGrid::extensible_array(
shape,
Some(max),
chunk_dims,
)?)),
_ => Err(bad(
"more than one unlimited dimension needs a B-tree v2 chunk index, \
which the writer does not support",
)),
}
}
}
/// Place each written chunk at its linear index in `grid`. `chunks` are in
/// row-major order over the chunks of the current extent (`split_into_chunks`).
/// `len` fixes the slot count (Fixed Array); otherwise it is one past the
/// highest index used.
fn index_slots(
grid: &ChunkGrid,
shape: &[u64],
chunk_dims: &[u64],
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 idx = usize::try_from(grid.linear_index(&scaled))
.map_err(|_| FormatError::Overflow("chunk index slot".into()))?;
placed.push((idx, chunk));
}
let n = len.unwrap_or_else(|| placed.iter().map(|&(i, _)| i + 1).max().unwrap_or(0));
let mut slots = vec![None; n];
for (idx, chunk) in placed {
*slots
.get_mut(idx)
.ok_or_else(|| FormatError::Overflow("chunk index slot".into()))? = Some(chunk.clone());
}
Ok(slots)
}
/// Build chunked data with absolute addresses.
/// If `maxshape` has unlimited dims, uses Extensible Array index.
pub fn build_chunked_data_at(
@@ -824,11 +939,7 @@ pub fn build_chunked_data_at_ext(
maxshape: Option<&[u64]>,
) -> Result<ChunkedDataResult, FormatError> {
let pre = precompress_chunks(raw_data, shape, chunk_dims, element_size, options)?;
Ok(build_chunked_data_from_precompressed(
&pre,
base_address,
maxshape,
))
build_chunked_data_from_precompressed(&pre, base_address, maxshape)
}
/// Write selected elements into an existing in-memory dataset buffer.