feat(format): choose chunk dimensions automatically for large datasets

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 <[email protected]>
This commit is contained in:
osobh
2026-09-19 14:21:04 -07:00
co-authored by Claude Fable 5.1
parent b36c6ec2af
commit 05c665a898
4 changed files with 152 additions and 6 deletions
+8
View File
@@ -22,6 +22,14 @@
type 5 — what `libver='latest'` uses for two or more unlimited dimensions; type 5 — what `libver='latest'` uses for two or more unlimited dimensions;
previously "unsupported chunked layout"). The four copies of the chunk-index previously "unsupported chunked layout"). The four copies of the chunk-index
dispatch are now one shared function, so every read path gets it. 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 - **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 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 range wrapped into the next row and returned that element. Now
+82 -5
View File
@@ -49,6 +49,38 @@ pub struct ChunkOptions {
pub pcodec: bool, 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<u64> {
let mut dims: Vec<u64> = 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 { impl ChunkOptions {
/// Whether any chunking option is enabled. /// Whether any chunking option is enabled.
pub fn is_chunked(&self) -> bool { pub fn is_chunked(&self) -> bool {
@@ -135,11 +167,17 @@ impl ChunkOptions {
/// Determine chunk dimensions, using user-specified or auto-computing. /// Determine chunk dimensions, using user-specified or auto-computing.
pub fn resolve_chunk_dims(&self, shape: &[u64]) -> Vec<u64> { pub fn resolve_chunk_dims(&self, shape: &[u64]) -> Vec<u64> {
if let Some(ref dims) = self.chunk_dims { // Without the element size, assume 8 bytes (the widest common scalar);
dims.clone() // the writer uses `resolve_chunk_dims_for`.
} else { self.resolve_chunk_dims_for(shape, 8)
// Auto chunk: use the full dataset shape (single chunk) }
shape.to_vec()
/// 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<u64> {
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]); assert_eq!(dims, vec![100, 50]);
} }
#[test]
fn auto_chunking_splits_only_large_datasets() {
let bytes = |dims: &[u64], elem: u64| dims.iter().product::<u64>() * 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] #[test]
fn chunk_options_pipeline_deflate() { fn chunk_options_pipeline_deflate() {
// Auto-shuffle is applied before compression by default (matches h5py). // Auto-shuffle is applied before compression by default (matches h5py).
+3 -1
View File
@@ -1221,8 +1221,10 @@ impl FileWriter {
precompressed: None, precompressed: None,
}); });
} else if is_chunked[i] { } 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 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 // Compress once in Pass 1; cache the result so Pass 2 can skip
// re-compression and just rebuild the index with real addresses. // re-compression and just rebuild the index with real addresses.
let pre = precompress_chunks( let pre = precompress_chunks(
@@ -979,3 +979,62 @@ with h5py.File("{path_str}", "r") as f:
expected["slab"] 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<f64> = (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]
);
}