diff --git a/crates/clawhdf5-format/src/chunk_grid.rs b/crates/clawhdf5-format/src/chunk_grid.rs index c5a06fd..9e03b9d 100644 --- a/crates/clawhdf5-format/src/chunk_grid.rs +++ b/crates/clawhdf5-format/src/chunk_grid.rs @@ -147,7 +147,6 @@ impl ChunkGrid { /// Linear index of the chunk with scaled coordinates `scaled` /// (`offset / chunk_dim` per dimension, in dataset order). - #[allow(dead_code)] // used by the writer pub(crate) fn linear_index(&self, scaled: &[u64]) -> u64 { self.order .iter() diff --git a/crates/clawhdf5-format/src/chunked_write.rs b/crates/clawhdf5-format/src/chunked_write.rs index 66bc82d..3136f58 100644 --- a/crates/clawhdf5-format/src/chunked_write.rs +++ b/crates/clawhdf5-format/src/chunked_write.rs @@ -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 { + 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 = 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> = 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> = 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 { + 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, +) -> Result>, FormatError> { + let rank = shape.len(); + let cur: Vec = 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 { 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. diff --git a/crates/clawhdf5-format/src/file_writer.rs b/crates/clawhdf5-format/src/file_writer.rs index c262cd4..5350c88 100644 --- a/crates/clawhdf5-format/src/file_writer.rs +++ b/crates/clawhdf5-format/src/file_writer.rs @@ -1244,7 +1244,7 @@ impl FileWriter { &pre, dummy_cursor, d.maxshape.as_deref(), - ); + )?; dummy_cursor += result.data_bytes.len() as u64; let dense_blob = if ds_dense[i] { Some(build_dense_attrs(&d.attrs, 0)) @@ -1424,7 +1424,7 @@ impl FileWriter { .expect("chunked dataset missing precompressed cache"), base_address, d.maxshape.as_deref(), - ); + )?; cursor2 += result.data_bytes.len(); let oh = build_chunked_dataset_oh( &d.dt, diff --git a/crates/clawhdf5/tests/chunk_index_interop.rs b/crates/clawhdf5/tests/chunk_index_interop.rs index 0e7c845..41d7a5a 100644 --- a/crates/clawhdf5/tests/chunk_index_interop.rs +++ b/crates/clawhdf5/tests/chunk_index_interop.rs @@ -373,6 +373,73 @@ fn check_we_write(cases: &[WriteCase]) { "h5dump failed: {stderr}" ); } + + // Let libhdf5 grow every resizable dataset by two chunks per dimension + // (capped at the maxshape) and rewrite it, which updates our index in + // place and inserts new chunks into it. Then both readers must agree. + let script = format!( + r#" +import h5py, numpy as np +grown = {{}} +with h5py.File(r'{path_str}', 'r+') as f: + for name in f: + d = f[name] + if d.chunks is None: + continue + new = tuple(s + 2 * c if m is None else min(m, s + 2 * c) + for s, m, c in zip(d.shape, d.maxshape, d.chunks)) + if new == d.shape: + continue + old = d[()] + full = np.full(new, -7, 'i4') + full[tuple(slice(0, s) for s in old.shape)] = old + d.resize(new) + d[...] = full + grown[name] = (list(old.shape), list(new)) +with h5py.File(r'{path_str}', 'r') as f: + for name, (old, new) in grown.items(): + want = np.full(new, -7, 'i4') + want[tuple(slice(0, s) for s in old)] = np.arange(int(np.prod(old)), dtype='i4').reshape(old) + assert np.array_equal(f[name][()], want), name +for name, (old, new) in grown.items(): + print(name, ','.join(map(str, old)), ','.join(map(str, new))) +"# + ); + let out = run_python(&script); + let growable = cases + .iter() + .filter(|c| c.maxshape.as_ref().is_some_and(|m| *m != c.shape)) + .count(); + assert_eq!(out.lines().count(), growable, "libhdf5 grew: {out}"); + let dims = |s: &str| -> Vec { s.split(',').map(|x| x.parse().unwrap()).collect() }; + let file = File::open(&path).unwrap(); + for line in out.lines() { + let mut parts = line.split(' '); + let (name, old, new) = ( + parts.next().unwrap(), + dims(parts.next().unwrap()), + dims(parts.next().unwrap()), + ); + let got = file.dataset(name).unwrap().read_i32().unwrap(); + let n: usize = new.iter().product(); + let mut want = vec![-7i32; n]; + for (flat, w) in want.iter_mut().enumerate() { + let mut rem = flat; + let mut coords = vec![0usize; new.len()]; + for d in (0..new.len()).rev() { + coords[d] = rem % new[d]; + rem /= new[d]; + } + if coords.iter().zip(&old).all(|(c, o)| c < o) { + *w = coords.iter().zip(&old).fold(0, |acc, (c, o)| acc * o + c) as i32; + } + } + let bad = got.iter().zip(&want).filter(|(a, b)| a != b).count(); + assert!( + got.len() == n && bad == 0, + "{name}: after libhdf5 grew it, our reader got {bad} of {n} values wrong" + ); + } } /// A Fixed Array with more than 1024 elements must be paged, or libhdf5 @@ -411,3 +478,53 @@ fn we_write_extensible_array_past_index_block() { cases.push(wcase("ea_140000", &[140_000], &[1], Some(unl))); check_we_write(&cases); } + +/// A maxshape larger than the shape: the index must be laid out over the +/// chunks of the maximum extent (libhdf5 read our Fixed Array past its end: +/// "addr overflow"), and an Extensible Array whose unlimited dimension is not +/// the first must swizzle it to the slowest position (libhdf5 read our +/// `(20, None)` dataset scrambled). +#[test] +fn we_write_maxshape_larger_than_shape() { + const U: u64 = u64::MAX; + let mut cases = vec![ + // Fixed Array over the maximum extent. + wcase("fa2d_finite_max", &[20, 30], &[5, 5], Some(&[40, 60])), + wcase("fa1d_finite_max", &[40], &[4], Some(&[100])), + wcase("fa3d_edges", &[6, 7, 8], &[4, 3, 5], Some(&[10, 9, 20])), + wcase("fa_paged_max", &[30, 50], &[1, 1], Some(&[40, 60])), + wcase("fa_one_chunk_now", &[5], &[5], Some(&[50])), + // Extensible Array, unlimited dimension first (no swizzle) ... + wcase("ea2d_unl_fin", &[20, 30], &[5, 5], Some(&[U, 30])), + wcase("ea2d_unl_fin_max", &[20, 30], &[5, 5], Some(&[U, 60])), + // ... and not first (swizzled). + wcase("ea2d_fin_unl", &[20, 30], &[5, 5], Some(&[20, U])), + wcase("ea2d_fin_max_unl", &[20, 30], &[5, 5], Some(&[40, U])), + wcase("ea3d_mid", &[6, 7, 8], &[4, 3, 5], Some(&[10, U, 20])), + // Past the index block and into super blocks, swizzled. + wcase("ea2d_many", &[3, 2000], &[1, 1], Some(&[4, U])), + ]; + let mut filtered = wcase( + "ea3d_last_deflate", + &[6, 7, 8], + &[4, 3, 5], + Some(&[6, 8, U]), + ); + filtered.deflate = true; + cases.push(filtered); + 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. +#[test] +fn two_unlimited_dims_are_refused() { + let mut b = FileBuilder::new(); + b.create_dataset("d") + .with_i32_data(&(0..600).collect::>()) + .with_shape(&[20, 30]) + .with_chunks(&[5, 5]) + .with_maxshape(&[u64::MAX, u64::MAX]); + let dir = tempfile::tempdir().unwrap(); + assert!(b.write(dir.path().join("unl_unl.h5")).is_err()); +}