quantumclaw/clawhdf5
3
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

Merge branch 'fix/p0-chunked-read' into fix/phase0-correctness

Browse Source
This commit is contained in:
osobh
2026-09-25 21:20:59 -05:00
parent 4b23ad697c 9066d34eaa
commit 36d7a6f234
10 changed files with 1381 additions and 429 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/clawhdf5-format/src/chunked_read.rs
+162 -105
View File
@@ -15,7 +15,7 @@ use crate::datatype::Datatype;
use crate::error::FormatError;
use crate::extensible_array::{ExtensibleArrayHeader, read_extensible_array_chunks};
use crate::filter_pipeline::FilterPipeline;
use crate::filters::decompress_chunk;
use crate::filters::{all_filters_skipped, decompress_chunk_masked};
use crate::fixed_array::{FixedArrayHeader, read_fixed_array_chunks};
#[cfg(feature = "std")]
use std::sync::Arc;
@@ -65,11 +65,13 @@ fn decompress_all_chunks(
let raw_chunk = &file_data[c_addr..c_addr + size];
let decompressed = if let Some(pl) = pipeline {
if chunk_info.filter_mask == 0 {
decompress_chunk(raw_chunk, pl, chunk_total_bytes, element_size)?
} else {
raw_chunk.to_vec()
}
decompress_chunk_masked(
raw_chunk,
pl,
chunk_total_bytes,
element_size,
chunk_info.filter_mask,
)?
} else {
raw_chunk.to_vec()
};
@@ -223,6 +225,10 @@ pub fn collect_chunk_info(
collect_chunk_info_inner(file_data, btree_address, ndims, offset_size, length_size, 0)
}
/// Width of each chunk offset in a v1 chunk B-tree key, independent of the
/// file's size-of-offsets.
const CHUNK_KEY_OFFSET_SIZE: u8 = 8;
/// Maximum recursion depth for chunk B-tree traversal (malformed/cyclic data
/// protection), matching `btree_v1.rs`'s `MAX_BTREE_DEPTH`.
const MAX_CHUNK_BTREE_DEPTH: usize = 64;
@@ -260,8 +266,14 @@ fn collect_chunk_info_inner(
let mut pos = offset + 8 + os * 2; // skip left/right sibling
// Key size: chunk_size(4) + filter_mask(4) + ndims * offset_size
let key_size = 4 + 4 + ndims * os;
// Key: chunk_size(4) + filter_mask(4) + one offset per dimension. The
// offsets are always 8 bytes each — they are dataset coordinates, not file
// addresses, so they do not follow the superblock's size-of-offsets (only
// the sibling and child addresses do).
let key_size = ndims
.checked_mul(CHUNK_KEY_OFFSET_SIZE as usize)
.and_then(|n| n.checked_add(8))
.ok_or_else(|| FormatError::ChunkedReadError("chunk key too large".into()))?;
if node_level == 0 {
// Leaf node: keys and children interleaved
@@ -287,8 +299,8 @@ fn collect_chunk_info_inner(
let mut offsets = Vec::with_capacity(ndims);
let mut kp = pos + 8;
for _ in 0..ndims {
offsets.push(read_offset(file_data, kp, offset_size)?);
kp += os;
offsets.push(read_offset(file_data, kp, CHUNK_KEY_OFFSET_SIZE)?);
kp += CHUNK_KEY_OFFSET_SIZE as usize;
}
pos += key_size;
@@ -507,6 +519,7 @@ pub fn list_chunks(
addr_opt,
single_filtered_size,
single_filter_mask,
unfiltered_edges,
) = match layout {
DataLayout::Chunked {
chunk_dimensions,
@@ -515,6 +528,7 @@ pub fn list_chunks(
chunk_index_type,
single_chunk_filtered_size,
single_chunk_filter_mask,
dont_filter_partial_edge_chunks,
} => (
chunk_dimensions,
*version,
@@ -522,6 +536,7 @@ pub fn list_chunks(
*btree_address,
*single_chunk_filtered_size,
*single_chunk_filter_mask,
*dont_filter_partial_edge_chunks,
),
_ => {
return Err(FormatError::ChunkedReadError(
@@ -554,7 +569,7 @@ pub fn list_chunks(
}
// Collect chunks based on version and index type
let chunks = match (version, chunk_index_type) {
let mut chunks = match (version, chunk_index_type) {
(3, _) => {
let ndims = chunk_dimensions.len(); // rank+1
collect_chunk_info(file_data, addr, ndims, offset_size, length_size)?
@@ -635,6 +650,23 @@ pub fn list_chunks(
}
};
// With "don't filter partial edge chunks", a chunk that extends past the
// dataset's extent is stored raw while its filter mask still reads 0.
// Mark every filter skipped so all read paths copy it as-is.
if unfiltered_edges {
for chunk in &mut chunks {
let partial = chunk
.offsets
.iter()
.zip(&chunk_dims)
.zip(&ds_dims)
.any(|((&off, &cd), &dd)| off.saturating_add(cd as u64) > dd as u64);
if partial {
chunk.filter_mask = u32::MAX;
}
}
}
Ok((chunks, chunk_dims))
}
@@ -806,24 +838,20 @@ pub fn read_chunked_data_cached(
)));
}
// The per-file cache is shared across datasets; bind it to this one so a
// different dataset's chunk index is never reused for this read.
cache.ensure_dataset(addr);
// Populate chunk index on first access
if !cache.has_index() {
let (chunks, _) = list_chunks(
// The per-file cache is shared across datasets (and threads); every
// lookup is keyed by this dataset's chunk-index address, so another
// dataset's index or chunks are never used for this read.
let chunks = cache.chunks_for(addr, rank, || {
list_chunks(
file_data,
layout,
dataspace,
elem_size,
offset_size,
length_size,
)?;
cache.populate_index(&chunks, rank);
}
let chunks = cache.all_indexed_chunks().unwrap_or_default();
)
.map(|(chunks, _)| chunks)
})?;
// Assemble output
let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?;
@@ -878,10 +906,11 @@ pub fn read_chunked_data_cached(
};
// Chunks stored as-is (no pipeline, or the filter mask says this chunk
// skipped it) are copied straight from the file bytes: they are already in
// memory, so routing them through a Vec and then an aligned cache buffer
// was two extra copies of the whole dataset for nothing.
let stored_raw = |c: &ChunkInfo| pipeline.is_none() || c.filter_mask != 0;
// skipped every filter) are copied straight from the file bytes: they are
// already in memory, so routing them through a Vec and then an aligned
// cache buffer was two extra copies of the whole dataset for nothing.
let stored_raw =
|c: &ChunkInfo| pipeline.is_none_or(|pl| all_filters_skipped(pl, c.filter_mask));
let mut misses: Vec<&ChunkInfo> = Vec::new();
for chunk_info in &chunks {
if stored_raw(chunk_info) {
@@ -889,7 +918,7 @@ pub fn read_chunked_data_cached(
continue;
}
let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();
match cache.get_decompressed_aligned(&coord) {
match cache.get_decompressed_in(addr, &coord) {
Some(cached) => place(&cached, chunk_info),
None => misses.push(chunk_info),
}
@@ -903,7 +932,13 @@ pub fn read_chunked_data_cached(
let cache_them = total_bytes <= cache.max_bytes();
if let Some(pl) = pipeline {
let decode = |c: &&ChunkInfo| -> Result<Vec<u8>, FormatError> {
decompress_chunk(raw_bytes(c)?, pl, chunk_total_bytes, elem_size as u32)
decompress_chunk_masked(
raw_bytes(c)?,
pl,
chunk_total_bytes,
elem_size as u32,
c.filter_mask,
)
};
for batch in misses.chunks(DECODE_BATCH) {
#[cfg(feature = "parallel")]
@@ -920,7 +955,7 @@ pub fn read_chunked_data_cached(
let data = data?;
if cache_them {
let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();
let cached = cache.put_decompressed(coord, data);
let cached = cache.put_decompressed_in(addr, coord, data);
place(&cached, chunk_info);
} else {
place(&data, chunk_info);
@@ -1124,24 +1159,20 @@ pub fn read_chunked_data_sweep(
)));
}
// The per-file cache is shared across datasets; bind it to this one so a
// different dataset's chunk index is never reused for this read.
cache.ensure_dataset(addr);
// Populate chunk index on first access
if !cache.has_index() {
let (chunks, _) = list_chunks(
// The per-file cache is shared across datasets (and threads); every
// lookup is keyed by this dataset's chunk-index address, so another
// dataset's index or chunks are never used for this read.
let chunks = cache.chunks_for(addr, rank, || {
list_chunks(
file_data,
layout,
dataspace,
elem_size,
offset_size,
length_size,
)?;
cache.populate_index(&chunks, rank);
}
let chunks = cache.all_indexed_chunks().unwrap_or_default();
)
.map(|(chunks, _)| chunks)
})?;
// Assemble output
let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?;
@@ -1172,12 +1203,12 @@ pub fn read_chunked_data_sweep(
// Issue prefetch hint for predicted next chunks
if !sweep.predicted_next.is_empty() {
cache.prefetch_hint(&sweep.predicted_next);
cache.prefetch_hint_in(addr, &sweep.predicted_next);
cache.set_sweep_direction(sweep.direction);
}
// Try decompressed cache first
let decompressed = if let Some(cached) = cache.get_decompressed_aligned(&coord) {
let decompressed = if let Some(cached) = cache.get_decompressed_in(addr, &coord) {
cached
} else {
// Decompress from file
@@ -1186,15 +1217,17 @@ pub fn read_chunked_data_sweep(
ensure_len(file_data, c_addr, size)?;
let raw_chunk = &file_data[c_addr..c_addr + size];
let dec = if let Some(pl) = pipeline {
if chunk_info.filter_mask == 0 {
decompress_chunk(raw_chunk, pl, chunk_total_bytes, elem_size as u32)?
} else {
raw_chunk.to_vec()
}
decompress_chunk_masked(
raw_chunk,
pl,
chunk_total_bytes,
elem_size as u32,
chunk_info.filter_mask,
)?
} else {
raw_chunk.to_vec()
};
cache.put_decompressed(coord, dec)
cache.put_decompressed_in(addr, coord, dec)
};
let chunk_offsets: Vec<usize> = chunk_info
@@ -1278,48 +1311,34 @@ pub fn read_chunked_data_indexed(
)));
}
// The per-file cache is shared across datasets; bind it to this one so a
// different dataset's chunk index is never reused for this read.
cache.ensure_dataset(addr);
// Build chunk index on first access
if !cache.has_chunk_index() {
let (chunks, _) = list_chunks(
file_data,
layout,
dataspace,
elem_size,
offset_size,
length_size,
)?;
cache.populate_chunk_index(&chunks, rank);
// Also populate the legacy index for compatibility
if !cache.has_index() {
cache.populate_index(&chunks, rank);
}
}
// Build chunk layout on first access
if !cache.has_chunk_layout() {
cache.populate_chunk_layout(&ds_dims, &chunk_dims, elem_size);
}
// Get the layout info (mappings, output size, chunk total bytes)
let (mappings_info, output_bytes, chunk_total_bytes) = cache
.with_chunk_layout(|layout| {
let info: Vec<_> = layout
.mappings
.iter()
.map(|m| (m.coord.clone(), m.file_offset, m.file_size, m.filter_mask))
.collect();
(info, layout.output_bytes, layout.chunk_total_bytes)
})
.ok_or_else(|| FormatError::ChunkedReadError("chunk layout not available".into()))?;
// Chunk index and assembly plan for this dataset, built on first access
// and kept per dataset (keyed by chunk-index address) in the shared cache.
let plan = cache.chunk_layout_for(
addr,
rank,
|| {
list_chunks(
file_data,
layout,
dataspace,
elem_size,
offset_size,
length_size,
)
.map(|(chunks, _)| chunks)
},
&ds_dims,
&chunk_dims,
elem_size,
)?;
let chunk_total_bytes = plan.chunk_total_bytes;
// Decompress chunks (using LRU cache where possible)
let mut chunk_buffers: Vec<Arc<CacheAlignedBuffer>> = Vec::with_capacity(mappings_info.len());
for (coord, file_offset, file_size, filter_mask) in &mappings_info {
if let Some(cached) = cache.get_decompressed_aligned(coord) {
let mut chunk_buffers: Vec<Arc<CacheAlignedBuffer>> = Vec::with_capacity(plan.mappings.len());
for m in &plan.mappings {
let (coord, file_offset, file_size, filter_mask) =
(&m.coord, &m.file_offset, &m.file_size, &m.filter_mask);
if let Some(cached) = cache.get_decompressed_in(addr, coord) {
chunk_buffers.push(cached);
} else {
let c_addr = *file_offset as usize;
@@ -1327,26 +1346,26 @@ pub fn read_chunked_data_indexed(
ensure_len(file_data, c_addr, size)?;
let raw_chunk = &file_data[c_addr..c_addr + size];
let decompressed = if let Some(pl) = pipeline {
if *filter_mask == 0 {
decompress_chunk(raw_chunk, pl, chunk_total_bytes, elem_size as u32)?
} else {
raw_chunk.to_vec()
}
decompress_chunk_masked(
raw_chunk,
pl,
chunk_total_bytes,
elem_size as u32,
*filter_mask,
)?
} else {
raw_chunk.to_vec()
};
let aligned = CacheAlignedBuffer::from_vec(decompressed);
let arc = cache.put_decompressed_aligned(coord.clone(), aligned);
let arc = cache.put_decompressed_aligned_in(addr, coord.clone(), aligned);
chunk_buffers.push(arc);
}
}
// Assemble using pre-computed layout
let mut output = vec![0u8; output_bytes];
let mut output = vec![0u8; plan.output_bytes];
let data_refs: Vec<&[u8]> = chunk_buffers.iter().map(|b| b.as_slice()).collect();
cache.with_chunk_layout(|layout| {
layout.assemble(&data_refs, &mut output);
});
plan.assemble(&data_refs, &mut output);
Ok(output)
}
@@ -1594,7 +1613,8 @@ mod tests {
} else {
0
};
write_offset(&mut buf, off, offset_size);
// Key offsets are always 8 bytes (they are coordinates).
write_offset(&mut buf, off, 8);
}
// Child: address
write_offset(&mut buf, chunk.address, offset_size);
@@ -1604,7 +1624,7 @@ mod tests {
buf.extend_from_slice(&0u32.to_le_bytes()); // chunk_size
buf.extend_from_slice(&0u32.to_le_bytes()); // filter_mask
for _ in 0..ndims {
write_offset(&mut buf, u64::MAX, offset_size);
write_offset(&mut buf, u64::MAX, 8);
}
buf
@@ -1682,6 +1702,37 @@ mod tests {
assert_eq!(result[2].address, 0x300);
}
#[test]
fn collect_chunks_with_four_byte_addresses() {
// Sibling and child addresses are 4 bytes; the key offsets stay 8.
let ndims = 3;
let os: u8 = 4;
let chunks = vec![
ChunkInfo {
chunk_size: 80,
filter_mask: 2,
offsets: vec![0, 5, 0],
address: 0x1000,
},
ChunkInfo {
chunk_size: 96,
filter_mask: 0,
offsets: vec![8, 10, 0],
address: 0x2000,
},
];
let btree = build_chunk_btree_leaf(&chunks, ndims, os);
assert_eq!(btree.len(), 8 + 2 * 4 + 2 * (8 + 3 * 8 + 4) + (8 + 3 * 8));
let result = collect_chunk_info(&btree, 0, ndims, os, os).unwrap();
assert_eq!(result.len(), 2);
for (got, want) in result.iter().zip(&chunks) {
assert_eq!(got.offsets, want.offsets);
assert_eq!(got.address, want.address);
assert_eq!(got.chunk_size, want.chunk_size);
assert_eq!(got.filter_mask, want.filter_mask);
}
}
#[test]
fn collect_empty_btree() {
let ndims = 2;
@@ -1776,6 +1827,7 @@ mod tests {
chunk_index_type: None,
single_chunk_filtered_size: None,
single_chunk_filter_mask: None,
dont_filter_partial_edge_chunks: false,
};
let dataspace = Dataspace {
@@ -1799,6 +1851,7 @@ mod tests {
chunk_index_type: None,
single_chunk_filtered_size: None,
single_chunk_filter_mask: None,
dont_filter_partial_edge_chunks: false,
};
let dataspace = Dataspace {
space_type: DataspaceType::Simple,
@@ -1956,6 +2009,7 @@ mod tests {
chunk_index_type: None,
single_chunk_filtered_size: None,
single_chunk_filter_mask: None,
dont_filter_partial_edge_chunks: false,
};
let dataspace = Dataspace {
space_type: DataspaceType::Simple,
@@ -2038,6 +2092,7 @@ mod tests {
chunk_index_type: None,
single_chunk_filtered_size: None,
single_chunk_filter_mask: None,
dont_filter_partial_edge_chunks: false,
};
let dataspace = Dataspace {
space_type: DataspaceType::Simple,
@@ -2200,6 +2255,7 @@ mod tests {
chunk_index_type: Some(1),
single_chunk_filtered_size: None,
single_chunk_filter_mask: None,
dont_filter_partial_edge_chunks: false,
};
let dataspace = Dataspace {
space_type: DataspaceType::Simple,
@@ -2229,12 +2285,12 @@ mod tests {
let datatype = make_f64_type();
let cache = ChunkCache::new();
assert!(!cache.has_index());
assert_eq!(cache.indexed_dataset_count(), 0);
let raw = read_chunked_data_cached(
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
)
.unwrap();
assert!(cache.has_index());
assert_eq!(cache.indexed_dataset_count(), 1);
assert_eq!(raw.len(), 20 * 8);
for i in 0..20 {
let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
@@ -2256,7 +2312,7 @@ mod tests {
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
)
.unwrap();
assert!(cache.has_index());
assert_eq!(cache.indexed_dataset_count(), 1);
assert_eq!(cache.cached_chunk_count(), 0);
// Second read — reuses the cached index
@@ -2265,6 +2321,7 @@ mod tests {
)
.unwrap();
assert_eq!(raw1, raw2);
assert_eq!(cache.indexed_dataset_count(), 1);
}
#[test]