diff --git a/crates/clawhdf5-format/src/chunked_write.rs b/crates/clawhdf5-format/src/chunked_write.rs index 3136f58..6ff0c07 100644 --- a/crates/clawhdf5-format/src/chunked_write.rs +++ b/crates/clawhdf5-format/src/chunked_write.rs @@ -444,6 +444,27 @@ fn serialize_v4_fixed_array( element_size: u32, max_bits: u8, ) -> Vec { + 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 { 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,65 +740,91 @@ pub fn build_chunked_data_from_precompressed( data_buf.resize(aligned_idx, 0u8); } - let layout_message = if let ChunkIndexPlan::ExtensibleArray(grid) = &index { - 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( - &slots, - offset_size, - length_size, - pre.has_filters, - ea_address, - ); - data_buf.extend_from_slice(&ea_bytes); - ea_writer::serialize_v4_extensible_array( - &chunk_dims_u32, - ea_address, - offset_size, - element_size as u32, - ) - } 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) - } else { - None - }; - let filter_mask = if pre.has_filters { Some(0u32) } else { None }; - serialize_v4_single_chunk( - &chunk_dims_u32, - chunk_addr, - filtered_size, - filter_mask, - offset_size, - element_size as u32, - ) - } else if let ChunkIndexPlan::FixedArray(grid, nslots) = &index { - let fa_address = base_address + data_buf.len() as u64; - let slots = index_slots( - grid, - &pre.shape, - &pre.chunk_dims, - &written_chunks, - Some(*nslots), - )?; - let fa_bytes = build_fixed_array_at( - &slots, - offset_size, - length_size, - pre.has_filters, - fa_address, - ); - data_buf.extend_from_slice(&fa_bytes); - serialize_v4_fixed_array( - &chunk_dims_u32, - fa_address, - offset_size, - element_size as u32, - FA_PAGE_BITS, - ) - } else { - unreachable!("every chunk index plan is handled above") + 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( + &slots, + offset_size, + length_size, + pre.has_filters, + ea_address, + ); + data_buf.extend_from_slice(&ea_bytes); + ea_writer::serialize_v4_extensible_array( + &chunk_dims_u32, + ea_address, + offset_size, + element_size as u32, + ) + } + ChunkIndexPlan::SingleChunk => { + let chunk_addr = written_chunks[0].address; + let filtered_size = if pre.has_filters { + Some(written_chunks[0].compressed_size) + } else { + None + }; + let filter_mask = if pre.has_filters { Some(0u32) } else { None }; + serialize_v4_single_chunk( + &chunk_dims_u32, + chunk_addr, + filtered_size, + filter_mask, + offset_size, + element_size as u32, + ) + } + ChunkIndexPlan::FixedArray(grid, nslots) => { + let fa_address = base_address + data_buf.len() as u64; + let slots = index_slots( + grid, + &pre.shape, + &pre.chunk_dims, + &written_chunks, + Some(*nslots), + )?; + let fa_bytes = build_fixed_array_at( + &slots, + offset_size, + length_size, + pre.has_filters, + fa_address, + ); + data_buf.extend_from_slice(&fa_bytes); + serialize_v4_fixed_array( + &chunk_dims_u32, + fa_address, + offset_size, + element_size as u32, + FA_PAGE_BITS, + ) + } + ChunkIndexPlan::BTreeV2 => { + let bt_address = base_address + data_buf.len() as u64; + let records: Vec<(Vec, &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, ) -> 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 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 { + 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, &WrittenChunk)], + offset_size: u8, + length_size: u8, + has_filters: bool, + base_address: u64, +) -> Result<(Vec, 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> = + 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 { + 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( diff --git a/crates/clawhdf5/tests/chunk_index_interop.rs b/crates/clawhdf5/tests/chunk_index_interop.rs index 41d7a5a..2fb5aeb 100644 --- a/crates/clawhdf5/tests/chunk_index_interop.rs +++ b/crates/clawhdf5/tests/chunk_index_interop.rs @@ -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::>()) - .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()); }