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

Merge pull request 'docs(clawhdf5): document DType variants, fix unresolved doc links' (#17) from sdlc-docs/clawhdf5-types-20260514-165210 into main

Browse Source
This commit is contained in:
redclawsystems
2026-05-14 23:54:48 +00:00
commit 3f222f6956
3030 changed files with 89917 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/clawhdf5-format/src/chunk_index.rs
+677
View File
@@ -0,0 +1,677 @@
//! Chunk B-tree index cache and pre-computed layout for fast repeated reads.
//!
//! [`ChunkIndex`] scans the chunk index structure (B-tree v1, fixed array, etc.)
//! once and builds a flat `HashMap<ChunkCoord, ChunkInfo>` for O(1) lookups.
//!
//! [`ChunkLayout`] pre-computes contiguous row-copy operations for each chunk,
//! turning the N-D → flat coordinate math into a series of `memcpy` calls.
//! This eliminates per-element coordinate arithmetic during output assembly.
#[cfg(not(feature = "std"))]
extern crate alloc;
#[cfg(not(feature = "std"))]
use alloc::{vec, vec::Vec};
#[cfg(not(feature = "std"))]
use alloc::collections::BTreeMap;
#[cfg(feature = "std")]
use std::collections::HashMap;
use crate::chunk_cache::ChunkCoord;
use crate::chunked_read::ChunkInfo;
// ---------------------------------------------------------------------------
// ChunkIndex — flat HashMap of coord → ChunkInfo
// ---------------------------------------------------------------------------
/// A cached chunk index that maps chunk coordinates to file location metadata.
///
/// Built once from a B-tree (or other index structure) traversal, then provides
/// O(1) lookups for any chunk coordinate.
pub struct ChunkIndex {
#[cfg(feature = "std")]
map: HashMap<ChunkCoord, ChunkInfo>,
#[cfg(not(feature = "std"))]
map: BTreeMap<ChunkCoord, ChunkInfo>,
rank: usize,
}
impl ChunkIndex {
/// Build a chunk index from a pre-collected list of `ChunkInfo`.
///
/// `rank` is the spatial dimensionality — offsets are truncated to the first
/// `rank` elements (B-tree v1 stores rank+1 offsets with element size as the
/// last dimension).
pub fn build(chunks: &[ChunkInfo], rank: usize) -> Self {
#[cfg(feature = "std")]
let mut map = HashMap::with_capacity(chunks.len());
#[cfg(not(feature = "std"))]
let mut map = BTreeMap::new();
for ci in chunks {
let coord: ChunkCoord = ci.offsets.iter().take(rank).copied().collect();
map.insert(coord, ci.clone());
}
Self { map, rank }
}
/// O(1) lookup of a chunk by its spatial coordinate.
#[inline]
pub fn lookup(&self, coords: &[u64]) -> Option<&ChunkInfo> {
self.map.get(coords)
}
/// Number of chunks in the index.
#[inline]
pub fn num_chunks(&self) -> usize {
self.map.len()
}
/// The spatial rank this index was built for.
#[inline]
pub fn rank(&self) -> usize {
self.rank
}
/// Iterate over all indexed chunks.
pub fn iter(&self) -> impl Iterator<Item = (&ChunkCoord, &ChunkInfo)> {
self.map.iter()
}
/// Return all chunk infos as a Vec (order unspecified).
pub fn all_chunks(&self) -> Vec<ChunkInfo> {
self.map.values().cloned().collect()
}
}
// ---------------------------------------------------------------------------
// ChunkLayout — pre-computed row-copy operations
// ---------------------------------------------------------------------------
/// A single contiguous memory copy within the output assembly.
#[derive(Debug, Clone)]
pub struct RowCopy {
/// Byte offset within the decompressed chunk buffer.
pub src_offset: usize,
/// Byte offset within the output buffer.
pub dst_offset: usize,
/// Number of bytes to copy.
pub len: usize,
}
/// Pre-computed mapping for one chunk: its metadata plus the row-copy plan.
#[derive(Debug, Clone)]
pub struct ChunkMapping {
/// Spatial coordinate of this chunk.
pub coord: ChunkCoord,
/// File byte address of the compressed chunk.
pub file_offset: u64,
/// Size of the compressed chunk in the file.
pub file_size: u32,
/// Filter mask (0 = all filters applied).
pub filter_mask: u32,
/// Pre-computed row-copy operations for assembling this chunk into output.
pub copies: Vec<RowCopy>,
}
/// Pre-computed layout for assembling all chunks of a dataset into a flat
/// output buffer.
///
/// Built once from the chunk index + dataset/chunk dimensions, then reused
/// on every read. Turns the per-element N-D → flat coordinate math into a
/// sequence of `memcpy` calls.
pub struct ChunkLayout {
pub mappings: Vec<ChunkMapping>,
/// Total bytes of the output buffer.
pub output_bytes: usize,
/// Element size in bytes.
pub elem_size: usize,
/// Chunk total bytes (uncompressed).
pub chunk_total_bytes: usize,
}
impl ChunkLayout {
/// Build the chunk layout from a chunk index and dataset geometry.
///
/// - `index`: the chunk index (coord → ChunkInfo)
/// - `ds_dims`: dataset dimensions (e.g. [100, 200])
/// - `chunk_dims`: chunk dimensions (e.g. [10, 20])
/// - `elem_size`: size of a single element in bytes
pub fn build(
index: &ChunkIndex,
ds_dims: &[usize],
chunk_dims: &[usize],
elem_size: usize,
) -> Self {
let rank = ds_dims.len();
// Compute dataset strides (row-major)
let mut ds_strides = vec![1usize; rank];
for i in (0..rank.saturating_sub(1)).rev() {
ds_strides[i] = ds_strides[i + 1] * ds_dims[i + 1];
}
// Compute chunk strides (row-major within a chunk)
let mut chunk_strides = vec![1usize; rank];
for i in (0..rank.saturating_sub(1)).rev() {
chunk_strides[i] = chunk_strides[i + 1] * chunk_dims[i + 1];
}
let chunk_total_elements: usize = chunk_dims.iter().product();
let chunk_total_bytes = chunk_total_elements * elem_size;
let total_elements: usize = ds_dims.iter().product();
let output_bytes = total_elements * elem_size;
let mut mappings = Vec::with_capacity(index.num_chunks());
for (_coord, ci) in index.iter() {
let coord: ChunkCoord = ci.offsets.iter().take(rank).copied().collect();
let chunk_offsets: Vec<usize> = coord.iter().map(|&o| o as usize).collect();
let copies = if rank == 0 {
// Scalar dataset — single copy
let len = chunk_total_bytes.min(output_bytes);
vec![RowCopy {
src_offset: 0,
dst_offset: 0,
len,
}]
} else {
compute_row_copies(
&chunk_offsets,
chunk_dims,
ds_dims,
&ds_strides,
&chunk_strides,
elem_size,
rank,
)
};
mappings.push(ChunkMapping {
coord,
file_offset: ci.address,
file_size: ci.chunk_size,
filter_mask: ci.filter_mask,
copies,
});
}
Self {
mappings,
output_bytes,
elem_size,
chunk_total_bytes,
}
}
/// Assemble output from pre-decompressed chunk buffers using the
/// pre-computed row-copy plan.
///
/// `chunk_data` must be in the same order as `self.mappings`.
pub fn assemble(&self, chunk_data: &[&[u8]], output: &mut [u8]) {
debug_assert_eq!(chunk_data.len(), self.mappings.len());
for (mapping, data) in self.mappings.iter().zip(chunk_data.iter()) {
for copy in &mapping.copies {
let src_end = copy.src_offset + copy.len;
let dst_end = copy.dst_offset + copy.len;
debug_assert!(
src_end <= data.len() && dst_end <= output.len(),
"chunk copy out of bounds: src_end={src_end} data_len={} dst_end={dst_end} out_len={}",
data.len(),
output.len(),
);
if src_end <= data.len() && dst_end <= output.len() {
output[copy.dst_offset..dst_end]
.copy_from_slice(&data[copy.src_offset..src_end]);
}
}
}
}
}
/// Compute contiguous row-copy operations for a single chunk.
///
/// The innermost dimension is always contiguous in both chunk and output
/// memory (row-major layout). We iterate over the outer dimensions to find
/// the valid range of the innermost dim, then emit one `RowCopy` per
/// contiguous run.
fn compute_row_copies(
chunk_offsets: &[usize],
chunk_dims: &[usize],
ds_dims: &[usize],
ds_strides: &[usize],
chunk_strides: &[usize],
elem_size: usize,
rank: usize,
) -> Vec<RowCopy> {
// Number of "rows" = product of all dimensions except the innermost
let inner_dim = rank - 1;
let inner_chunk_len = chunk_dims[inner_dim];
let inner_ds_len = ds_dims[inner_dim];
let inner_chunk_offset = chunk_offsets[inner_dim];
// How many elements of the innermost dim are valid (not out-of-bounds)?
let inner_valid = if inner_chunk_offset >= inner_ds_len {
0
} else {
inner_chunk_len.min(inner_ds_len - inner_chunk_offset)
};
if inner_valid == 0 {
return Vec::new();
}
// Number of outer "rows" (product of dims 0..rank-1)
let outer_dims = &chunk_dims[..inner_dim];
let outer_count: usize = outer_dims.iter().product();
// Compute outer-only strides (row-major within the outer dimensions only).
// These are different from chunk_strides which include the innermost dim.
let mut outer_strides = vec![1usize; inner_dim];
for i in (0..inner_dim.saturating_sub(1)).rev() {
outer_strides[i] = outer_strides[i + 1] * outer_dims[i + 1];
}
let mut copies = Vec::with_capacity(outer_count);
for outer_flat in 0..outer_count {
// Convert outer flat index to N-D coordinates in dims 0..rank-1
let mut remaining = outer_flat;
let mut ds_flat_base = 0usize;
let mut chunk_flat_base = 0usize;
let mut out_of_bounds = false;
for d in 0..inner_dim {
let coord_in_chunk = remaining / outer_strides[d];
remaining %= outer_strides[d];
let global_coord = chunk_offsets[d] + coord_in_chunk;
if global_coord >= ds_dims[d] {
out_of_bounds = true;
break;
}
ds_flat_base += global_coord * ds_strides[d];
chunk_flat_base += coord_in_chunk * chunk_strides[d];
}
if out_of_bounds {
continue;
}
// Add the innermost dimension contribution
ds_flat_base += inner_chunk_offset * ds_strides[inner_dim];
// chunk_flat_base already aligned to start of inner row
let src_offset = chunk_flat_base * elem_size;
let dst_offset = ds_flat_base * elem_size;
let len = inner_valid * elem_size;
copies.push(RowCopy {
src_offset,
dst_offset,
len,
});
}
copies
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
use super::*;
use crate::chunked_read::ChunkInfo;
fn make_chunk(offsets: Vec<u64>, address: u64, size: u32) -> ChunkInfo {
ChunkInfo {
chunk_size: size,
filter_mask: 0,
offsets,
address,
}
}
// --- ChunkIndex tests ---
#[test]
fn index_build_and_lookup() {
let chunks = vec![
make_chunk(vec![0, 0, 4], 0x1000, 80),
make_chunk(vec![10, 0, 4], 0x2000, 80),
make_chunk(vec![10, 5, 4], 0x3000, 80),
];
let index = ChunkIndex::build(&chunks, 2);
assert_eq!(index.num_chunks(), 3);
assert_eq!(index.rank(), 2);
let c0 = index.lookup(&[0, 0]).unwrap();
assert_eq!(c0.address, 0x1000);
let c1 = index.lookup(&[10, 0]).unwrap();
assert_eq!(c1.address, 0x2000);
let c2 = index.lookup(&[10, 5]).unwrap();
assert_eq!(c2.address, 0x3000);
assert!(index.lookup(&[5, 0]).is_none());
}
#[test]
fn index_empty() {
let index = ChunkIndex::build(&[], 2);
assert_eq!(index.num_chunks(), 0);
assert!(index.lookup(&[0, 0]).is_none());
}
#[test]
fn index_1d() {
let chunks = vec![
make_chunk(vec![0, 8], 0x100, 40),
make_chunk(vec![5, 8], 0x200, 40),
];
let index = ChunkIndex::build(&chunks, 1);
assert_eq!(index.num_chunks(), 2);
assert_eq!(index.lookup(&[0]).unwrap().address, 0x100);
assert_eq!(index.lookup(&[5]).unwrap().address, 0x200);
}
#[test]
fn index_3d() {
let chunks = vec![
make_chunk(vec![0, 0, 0, 4], 0x100, 100),
make_chunk(vec![0, 0, 5, 4], 0x200, 100),
make_chunk(vec![2, 3, 0, 4], 0x300, 100),
];
let index = ChunkIndex::build(&chunks, 3);
assert_eq!(index.num_chunks(), 3);
assert_eq!(index.lookup(&[0, 0, 0]).unwrap().address, 0x100);
assert_eq!(index.lookup(&[0, 0, 5]).unwrap().address, 0x200);
assert_eq!(index.lookup(&[2, 3, 0]).unwrap().address, 0x300);
}
// --- ChunkLayout tests ---
#[test]
fn layout_1d_two_chunks() {
// 20-element 1D dataset, chunk size 10, f64
let chunks = vec![
make_chunk(vec![0, 8], 0x1000, 80),
make_chunk(vec![10, 8], 0x2000, 80),
];
let index = ChunkIndex::build(&chunks, 1);
let layout = ChunkLayout::build(&index, &[20], &[10], 8);
assert_eq!(layout.output_bytes, 160); // 20 * 8
assert_eq!(layout.mappings.len(), 2);
// Each 1D chunk should have exactly 1 row-copy
for mapping in &layout.mappings {
assert_eq!(mapping.copies.len(), 1);
assert_eq!(mapping.copies[0].len, 80); // 10 * 8
}
// Verify assembly
let chunk0: Vec<u8> = (0..10u64).flat_map(|i| (i as f64).to_le_bytes()).collect();
let chunk1: Vec<u8> = (10..20u64).flat_map(|i| (i as f64).to_le_bytes()).collect();
let mut output = vec![0u8; 160];
// Find which mapping goes where
let data_refs: Vec<&[u8]> = layout
.mappings
.iter()
.map(|m| {
if m.coord == vec![0] {
chunk0.as_slice()
} else {
chunk1.as_slice()
}
})
.collect();
layout.assemble(&data_refs, &mut output);
for i in 0..20u64 {
let val = f64::from_le_bytes(
output[i as usize * 8..(i as usize + 1) * 8]
.try_into()
.unwrap(),
);
assert_eq!(val, i as f64, "mismatch at index {i}");
}
}
#[test]
fn layout_2d_four_chunks() {
// 4x6 dataset, chunk 2x3, f32 (4 bytes)
let chunks = vec![
make_chunk(vec![0, 0, 4], 0x100, 24),
make_chunk(vec![0, 3, 4], 0x200, 24),
make_chunk(vec![2, 0, 4], 0x300, 24),
make_chunk(vec![2, 3, 4], 0x400, 24),
];
let index = ChunkIndex::build(&chunks, 2);
let layout = ChunkLayout::build(&index, &[4, 6], &[2, 3], 4);
assert_eq!(layout.output_bytes, 96); // 24 * 4
assert_eq!(layout.mappings.len(), 4);
// Each 2x3 chunk should have 2 row-copies (one per row)
for mapping in &layout.mappings {
assert_eq!(mapping.copies.len(), 2);
// Each copy should be 3 elements * 4 bytes = 12 bytes
for copy in &mapping.copies {
assert_eq!(copy.len, 12);
}
}
// Build chunk data and verify assembly
let mut chunk_data_map: Vec<(Vec<u64>, Vec<u8>)> = Vec::new();
for row_start in [0usize, 2] {
for col_start in [0usize, 3] {
let mut data = Vec::new();
for r in 0..2 {
for c in 0..3 {
let val = ((row_start + r) * 6 + (col_start + c)) as f32;
data.extend_from_slice(&val.to_le_bytes());
}
}
chunk_data_map.push((vec![row_start as u64, col_start as u64], data));
}
}
let data_refs: Vec<&[u8]> = layout
.mappings
.iter()
.map(|m| {
chunk_data_map
.iter()
.find(|(c, _)| c == &m.coord)
.unwrap()
.1
.as_slice()
})
.collect();
let mut output = vec![0u8; 96];
layout.assemble(&data_refs, &mut output);
for i in 0..24 {
let val = f32::from_le_bytes(output[i * 4..(i + 1) * 4].try_into().unwrap());
assert_eq!(val, i as f32, "mismatch at element {i}");
}
}
#[test]
fn layout_partial_last_chunk() {
// 25-element 1D, chunk 10 => 3 chunks, last has 5 valid
let chunks = vec![
make_chunk(vec![0, 8], 0x100, 80),
make_chunk(vec![10, 8], 0x200, 80),
make_chunk(vec![20, 8], 0x300, 80),
];
let index = ChunkIndex::build(&chunks, 1);
let layout = ChunkLayout::build(&index, &[25], &[10], 8);
assert_eq!(layout.output_bytes, 200); // 25 * 8
// Find the chunk at offset 20 — should only copy 5 elements
let partial = layout
.mappings
.iter()
.find(|m| m.coord == vec![20])
.unwrap();
assert_eq!(partial.copies.len(), 1);
assert_eq!(partial.copies[0].len, 40); // 5 * 8 bytes
}
#[test]
fn layout_2d_partial_boundary() {
// 5x5 dataset, chunk 3x3 => 4 chunks, some partial
let chunks = vec![
make_chunk(vec![0, 0, 4], 0x100, 36),
make_chunk(vec![0, 3, 4], 0x200, 36),
make_chunk(vec![3, 0, 4], 0x300, 36),
make_chunk(vec![3, 3, 4], 0x400, 36),
];
let index = ChunkIndex::build(&chunks, 2);
let layout = ChunkLayout::build(&index, &[5, 5], &[3, 3], 4);
assert_eq!(layout.output_bytes, 100); // 25 * 4
// Chunk at (0,0): 3x3 fully valid → 3 rows, 3 elems each
let c00 = layout
.mappings
.iter()
.find(|m| m.coord == vec![0, 0])
.unwrap();
assert_eq!(c00.copies.len(), 3);
for copy in &c00.copies {
assert_eq!(copy.len, 12); // 3 * 4
}
// Chunk at (0,3): 3x2 valid (cols 3,4) → 3 rows, 2 elems each
let c03 = layout
.mappings
.iter()
.find(|m| m.coord == vec![0, 3])
.unwrap();
assert_eq!(c03.copies.len(), 3);
for copy in &c03.copies {
assert_eq!(copy.len, 8); // 2 * 4
}
// Chunk at (3,0): 2x3 valid (rows 3,4) → 2 rows, 3 elems each
let c30 = layout
.mappings
.iter()
.find(|m| m.coord == vec![3, 0])
.unwrap();
assert_eq!(c30.copies.len(), 2);
// Chunk at (3,3): 2x2 valid → 2 rows, 2 elems each
let c33 = layout
.mappings
.iter()
.find(|m| m.coord == vec![3, 3])
.unwrap();
assert_eq!(c33.copies.len(), 2);
for copy in &c33.copies {
assert_eq!(copy.len, 8); // 2 * 4
}
}
#[test]
fn layout_3d_basic() {
// 4x4x4 dataset, chunk 2x2x2, f32 (4 bytes) => 8 chunks
let mut chunks = Vec::new();
let elem_size = 4u32;
let chunk_bytes = 2 * 2 * 2 * elem_size;
let mut addr = 0x1000u64;
for x in [0u64, 2] {
for y in [0u64, 2] {
for z in [0u64, 2] {
chunks.push(make_chunk(
vec![x, y, z, elem_size as u64],
addr,
chunk_bytes,
));
addr += chunk_bytes as u64;
}
}
}
let index = ChunkIndex::build(&chunks, 3);
let layout = ChunkLayout::build(&index, &[4, 4, 4], &[2, 2, 2], 4);
assert_eq!(layout.output_bytes, 256); // 64 * 4
assert_eq!(layout.mappings.len(), 8);
// Each 2x2x2 chunk: outer dims are 2x2=4 rows, each row copies 2 elements
for mapping in &layout.mappings {
assert_eq!(mapping.copies.len(), 4); // 2*2 outer rows
for copy in &mapping.copies {
assert_eq!(copy.len, 8); // 2 * 4 bytes
}
}
}
#[test]
fn assemble_matches_element_by_element() {
// Verify that layout assembly produces identical output to the
// original copy_chunk_to_output approach (2D case)
let ds_dims = [6usize, 8];
let chunk_dims = [3usize, 4];
let elem_size = 8usize; // f64
let mut chunks = Vec::new();
let mut addr = 0x1000u64;
for r in (0..6).step_by(3) {
for c in (0..8).step_by(4) {
chunks.push(make_chunk(
vec![r as u64, c as u64, elem_size as u64],
addr,
(3 * 4 * elem_size) as u32,
));
addr += (3 * 4 * elem_size) as u64;
}
}
let index = ChunkIndex::build(&chunks, 2);
let layout = ChunkLayout::build(&index, &ds_dims, &chunk_dims, elem_size);
// Build chunk data
let mut chunk_buffers: Vec<(Vec<u64>, Vec<u8>)> = Vec::new();
for mapping in &layout.mappings {
let row_start = mapping.coord[0] as usize;
let col_start = mapping.coord[1] as usize;
let mut buf = vec![0u8; 3 * 4 * elem_size];
for r in 0..3 {
for c in 0..4 {
let val = ((row_start + r) * 8 + (col_start + c)) as f64;
let offset = (r * 4 + c) * elem_size;
buf[offset..offset + 8].copy_from_slice(&val.to_le_bytes());
}
}
chunk_buffers.push((mapping.coord.clone(), buf));
}
let data_refs: Vec<&[u8]> = chunk_buffers.iter().map(|(_, b)| b.as_slice()).collect();
let mut output = vec![0u8; 6 * 8 * elem_size];
layout.assemble(&data_refs, &mut output);
// Verify every element
for r in 0..6 {
for c in 0..8 {
let idx = r * 8 + c;
let val = f64::from_le_bytes(output[idx * 8..(idx + 1) * 8].try_into().unwrap());
assert_eq!(val, idx as f64, "mismatch at ({r}, {c})");
}
}
}
}