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:
co-authored by
Claude Fable 5.1
parent
b36c6ec2af
commit
05c665a898
@@ -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<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 {
|
||||
/// 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<u64> {
|
||||
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<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]);
|
||||
}
|
||||
|
||||
#[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]
|
||||
fn chunk_options_pipeline_deflate() {
|
||||
// Auto-shuffle is applied before compression by default (matches h5py).
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user