From 05c665a898c084481cd88274ee9c6c84bc72a45a Mon Sep 17 00:00:00 2001 From: osobh Date: Sat, 19 Sep 2026 14:21:04 -0700 Subject: [PATCH] feat(format): choose chunk dimensions automatically for large datasets MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Requesting a filter without chunk dimensions made the whole dataset a single chunk. Any read, even one row, then decompresses everything, and a large dataset cannot be decoded in parallel — which also made the new partial reads pointless for such files. auto_chunk_dims keeps datasets up to 1 MiB as one chunk (unchanged behaviour) and splits larger ones by halving the dimensions in turn, so chunks keep roughly the dataset's proportions, until a chunk is at most 1 MiB — h5py's approach. An empty (unlimited, unwritten) dimension is treated as 1024. The writer passes the element size through resolve_chunk_dims_for; the old resolve_chunk_dims assumes 8-byte elements. Explicit with_chunks always wins. Interop test: h5py reads an auto-chunked 13 MB deflate dataset, sees chunks between 128 KiB and 1 MiB, and a small dataset still has one chunk. Co-Authored-By: Claude Fable 5.1 --- CHANGELOG.md | 8 ++ crates/clawhdf5-format/src/chunked_write.rs | 87 +++++++++++++++++++-- crates/clawhdf5-format/src/file_writer.rs | 4 +- crates/clawhdf5/tests/h5py_interop_tests.rs | 59 ++++++++++++++ 4 files changed, 152 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b600b8b..8e8ee83 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,14 @@ type 5 — what `libver='latest'` uses for two or more unlimited dimensions; previously "unsupported chunked layout"). The four copies of the chunk-index dispatch are now one shared function, so every read path gets it. +- **Automatic chunk sizes.** Asking for compression (or any filter) without + `with_chunks` used to store the whole dataset as one chunk, so any read had + to decompress everything and nothing could be decoded in parallel. Datasets up + to 1 MiB stay a single chunk, as before; larger ones are split by halving the + dimensions in turn until a chunk is at most 1 MiB (the approach h5py takes). + **Behaviour change:** large compressed datasets written without explicit + chunk dimensions get a different (standard, h5py-readable) layout. Explicit + `with_chunks` is unaffected. - **Out-of-range selections are errors.** They used to return data: a hyperslab past an edge came back padded with zeros, and a point whose column was out of range wrapped into the next row and returned that element. Now diff --git a/crates/clawhdf5-format/src/chunked_write.rs b/crates/clawhdf5-format/src/chunked_write.rs index 405312a..e6a9e44 100644 --- a/crates/clawhdf5-format/src/chunked_write.rs +++ b/crates/clawhdf5-format/src/chunked_write.rs @@ -49,6 +49,38 @@ pub struct ChunkOptions { pub pcodec: bool, } +/// Largest chunk the automatic choice produces, in bytes. +const AUTO_CHUNK_TARGET_BYTES: u64 = 1 << 20; + +/// Extent assumed for a dimension that is currently empty (an unlimited +/// dimension not yet written to) — the same stand-in h5py uses. +const AUTO_CHUNK_EMPTY_DIM: u64 = 1024; + +/// Choose chunk dimensions for a dataset nobody specified them for. +/// +/// Asking for compression (or any filter) without chunk dimensions used to +/// make the whole dataset one chunk. That defeats the point of chunking: any +/// read — even a single row — must decompress everything, and a large dataset +/// cannot be decompressed in parallel. Datasets up to the target size stay a +/// single chunk, exactly as before; larger ones are split by halving the +/// dimensions in turn (so chunks keep roughly the dataset's proportions, the +/// approach h5py takes) until a chunk fits the target. +pub fn auto_chunk_dims(shape: &[u64], elem_size: usize) -> Vec { + let mut dims: Vec = shape + .iter() + .map(|&d| if d == 0 { AUTO_CHUNK_EMPTY_DIM } else { d }) + .collect(); + let elem = elem_size.max(1) as u64; + let bytes = |dims: &[u64]| dims.iter().fold(elem, |acc, &d| acc.saturating_mul(d)); + let mut axis = 0; + while bytes(&dims) > AUTO_CHUNK_TARGET_BYTES && dims.iter().any(|&d| d > 1) { + let i = axis % dims.len(); + dims[i] = dims[i].div_ceil(2); + axis += 1; + } + dims +} + impl ChunkOptions { /// Whether any chunking option is enabled. pub fn is_chunked(&self) -> bool { @@ -135,11 +167,17 @@ impl ChunkOptions { /// Determine chunk dimensions, using user-specified or auto-computing. pub fn resolve_chunk_dims(&self, shape: &[u64]) -> Vec { - if let Some(ref dims) = self.chunk_dims { - dims.clone() - } else { - // Auto chunk: use the full dataset shape (single chunk) - shape.to_vec() + // Without the element size, assume 8 bytes (the widest common scalar); + // the writer uses `resolve_chunk_dims_for`. + self.resolve_chunk_dims_for(shape, 8) + } + + /// Chunk dimensions for a dataset of `shape` whose elements are `elem_size` + /// bytes: the caller's if given, otherwise chosen automatically. + pub fn resolve_chunk_dims_for(&self, shape: &[u64], elem_size: usize) -> Vec { + match self.chunk_dims { + Some(ref dims) => dims.clone(), + None => auto_chunk_dims(shape, elem_size), } } } @@ -1143,6 +1181,45 @@ mod tests { assert_eq!(dims, vec![100, 50]); } + #[test] + fn auto_chunking_splits_only_large_datasets() { + let bytes = |dims: &[u64], elem: u64| dims.iter().product::() * elem; + // Up to the target: one chunk, as before. + assert_eq!(auto_chunk_dims(&[100, 50], 8), [100, 50]); + assert_eq!(auto_chunk_dims(&[131_072], 8), [131_072]); // exactly 1 MiB + // Larger: split, keeping proportions, never above the target. + let big = auto_chunk_dims(&[4096, 2048], 8); + assert!(bytes(&big, 8) <= AUTO_CHUNK_TARGET_BYTES, "{big:?}"); + assert!(bytes(&big, 8) > AUTO_CHUNK_TARGET_BYTES / 4, "{big:?}"); + assert_eq!(big[0] / big[1], 2, "proportions kept: {big:?}"); + // Every dimension stays within the dataset and at least 1. + for shape in [ + vec![10_000_000u64], + vec![3, 5_000_000], + vec![1, 1, 9_000_000], + vec![7; 9], + ] { + let dims = auto_chunk_dims(&shape, 4); + assert!( + dims.iter().zip(&shape).all(|(c, s)| *c >= 1 && c <= s), + "{shape:?} -> {dims:?}" + ); + assert!( + bytes(&dims, 4) <= AUTO_CHUNK_TARGET_BYTES, + "{shape:?} -> {dims:?}" + ); + } + // An empty (unlimited, unwritten) dimension still gets a usable chunk. + let growable = auto_chunk_dims(&[0, 128], 8); + assert!(growable[0] >= 1 && bytes(&growable, 8) <= AUTO_CHUNK_TARGET_BYTES); + // Explicit dimensions always win. + let explicit = ChunkOptions { + chunk_dims: Some(vec![10, 10]), + ..Default::default() + }; + assert_eq!(explicit.resolve_chunk_dims_for(&[4096, 2048], 8), [10, 10]); + } + #[test] fn chunk_options_pipeline_deflate() { // Auto-shuffle is applied before compression by default (matches h5py). diff --git a/crates/clawhdf5-format/src/file_writer.rs b/crates/clawhdf5-format/src/file_writer.rs index 644c483..c57691b 100644 --- a/crates/clawhdf5-format/src/file_writer.rs +++ b/crates/clawhdf5-format/src/file_writer.rs @@ -1221,8 +1221,10 @@ impl FileWriter { precompressed: None, }); } else if is_chunked[i] { - let chunk_dims = d.chunk_options.resolve_chunk_dims(&d.ds.dimensions); let elem_size = d.dt.type_size() as usize; + let chunk_dims = d + .chunk_options + .resolve_chunk_dims_for(&d.ds.dimensions, elem_size); // Compress once in Pass 1; cache the result so Pass 2 can skip // re-compression and just rebuild the index with real addresses. let pre = precompress_chunks( diff --git a/crates/clawhdf5/tests/h5py_interop_tests.rs b/crates/clawhdf5/tests/h5py_interop_tests.rs index b5692d3..e177b72 100644 --- a/crates/clawhdf5/tests/h5py_interop_tests.rs +++ b/crates/clawhdf5/tests/h5py_interop_tests.rs @@ -979,3 +979,62 @@ with h5py.File("{path_str}", "r") as f: expected["slab"] ); } + +// --------------------------------------------------------------------------- +// clawhdf5 auto-chunks a large compressed dataset -> h5py reads +// --------------------------------------------------------------------------- + +/// Compression without explicit chunk dimensions used to store the whole +/// dataset as a single chunk. Large datasets are now split automatically; +/// h5py must read the result and see sensibly sized chunks. +#[test] +fn clawhdf5_auto_chunked_dataset_h5py_reads() { + skip_if_no_python!(); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("auto_chunk.h5"); + let path_str = path.display().to_string(); + + let (rows, cols) = (1500u64, 1100u64); // 13.2 MB of f64 + let data: Vec = (0..rows * cols).map(|i| (i % 9973) as f64 * 0.25).collect(); + let mut builder = FileBuilder::new(); + builder + .create_dataset("big") + .with_f64_data(&data) + .with_shape(&[rows, cols]) + .with_deflate(4); + builder + .create_dataset("small") + .with_f64_data(&data[..600]) + .with_shape(&[20, 30]) + .with_deflate(4); + builder.write(&path).unwrap(); + + let out = run_python_output(&format!( + r#" +import h5py, numpy as np +with h5py.File("{path_str}", "r") as f: + big, small = f["big"], f["small"] + expect = (np.arange(1500 * 1100) % 9973) * 0.25 + ok = bool(np.array_equal(big[...].ravel(), expect)) and bool(np.array_equal(small[...].ravel(), expect[:600])) + chunk_bytes = int(np.prod(big.chunks)) * 8 + print(ok, chunk_bytes <= 1 << 20, chunk_bytes >= 1 << 17, small.chunks == (20, 30), big.compression) +"# + )); + assert_eq!(out.trim(), "True True True True gzip"); + + // And it reads back here, in full and partially. + let file = File::open(&path).unwrap(); + let ds = file.dataset("big").unwrap(); + assert_eq!(ds.read_f64().unwrap(), data); + let row = clawhdf5_format::selection::Selection::Hyperslab { + start: vec![777, 0], + stride: vec![1, 1], + count: vec![1, cols], + block: vec![1, 1], + }; + let start = (777 * cols) as usize; + assert_eq!( + ds.read_f64_selection(&row).unwrap(), + data[start..start + cols as usize] + ); +}