osobhandClaude Fable 5.1
0addf328bc
Same-moment A/B on a 64 MB f64 dataset: chunked+deflate 110 -> 69 ms, chunked 72 -> 60 ms, contiguous 56 -> 30 ms. - read_chunked_data_cached — the path the facade uses — decompressed chunks one at a time; only the uncached reader was parallel. Cache misses are now decoded in bounded batches (128), in parallel with the `parallel` feature. - Every chunk was pushed into the 16 MiB chunk cache, which a larger dataset just churns (insert, evict moments later). Chunks are cached only when the whole dataset fits (new ChunkCache::max_bytes). - Unfiltered chunks went file -> Vec -> aligned cache buffer -> output. They are copied straight from the file bytes. - The facade's typed reads convert a contiguous dataset straight from the borrowed file bytes instead of copying it into a Vec first. - The native little-endian fast paths allocated vec![0; n] and then overwrote it; they now fill an uninitialised buffer in one copy (native_le_to_vec). alloc_output requests zeroed memory from the allocator instead of reserving and filling. The unit test that expected unfiltered chunks to land in the decompressed cache now asserts the new design (index reused, cache not involved). Co-Authored-By: Claude Fable 5.1 <[email protected]>
2363 lines
79 KiB
Rust
2363 lines
79 KiB
Rust
//! Chunked dataset reading: B-tree v1 type 1 traversal and chunk assembly.
|
|
|
|
#[cfg(not(feature = "std"))]
|
|
extern crate alloc;
|
|
|
|
#[cfg(not(feature = "std"))]
|
|
use alloc::{format, vec, vec::Vec};
|
|
|
|
use crate::chunk_cache::CacheAlignedBuffer;
|
|
#[cfg(feature = "std")]
|
|
use crate::chunk_cache::ChunkCache;
|
|
use crate::data_layout::DataLayout;
|
|
use crate::dataspace::Dataspace;
|
|
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::fixed_array::{FixedArrayHeader, read_fixed_array_chunks};
|
|
#[cfg(feature = "std")]
|
|
use std::sync::Arc;
|
|
|
|
#[cfg(feature = "parallel")]
|
|
use crate::parallel_read;
|
|
|
|
#[cfg(feature = "parallel")]
|
|
use crate::lane_partition::PartitionStats;
|
|
|
|
/// Decompress all chunks into cache-line-aligned buffers, using lane-partitioned
|
|
/// parallel decompression when the `parallel` feature is enabled and the chunk
|
|
/// count exceeds the threshold.
|
|
fn decompress_all_chunks(
|
|
file_data: &[u8],
|
|
chunks: &[ChunkInfo],
|
|
pipeline: Option<&FilterPipeline>,
|
|
chunk_total_bytes: usize,
|
|
element_size: u32,
|
|
) -> Result<Vec<CacheAlignedBuffer>, FormatError> {
|
|
#[cfg(feature = "parallel")]
|
|
{
|
|
if let Some(pl) = pipeline
|
|
&& parallel_read::should_use_parallel(chunks.len())
|
|
{
|
|
// Seed from the first chunk's address and count for determinism.
|
|
let seed = chunks.first().map(|c| c.address).unwrap_or(0) ^ (chunks.len() as u64);
|
|
let (data, _stats) = parallel_read::decompress_chunks_lane_partitioned(
|
|
file_data,
|
|
chunks,
|
|
pl,
|
|
chunk_total_bytes,
|
|
element_size,
|
|
seed,
|
|
None, // auto-detect lane count
|
|
)?;
|
|
return Ok(data.into_iter().map(CacheAlignedBuffer::from_vec).collect());
|
|
}
|
|
}
|
|
|
|
// Sequential fallback — allocate into aligned buffers
|
|
let mut result = Vec::with_capacity(chunks.len());
|
|
for chunk_info in chunks {
|
|
let c_addr = chunk_info.address as usize;
|
|
let size = chunk_info.chunk_size as usize;
|
|
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 chunk_info.filter_mask == 0 {
|
|
decompress_chunk(raw_chunk, pl, chunk_total_bytes, element_size)?
|
|
} else {
|
|
raw_chunk.to_vec()
|
|
}
|
|
} else {
|
|
raw_chunk.to_vec()
|
|
};
|
|
result.push(CacheAlignedBuffer::from_vec(decompressed));
|
|
}
|
|
Ok(result)
|
|
}
|
|
|
|
/// Decompress all chunks with lane-partitioned parallelism and return
|
|
/// per-lane diagnostics.
|
|
///
|
|
/// This is the stats-returning variant for callers who want to inspect
|
|
/// the partition balance. Only available with the `parallel` feature.
|
|
#[cfg(feature = "parallel")]
|
|
pub fn decompress_all_chunks_with_stats(
|
|
file_data: &[u8],
|
|
chunks: &[ChunkInfo],
|
|
pipeline: &FilterPipeline,
|
|
chunk_total_bytes: usize,
|
|
element_size: u32,
|
|
seed: u64,
|
|
num_lanes: Option<usize>,
|
|
) -> Result<(Vec<Vec<u8>>, PartitionStats), FormatError> {
|
|
parallel_read::decompress_chunks_lane_partitioned(
|
|
file_data,
|
|
chunks,
|
|
pipeline,
|
|
chunk_total_bytes,
|
|
element_size,
|
|
seed,
|
|
num_lanes,
|
|
)
|
|
}
|
|
|
|
/// Information about a single chunk in a chunked dataset.
|
|
#[derive(Debug, Clone)]
|
|
pub struct ChunkInfo {
|
|
/// Size of chunk data in the file (after compression).
|
|
pub chunk_size: u32,
|
|
/// Bitmask of filters that were NOT applied (0 = all applied).
|
|
pub filter_mask: u32,
|
|
/// N-dimensional offset of this chunk in dataset space.
|
|
pub offsets: Vec<u64>,
|
|
/// File address of the chunk data.
|
|
pub address: u64,
|
|
}
|
|
|
|
/// Checks that `[offset, offset + needed)` fits within `data`, guarding the
|
|
/// addition against `usize` overflow from a crafted near-`usize::MAX` offset.
|
|
fn ensure_len(data: &[u8], offset: usize, needed: usize) -> Result<(), FormatError> {
|
|
if offset
|
|
.checked_add(needed)
|
|
.is_none_or(|end| end > data.len())
|
|
{
|
|
return Err(FormatError::UnexpectedEof {
|
|
expected: offset.saturating_add(needed),
|
|
available: data.len(),
|
|
});
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// `elements * elem_size` for sizes that come from the file. Dataspace and
|
|
/// chunk dimensions are untrusted 64-bit fields, so a crafted file can make
|
|
/// the plain product wrap to a small number (or to something enormous).
|
|
pub(crate) fn checked_byte_len(elements: u64, elem_size: usize) -> Result<usize, FormatError> {
|
|
usize::try_from(elements)
|
|
.ok()
|
|
.and_then(|n| n.checked_mul(elem_size))
|
|
.ok_or_else(|| {
|
|
FormatError::Overflow(format!(
|
|
"{elements} elements of {elem_size} bytes exceeds the addressable size"
|
|
))
|
|
})
|
|
}
|
|
|
|
/// Product of chunk dimensions times the element size, overflow-checked.
|
|
pub(crate) fn checked_chunk_byte_len(
|
|
chunk_dims: &[usize],
|
|
elem_size: usize,
|
|
) -> Result<usize, FormatError> {
|
|
chunk_dims
|
|
.iter()
|
|
.try_fold(elem_size, |acc, &d| acc.checked_mul(d))
|
|
.ok_or_else(|| {
|
|
FormatError::Overflow(format!(
|
|
"chunk dimensions {chunk_dims:?} x {elem_size} bytes exceeds the addressable size"
|
|
))
|
|
})
|
|
}
|
|
|
|
/// A zero-filled output buffer of `len` bytes. `vec![0; len]` aborts the
|
|
/// process when the allocation fails; a size taken from the file must surface
|
|
/// as an error instead.
|
|
pub(crate) fn alloc_output(len: usize) -> Result<Vec<u8>, FormatError> {
|
|
if len == 0 {
|
|
return Ok(Vec::new());
|
|
}
|
|
let failed =
|
|
|| FormatError::Overflow(format!("cannot allocate {len} bytes for dataset output"));
|
|
let layout = core::alloc::Layout::array::<u8>(len).map_err(|_| failed())?;
|
|
// Ask the allocator for zeroed memory instead of reserving and then
|
|
// writing zeros: for a large buffer the OS hands out already-zero pages
|
|
// lazily, where an explicit fill touches every page up front — and most of
|
|
// the buffer is about to be overwritten with chunk data anyway.
|
|
//
|
|
// SAFETY (both arms): `layout` has non-zero size (len > 0) and alignment 1.
|
|
#[cfg(feature = "std")]
|
|
let ptr = unsafe { std::alloc::alloc_zeroed(layout) };
|
|
#[cfg(not(feature = "std"))]
|
|
let ptr = unsafe { alloc::alloc::alloc_zeroed(layout) };
|
|
if ptr.is_null() {
|
|
return Err(failed());
|
|
}
|
|
// SAFETY: `ptr` came from the global allocator with the layout of
|
|
// `[u8; len]`, which is exactly what `Vec<u8>` with capacity `len` frees;
|
|
// all `len` bytes are initialised (zero).
|
|
Ok(unsafe { Vec::from_raw_parts(ptr, len, len) })
|
|
}
|
|
|
|
fn read_offset(data: &[u8], pos: usize, size: u8) -> Result<u64, FormatError> {
|
|
let s = size as usize;
|
|
if pos.checked_add(s).is_none_or(|end| end > data.len()) {
|
|
return Err(FormatError::UnexpectedEof {
|
|
expected: pos.saturating_add(s),
|
|
available: data.len(),
|
|
});
|
|
}
|
|
let slice = &data[pos..pos + s];
|
|
Ok(match size {
|
|
2 => u16::from_le_bytes([slice[0], slice[1]]) as u64,
|
|
4 => u32::from_le_bytes([slice[0], slice[1], slice[2], slice[3]]) as u64,
|
|
8 => u64::from_le_bytes([
|
|
slice[0], slice[1], slice[2], slice[3], slice[4], slice[5], slice[6], slice[7],
|
|
]),
|
|
_ => return Err(FormatError::InvalidOffsetSize(size)),
|
|
})
|
|
}
|
|
|
|
/// Traverse B-tree v1 type 1 to collect all chunk locations.
|
|
///
|
|
/// `ndims` is the number of offset dimensions in each key, which equals
|
|
/// `chunk_dimensions.len()` from the DataLayout::Chunked message (rank+1).
|
|
pub fn collect_chunk_info(
|
|
file_data: &[u8],
|
|
btree_address: u64,
|
|
ndims: usize,
|
|
offset_size: u8,
|
|
length_size: u8,
|
|
) -> Result<Vec<ChunkInfo>, FormatError> {
|
|
collect_chunk_info_inner(file_data, btree_address, ndims, offset_size, length_size, 0)
|
|
}
|
|
|
|
/// 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;
|
|
|
|
fn collect_chunk_info_inner(
|
|
file_data: &[u8],
|
|
btree_address: u64,
|
|
ndims: usize,
|
|
offset_size: u8,
|
|
_length_size: u8,
|
|
depth: usize,
|
|
) -> Result<Vec<ChunkInfo>, FormatError> {
|
|
if depth > MAX_CHUNK_BTREE_DEPTH {
|
|
return Err(FormatError::NestingDepthExceeded);
|
|
}
|
|
|
|
let offset = btree_address as usize;
|
|
let os = offset_size as usize;
|
|
|
|
// Parse B-tree v1 header
|
|
let header_size = 8 + os * 2;
|
|
ensure_len(file_data, offset, header_size)?;
|
|
|
|
if &file_data[offset..offset + 4] != b"TREE" {
|
|
return Err(FormatError::InvalidBTreeSignature);
|
|
}
|
|
|
|
let node_type = file_data[offset + 4];
|
|
if node_type != 1 {
|
|
return Err(FormatError::InvalidBTreeNodeType(node_type));
|
|
}
|
|
|
|
let node_level = file_data[offset + 5];
|
|
let entries_used = u16::from_le_bytes([file_data[offset + 6], file_data[offset + 7]]) as usize;
|
|
|
|
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;
|
|
|
|
if node_level == 0 {
|
|
// Leaf node: keys and children interleaved
|
|
// key[0], child[0], key[1], child[1], ..., key[N-1], child[N-1], key[N]
|
|
let needed = entries_used * (key_size + os) + key_size;
|
|
ensure_len(file_data, pos, needed)?;
|
|
|
|
let mut chunks = Vec::with_capacity(entries_used);
|
|
for _ in 0..entries_used {
|
|
// Parse key
|
|
let chunk_size = u32::from_le_bytes([
|
|
file_data[pos],
|
|
file_data[pos + 1],
|
|
file_data[pos + 2],
|
|
file_data[pos + 3],
|
|
]);
|
|
let filter_mask = u32::from_le_bytes([
|
|
file_data[pos + 4],
|
|
file_data[pos + 5],
|
|
file_data[pos + 6],
|
|
file_data[pos + 7],
|
|
]);
|
|
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;
|
|
}
|
|
pos += key_size;
|
|
|
|
// Parse child address
|
|
let address = read_offset(file_data, pos, offset_size)?;
|
|
pos += os;
|
|
|
|
chunks.push(ChunkInfo {
|
|
chunk_size,
|
|
filter_mask,
|
|
offsets,
|
|
address,
|
|
});
|
|
}
|
|
// Skip final key
|
|
Ok(chunks)
|
|
} else {
|
|
// Internal node: recurse into children
|
|
let needed = entries_used * (key_size + os) + key_size;
|
|
ensure_len(file_data, pos, needed)?;
|
|
|
|
let mut child_addrs = Vec::with_capacity(entries_used);
|
|
for _ in 0..entries_used {
|
|
pos += key_size; // skip key
|
|
let child_addr = read_offset(file_data, pos, offset_size)?;
|
|
child_addrs.push(child_addr);
|
|
pos += os;
|
|
}
|
|
|
|
let mut all_chunks = Vec::new();
|
|
for child_addr in child_addrs {
|
|
let child_chunks = collect_chunk_info_inner(
|
|
file_data,
|
|
child_addr,
|
|
ndims,
|
|
offset_size,
|
|
_length_size,
|
|
depth + 1,
|
|
)?;
|
|
all_chunks.extend(child_chunks);
|
|
}
|
|
Ok(all_chunks)
|
|
}
|
|
}
|
|
|
|
/// Generate ChunkInfo entries for an implicit index (v4 index type 2).
|
|
///
|
|
/// Chunks are stored contiguously starting at `base_address`. No stored index;
|
|
/// addresses are computed from the chunk position.
|
|
pub fn generate_implicit_chunks(
|
|
base_address: u64,
|
|
dataset_dims: &[u64],
|
|
chunk_dimensions: &[u32],
|
|
element_size: u32,
|
|
) -> Vec<ChunkInfo> {
|
|
let rank = chunk_dimensions.len();
|
|
let chunk_byte_size: u64 =
|
|
chunk_dimensions.iter().map(|&d| d as u64).product::<u64>() * element_size as u64;
|
|
|
|
let mut num_chunks_per_dim = Vec::with_capacity(rank);
|
|
for d in 0..rank {
|
|
let ds = dataset_dims[d];
|
|
let ch = chunk_dimensions[d] as u64;
|
|
num_chunks_per_dim.push(ds.div_ceil(ch));
|
|
}
|
|
let total_chunks: u64 = num_chunks_per_dim.iter().product();
|
|
|
|
let mut chunks = Vec::with_capacity(total_chunks as usize);
|
|
for linear_idx in 0..total_chunks {
|
|
let mut offsets = vec![0u64; rank];
|
|
let mut remaining = linear_idx;
|
|
for d in (0..rank).rev() {
|
|
let nchunks = num_chunks_per_dim[d];
|
|
let chunk_idx = remaining % nchunks;
|
|
remaining /= nchunks;
|
|
offsets[d] = chunk_idx * chunk_dimensions[d] as u64;
|
|
}
|
|
|
|
chunks.push(ChunkInfo {
|
|
chunk_size: chunk_byte_size as u32,
|
|
filter_mask: 0,
|
|
offsets,
|
|
address: base_address + linear_idx * chunk_byte_size,
|
|
});
|
|
}
|
|
|
|
chunks
|
|
}
|
|
|
|
/// Read a chunked dataset, decompressing chunks as needed.
|
|
/// Chunks decompressed together before being copied out, bounding the extra
|
|
/// memory a parallel full read holds at once.
|
|
const DECODE_BATCH: usize = 128;
|
|
|
|
/// B-tree v2 record types used for chunk indexing.
|
|
const BT2_CHUNK_UNFILTERED: u8 = 10;
|
|
const BT2_CHUNK_FILTERED: u8 = 11;
|
|
|
|
/// Chunks indexed by a version-2 B-tree (layout v4, index type 5).
|
|
///
|
|
/// Record layouts (all little endian):
|
|
/// * type 10, unfiltered: address, then one 8-byte *scaled* offset per
|
|
/// dimension (offset / chunk dimension);
|
|
/// * type 11, filtered: address, stored chunk size (a variable number of
|
|
/// bytes), 4-byte filter mask, then the scaled offsets.
|
|
///
|
|
/// The width of the stored-size field depends on the largest possible chunk;
|
|
/// rather than re-derive the library's formula it is taken from the record
|
|
/// size the tree header declares, which is what actually governs the bytes.
|
|
fn read_btree_v2_chunks(
|
|
file_data: &[u8],
|
|
addr: u64,
|
|
chunk_dims: &[usize],
|
|
elem_size: usize,
|
|
offset_size: u8,
|
|
length_size: u8,
|
|
) -> Result<Vec<ChunkInfo>, FormatError> {
|
|
use crate::btree_v2::{BTreeV2Header, collect_btree_v2_records};
|
|
|
|
let bad = |what: &str| FormatError::ChunkedReadError(format!("B-tree v2 chunk index: {what}"));
|
|
let header = BTreeV2Header::parse(file_data, addr as usize, offset_size, length_size)?;
|
|
let rank = chunk_dims.len();
|
|
let os = offset_size as usize;
|
|
let record_size = header.record_size as usize;
|
|
let size_len = match header.tree_type {
|
|
BT2_CHUNK_UNFILTERED => {
|
|
if record_size != os + 8 * rank {
|
|
return Err(bad("unexpected record size for unfiltered chunks"));
|
|
}
|
|
0
|
|
}
|
|
BT2_CHUNK_FILTERED => {
|
|
let fixed = os + 4 + 8 * rank;
|
|
let size_len = record_size
|
|
.checked_sub(fixed)
|
|
.ok_or_else(|| bad("record too small"))?;
|
|
if !(1..=8).contains(&size_len) {
|
|
return Err(bad("implausible chunk-size field width"));
|
|
}
|
|
size_len
|
|
}
|
|
_ => return Err(bad("tree is not a chunk index")),
|
|
};
|
|
let unfiltered_bytes = checked_chunk_byte_len(chunk_dims, elem_size)?;
|
|
let unfiltered_bytes =
|
|
u32::try_from(unfiltered_bytes).map_err(|_| bad("chunk larger than 4 GiB"))?;
|
|
|
|
let records = collect_btree_v2_records(file_data, &header, offset_size, length_size)?;
|
|
let mut chunks = Vec::with_capacity(records.len());
|
|
for record in &records {
|
|
let data = record.data.as_slice();
|
|
if data.len() < record_size {
|
|
return Err(bad("truncated record"));
|
|
}
|
|
let address = read_offset(data, 0, offset_size)?;
|
|
let mut pos = os;
|
|
let (chunk_size, filter_mask) = if size_len == 0 {
|
|
(unfiltered_bytes, 0)
|
|
} else {
|
|
let mut size = 0u64;
|
|
for (i, &b) in data[pos..pos + size_len].iter().enumerate() {
|
|
size |= u64::from(b) << (8 * i);
|
|
}
|
|
pos += size_len;
|
|
let mask = u32::from_le_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]]);
|
|
pos += 4;
|
|
(
|
|
u32::try_from(size).map_err(|_| bad("stored chunk larger than 4 GiB"))?,
|
|
mask,
|
|
)
|
|
};
|
|
let mut offsets = Vec::with_capacity(rank);
|
|
for &dim in chunk_dims {
|
|
let scaled = u64::from_le_bytes([
|
|
data[pos],
|
|
data[pos + 1],
|
|
data[pos + 2],
|
|
data[pos + 3],
|
|
data[pos + 4],
|
|
data[pos + 5],
|
|
data[pos + 6],
|
|
data[pos + 7],
|
|
]);
|
|
pos += 8;
|
|
offsets.push(
|
|
scaled
|
|
.checked_mul(dim as u64)
|
|
.ok_or_else(|| bad("chunk offset overflows"))?,
|
|
);
|
|
}
|
|
chunks.push(ChunkInfo {
|
|
chunk_size,
|
|
filter_mask,
|
|
offsets,
|
|
address,
|
|
});
|
|
}
|
|
Ok(chunks)
|
|
}
|
|
|
|
/// Every allocated chunk of a chunked dataset, for any supported chunk index,
|
|
/// plus the spatial chunk dimensions. Chunks the file never allocated (sparse
|
|
/// datasets) are simply absent from the list.
|
|
pub fn list_chunks(
|
|
file_data: &[u8],
|
|
layout: &DataLayout,
|
|
dataspace: &Dataspace,
|
|
elem_size: usize,
|
|
offset_size: u8,
|
|
length_size: u8,
|
|
) -> Result<(Vec<ChunkInfo>, Vec<usize>), FormatError> {
|
|
let (
|
|
chunk_dimensions,
|
|
version,
|
|
chunk_index_type,
|
|
addr_opt,
|
|
single_filtered_size,
|
|
single_filter_mask,
|
|
) = match layout {
|
|
DataLayout::Chunked {
|
|
chunk_dimensions,
|
|
btree_address,
|
|
version,
|
|
chunk_index_type,
|
|
single_chunk_filtered_size,
|
|
single_chunk_filter_mask,
|
|
} => (
|
|
chunk_dimensions,
|
|
*version,
|
|
*chunk_index_type,
|
|
*btree_address,
|
|
*single_chunk_filtered_size,
|
|
*single_chunk_filter_mask,
|
|
),
|
|
_ => {
|
|
return Err(FormatError::ChunkedReadError(
|
|
"expected chunked layout".into(),
|
|
));
|
|
}
|
|
};
|
|
|
|
let addr = addr_opt
|
|
.ok_or_else(|| FormatError::ChunkedReadError("no address for chunked layout".into()))?;
|
|
|
|
// Both v3 and v4 include element size as last dim (rank+1)
|
|
let ndims = chunk_dimensions.len();
|
|
let rank = ndims
|
|
.checked_sub(1)
|
|
.ok_or_else(|| FormatError::ChunkedReadError("chunked layout has no dimensions".into()))?;
|
|
let chunk_dims: Vec<usize> = chunk_dimensions[..rank]
|
|
.iter()
|
|
.map(|&d| d as usize)
|
|
.collect();
|
|
|
|
let ds_dims: Vec<usize> = dataspace.dimensions.iter().map(|&d| d as usize).collect();
|
|
if ds_dims.len() != rank {
|
|
return Err(FormatError::ChunkedReadError(format!(
|
|
"rank mismatch: dataspace has {} dims, layout has {} chunk dims (rank={})",
|
|
ds_dims.len(),
|
|
chunk_dimensions.len(),
|
|
rank
|
|
)));
|
|
}
|
|
|
|
// Collect chunks based on version and index type
|
|
let 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)?
|
|
}
|
|
(4, Some(1)) => {
|
|
// Single chunk — one chunk covering the entire dataset
|
|
let chunk_byte_size = checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
|
let (csize, fmask) = if let Some(fs) = single_filtered_size {
|
|
(fs as u32, single_filter_mask.unwrap_or(0))
|
|
} else {
|
|
(chunk_byte_size as u32, 0)
|
|
};
|
|
vec![ChunkInfo {
|
|
chunk_size: csize,
|
|
filter_mask: fmask,
|
|
offsets: vec![0u64; rank],
|
|
address: addr,
|
|
}]
|
|
}
|
|
(4, Some(2)) => {
|
|
// Implicit index — use spatial chunk dims only
|
|
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
|
generate_implicit_chunks(
|
|
addr,
|
|
&dataspace.dimensions,
|
|
spatial_chunk_dims,
|
|
elem_size as u32,
|
|
)
|
|
}
|
|
(4, Some(3)) => {
|
|
// Fixed Array — use spatial chunk dims only
|
|
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
|
let header =
|
|
FixedArrayHeader::parse(file_data, addr as usize, offset_size, length_size)?;
|
|
read_fixed_array_chunks(
|
|
file_data,
|
|
&header,
|
|
&dataspace.dimensions,
|
|
spatial_chunk_dims,
|
|
elem_size as u32,
|
|
offset_size,
|
|
length_size,
|
|
)?
|
|
}
|
|
(4, Some(4)) => {
|
|
// Extensible Array — use spatial chunk dims only
|
|
let spatial_chunk_dims: &[u32] = &chunk_dimensions[..rank];
|
|
let header =
|
|
ExtensibleArrayHeader::parse(file_data, addr as usize, offset_size, length_size)?;
|
|
read_extensible_array_chunks(
|
|
file_data,
|
|
&header,
|
|
&dataspace.dimensions,
|
|
spatial_chunk_dims,
|
|
elem_size as u32,
|
|
offset_size,
|
|
length_size,
|
|
)?
|
|
}
|
|
(4, Some(5)) => {
|
|
// Version-2 B-tree: what the library uses for a dataset with two
|
|
// or more unlimited dimensions.
|
|
read_btree_v2_chunks(
|
|
file_data,
|
|
addr,
|
|
&chunk_dims,
|
|
elem_size,
|
|
offset_size,
|
|
length_size,
|
|
)?
|
|
}
|
|
(v, idx) => {
|
|
return Err(FormatError::ChunkedReadError(format!(
|
|
"unsupported chunked layout version={v}, index_type={idx:?}"
|
|
)));
|
|
}
|
|
};
|
|
|
|
Ok((chunks, chunk_dims))
|
|
}
|
|
|
|
pub fn read_chunked_data(
|
|
file_data: &[u8],
|
|
layout: &DataLayout,
|
|
dataspace: &Dataspace,
|
|
datatype: &Datatype,
|
|
pipeline: Option<&FilterPipeline>,
|
|
offset_size: u8,
|
|
length_size: u8,
|
|
) -> Result<Vec<u8>, FormatError> {
|
|
let elem_size = datatype.type_size() as usize;
|
|
let (chunks, chunk_dims) = list_chunks(
|
|
file_data,
|
|
layout,
|
|
dataspace,
|
|
elem_size,
|
|
offset_size,
|
|
length_size,
|
|
)?;
|
|
let rank = chunk_dims.len();
|
|
let ds_dims: Vec<usize> = dataspace.dimensions.iter().map(|&d| d as usize).collect();
|
|
|
|
// Assemble output
|
|
let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?;
|
|
if total_bytes == 0 {
|
|
// Also keeps the stride products below in range: with a zero-sized
|
|
// dimension the total is 0 even if other dimensions are huge.
|
|
return Ok(Vec::new());
|
|
}
|
|
let mut output = alloc_output(total_bytes)?;
|
|
|
|
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];
|
|
}
|
|
|
|
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_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
|
|
|
// Fast path: no filters — copy directly from file_data without intermediate alloc
|
|
if pipeline.is_none() {
|
|
for chunk_info in &chunks {
|
|
let chunk_offsets: Vec<usize> = chunk_info
|
|
.offsets
|
|
.iter()
|
|
.take(rank)
|
|
.map(|&o| o as usize)
|
|
.collect();
|
|
|
|
let c_addr = chunk_info.address as usize;
|
|
let size = chunk_info.chunk_size as usize;
|
|
ensure_len(file_data, c_addr, size)?;
|
|
let chunk_data = &file_data[c_addr..c_addr + size];
|
|
|
|
if rank == 0 {
|
|
let copy_len = chunk_data.len().min(output.len());
|
|
output[..copy_len].copy_from_slice(&chunk_data[..copy_len]);
|
|
} else {
|
|
copy_chunk_to_output(
|
|
chunk_data,
|
|
&mut output,
|
|
&chunk_offsets,
|
|
&chunk_dims,
|
|
&ds_dims,
|
|
&ds_strides,
|
|
&chunk_strides,
|
|
elem_size,
|
|
rank,
|
|
);
|
|
}
|
|
}
|
|
return Ok(output);
|
|
}
|
|
|
|
// Filtered path: decompress all chunks then assemble
|
|
let decompressed_chunks = decompress_all_chunks(
|
|
file_data,
|
|
&chunks,
|
|
pipeline,
|
|
chunk_total_bytes,
|
|
elem_size as u32,
|
|
)?;
|
|
|
|
for (chunk_info, decompressed) in chunks.iter().zip(decompressed_chunks.iter()) {
|
|
let chunk_offsets: Vec<usize> = chunk_info
|
|
.offsets
|
|
.iter()
|
|
.take(rank)
|
|
.map(|&o| o as usize)
|
|
.collect();
|
|
|
|
if rank == 0 {
|
|
let copy_len = decompressed.len().min(output.len());
|
|
output[..copy_len].copy_from_slice(&decompressed[..copy_len]);
|
|
} else {
|
|
copy_chunk_to_output(
|
|
decompressed,
|
|
&mut output,
|
|
&chunk_offsets,
|
|
&chunk_dims,
|
|
&ds_dims,
|
|
&ds_strides,
|
|
&chunk_strides,
|
|
elem_size,
|
|
rank,
|
|
);
|
|
}
|
|
}
|
|
|
|
Ok(output)
|
|
}
|
|
|
|
/// Read a chunked dataset with caching support.
|
|
///
|
|
/// On the first call, scans the chunk index (B-tree / fixed array / etc.) once
|
|
/// and populates the cache's hash index. Subsequent calls skip the index scan
|
|
/// entirely. Decompressed chunk data is also cached with LRU eviction.
|
|
#[cfg(feature = "std")]
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn read_chunked_data_cached(
|
|
file_data: &[u8],
|
|
layout: &DataLayout,
|
|
dataspace: &Dataspace,
|
|
datatype: &Datatype,
|
|
pipeline: Option<&FilterPipeline>,
|
|
offset_size: u8,
|
|
length_size: u8,
|
|
cache: &ChunkCache,
|
|
) -> Result<Vec<u8>, FormatError> {
|
|
let (chunk_dimensions, addr_opt) = match layout {
|
|
DataLayout::Chunked {
|
|
chunk_dimensions,
|
|
btree_address,
|
|
..
|
|
} => (chunk_dimensions, *btree_address),
|
|
_ => {
|
|
return Err(FormatError::ChunkedReadError(
|
|
"expected chunked layout".into(),
|
|
));
|
|
}
|
|
};
|
|
|
|
let addr = addr_opt
|
|
.ok_or_else(|| FormatError::ChunkedReadError("no address for chunked layout".into()))?;
|
|
|
|
let elem_size = datatype.type_size() as usize;
|
|
let ndims = chunk_dimensions.len();
|
|
let rank = ndims
|
|
.checked_sub(1)
|
|
.ok_or_else(|| FormatError::ChunkedReadError("chunked layout has no dimensions".into()))?;
|
|
let chunk_dims: Vec<usize> = chunk_dimensions[..rank]
|
|
.iter()
|
|
.map(|&d| d as usize)
|
|
.collect();
|
|
|
|
let ds_dims: Vec<usize> = dataspace.dimensions.iter().map(|&d| d as usize).collect();
|
|
if ds_dims.len() != rank {
|
|
return Err(FormatError::ChunkedReadError(format!(
|
|
"rank mismatch: dataspace has {} dims, layout has {} chunk dims (rank={})",
|
|
ds_dims.len(),
|
|
chunk_dimensions.len(),
|
|
rank
|
|
)));
|
|
}
|
|
|
|
// 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(
|
|
file_data,
|
|
layout,
|
|
dataspace,
|
|
elem_size,
|
|
offset_size,
|
|
length_size,
|
|
)?;
|
|
cache.populate_index(&chunks, rank);
|
|
}
|
|
|
|
let chunks = cache.all_indexed_chunks().unwrap_or_default();
|
|
|
|
// Assemble output
|
|
let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?;
|
|
if total_bytes == 0 {
|
|
// Also keeps the stride products below in range: with a zero-sized
|
|
// dimension the total is 0 even if other dimensions are huge.
|
|
return Ok(Vec::new());
|
|
}
|
|
let mut output = alloc_output(total_bytes)?;
|
|
|
|
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];
|
|
}
|
|
|
|
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_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
|
|
|
let mut place = |data: &[u8], chunk_info: &ChunkInfo| {
|
|
if rank == 0 {
|
|
let copy_len = data.len().min(output.len());
|
|
output[..copy_len].copy_from_slice(&data[..copy_len]);
|
|
return;
|
|
}
|
|
let chunk_offsets: Vec<usize> = chunk_info
|
|
.offsets
|
|
.iter()
|
|
.take(rank)
|
|
.map(|&o| o as usize)
|
|
.collect();
|
|
copy_chunk_to_output(
|
|
data,
|
|
&mut output,
|
|
&chunk_offsets,
|
|
&chunk_dims,
|
|
&ds_dims,
|
|
&ds_strides,
|
|
&chunk_strides,
|
|
elem_size,
|
|
rank,
|
|
);
|
|
};
|
|
let raw_bytes = |chunk_info: &ChunkInfo| -> Result<&[u8], FormatError> {
|
|
let c_addr = chunk_info.address as usize;
|
|
let size = chunk_info.chunk_size as usize;
|
|
ensure_len(file_data, c_addr, size)?;
|
|
Ok(&file_data[c_addr..c_addr + size])
|
|
};
|
|
|
|
// 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;
|
|
let mut misses: Vec<&ChunkInfo> = Vec::new();
|
|
for chunk_info in &chunks {
|
|
if stored_raw(chunk_info) {
|
|
place(raw_bytes(chunk_info)?, chunk_info);
|
|
continue;
|
|
}
|
|
let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();
|
|
match cache.get_decompressed_aligned(&coord) {
|
|
Some(cached) => place(&cached, chunk_info),
|
|
None => misses.push(chunk_info),
|
|
}
|
|
}
|
|
|
|
// Decompress what the cache didn't have, a bounded batch at a time — in
|
|
// parallel with the `parallel` feature (this path, the one the facade
|
|
// uses, was sequential; only the uncached reader was parallel). Chunks are
|
|
// cached only when the whole dataset fits: pushing a larger dataset
|
|
// through the cache just evicts each chunk moments after inserting it.
|
|
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)
|
|
};
|
|
for batch in misses.chunks(DECODE_BATCH) {
|
|
#[cfg(feature = "parallel")]
|
|
let decoded: Vec<Result<Vec<u8>, FormatError>> = if batch.len() >= 4 {
|
|
use rayon::prelude::*;
|
|
batch.par_iter().map(decode).collect()
|
|
} else {
|
|
batch.iter().map(decode).collect()
|
|
};
|
|
#[cfg(not(feature = "parallel"))]
|
|
let decoded: Vec<Result<Vec<u8>, FormatError>> = batch.iter().map(decode).collect();
|
|
|
|
for (chunk_info, data) in batch.iter().zip(decoded) {
|
|
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);
|
|
place(&cached, chunk_info);
|
|
} else {
|
|
place(&data, chunk_info);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(output)
|
|
}
|
|
|
|
/// Sweep context passed into `read_chunked_data_sweep` to enable adaptive
|
|
/// prefetching based on detected access patterns.
|
|
///
|
|
/// The caller is responsible for maintaining the `SweepContext` across
|
|
/// multiple reads on the same dataset. After each read, the context will
|
|
/// contain updated sweep detection state and any predicted next-chunk
|
|
/// coordinates.
|
|
pub struct SweepContext {
|
|
/// Sliding window of recent chunk coordinates.
|
|
pub history: Vec<Vec<u64>>,
|
|
/// Maximum window size.
|
|
pub window_size: usize,
|
|
/// Currently detected sweep direction label.
|
|
pub direction: &'static str,
|
|
/// How many chunks ahead to predict.
|
|
pub prefetch_count: usize,
|
|
/// Predicted next chunk coordinates (populated after each read).
|
|
pub predicted_next: Vec<Vec<u64>>,
|
|
}
|
|
|
|
impl SweepContext {
|
|
/// Create a new sweep context with the given window size and prefetch count.
|
|
pub fn new(window_size: usize, prefetch_count: usize) -> Self {
|
|
Self {
|
|
history: Vec::with_capacity(window_size),
|
|
window_size,
|
|
direction: "random",
|
|
prefetch_count,
|
|
predicted_next: Vec::new(),
|
|
}
|
|
}
|
|
|
|
/// Create with default settings (window=12, prefetch=4).
|
|
pub fn with_defaults() -> Self {
|
|
Self::new(12, 4)
|
|
}
|
|
|
|
/// Record a chunk coordinate access and update predictions.
|
|
fn record(&mut self, coord: Vec<u64>, ndims: usize) {
|
|
if self.history.len() >= self.window_size {
|
|
self.history.remove(0);
|
|
}
|
|
self.history.push(coord);
|
|
|
|
if self.history.len() < 3 || ndims == 0 {
|
|
self.direction = "random";
|
|
self.predicted_next.clear();
|
|
return;
|
|
}
|
|
|
|
// Inline sweep detection matching the algorithm in clawhdf5-io/sweep.rs
|
|
let num_deltas = self.history.len() - 1;
|
|
let mut changing = vec![0usize; ndims];
|
|
for i in 0..num_deltas {
|
|
let prev = &self.history[i];
|
|
let curr = &self.history[i + 1];
|
|
if prev.len() < ndims || curr.len() < ndims {
|
|
self.direction = "random";
|
|
self.predicted_next.clear();
|
|
return;
|
|
}
|
|
for d in 0..ndims {
|
|
if curr[d] != prev[d] {
|
|
changing[d] += 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
let threshold = num_deltas.div_ceil(2);
|
|
let (max_dim, max_changes) = changing.iter().enumerate().max_by_key(|(_, c)| *c).unwrap();
|
|
|
|
if *max_changes < threshold {
|
|
self.direction = "random";
|
|
self.predicted_next.clear();
|
|
return;
|
|
}
|
|
|
|
let others_max = changing
|
|
.iter()
|
|
.enumerate()
|
|
.filter(|(d, _)| *d != max_dim)
|
|
.map(|(_, c)| *c)
|
|
.max()
|
|
.unwrap_or(0);
|
|
|
|
if others_max > 0 && *max_changes < others_max * 2 {
|
|
self.direction = "random";
|
|
self.predicted_next.clear();
|
|
return;
|
|
}
|
|
|
|
self.direction = if max_dim == ndims - 1 {
|
|
"row_major"
|
|
} else if max_dim == 0 {
|
|
"column_major"
|
|
} else {
|
|
"slice_major"
|
|
};
|
|
|
|
// Predict next chunks
|
|
let sweep_dim = max_dim;
|
|
let mut total_step: i64 = 0;
|
|
let mut step_count: usize = 0;
|
|
for i in 1..self.history.len() {
|
|
let prev = self.history[i - 1][sweep_dim] as i64;
|
|
let curr = self.history[i][sweep_dim] as i64;
|
|
let diff = curr - prev;
|
|
if diff != 0 {
|
|
total_step += diff;
|
|
step_count += 1;
|
|
}
|
|
}
|
|
|
|
if step_count == 0 {
|
|
self.predicted_next.clear();
|
|
return;
|
|
}
|
|
|
|
let avg_step = total_step / step_count as i64;
|
|
if avg_step == 0 {
|
|
self.predicted_next.clear();
|
|
return;
|
|
}
|
|
|
|
let last = self.history.last().unwrap();
|
|
self.predicted_next.clear();
|
|
for i in 1..=self.prefetch_count {
|
|
let mut pred = last.clone();
|
|
let new_val = last[sweep_dim] as i64 + avg_step * i as i64;
|
|
if new_val < 0 {
|
|
break;
|
|
}
|
|
pred[sweep_dim] = new_val as u64;
|
|
self.predicted_next.push(pred);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Read a chunked dataset with caching and sweep-aware prefetching.
|
|
///
|
|
/// Extends `read_chunked_data_cached` by feeding each chunk coordinate to a
|
|
/// [`SweepContext`]. When a sweep pattern is detected, predicted next-chunk
|
|
/// coordinates are pre-populated in the cache index via `prefetch_hint`.
|
|
#[cfg(feature = "std")]
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn read_chunked_data_sweep(
|
|
file_data: &[u8],
|
|
layout: &DataLayout,
|
|
dataspace: &Dataspace,
|
|
datatype: &Datatype,
|
|
pipeline: Option<&FilterPipeline>,
|
|
offset_size: u8,
|
|
length_size: u8,
|
|
cache: &ChunkCache,
|
|
sweep: &mut SweepContext,
|
|
) -> Result<Vec<u8>, FormatError> {
|
|
let (chunk_dimensions, addr_opt) = match layout {
|
|
DataLayout::Chunked {
|
|
chunk_dimensions,
|
|
btree_address,
|
|
..
|
|
} => (chunk_dimensions, *btree_address),
|
|
_ => {
|
|
return Err(FormatError::ChunkedReadError(
|
|
"expected chunked layout".into(),
|
|
));
|
|
}
|
|
};
|
|
|
|
let addr = addr_opt
|
|
.ok_or_else(|| FormatError::ChunkedReadError("no address for chunked layout".into()))?;
|
|
|
|
let elem_size = datatype.type_size() as usize;
|
|
let ndims = chunk_dimensions.len();
|
|
let rank = ndims
|
|
.checked_sub(1)
|
|
.ok_or_else(|| FormatError::ChunkedReadError("chunked layout has no dimensions".into()))?;
|
|
let chunk_dims: Vec<usize> = chunk_dimensions[..rank]
|
|
.iter()
|
|
.map(|&d| d as usize)
|
|
.collect();
|
|
|
|
let ds_dims: Vec<usize> = dataspace.dimensions.iter().map(|&d| d as usize).collect();
|
|
if ds_dims.len() != rank {
|
|
return Err(FormatError::ChunkedReadError(format!(
|
|
"rank mismatch: dataspace has {} dims, layout has {} chunk dims (rank={})",
|
|
ds_dims.len(),
|
|
chunk_dimensions.len(),
|
|
rank
|
|
)));
|
|
}
|
|
|
|
// 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(
|
|
file_data,
|
|
layout,
|
|
dataspace,
|
|
elem_size,
|
|
offset_size,
|
|
length_size,
|
|
)?;
|
|
cache.populate_index(&chunks, rank);
|
|
}
|
|
|
|
let chunks = cache.all_indexed_chunks().unwrap_or_default();
|
|
|
|
// Assemble output
|
|
let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?;
|
|
if total_bytes == 0 {
|
|
// Also keeps the stride products below in range: with a zero-sized
|
|
// dimension the total is 0 even if other dimensions are huge.
|
|
return Ok(Vec::new());
|
|
}
|
|
let mut output = alloc_output(total_bytes)?;
|
|
|
|
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];
|
|
}
|
|
|
|
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_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?;
|
|
|
|
for chunk_info in &chunks {
|
|
let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();
|
|
|
|
// Feed coordinate to sweep detector
|
|
sweep.record(coord.clone(), rank);
|
|
|
|
// Issue prefetch hint for predicted next chunks
|
|
if !sweep.predicted_next.is_empty() {
|
|
cache.prefetch_hint(&sweep.predicted_next);
|
|
cache.set_sweep_direction(sweep.direction);
|
|
}
|
|
|
|
// Try decompressed cache first
|
|
let decompressed = if let Some(cached) = cache.get_decompressed_aligned(&coord) {
|
|
cached
|
|
} else {
|
|
// Decompress from file
|
|
let c_addr = chunk_info.address as usize;
|
|
let size = chunk_info.chunk_size as usize;
|
|
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()
|
|
}
|
|
} else {
|
|
raw_chunk.to_vec()
|
|
};
|
|
cache.put_decompressed(coord, dec)
|
|
};
|
|
|
|
let chunk_offsets: Vec<usize> = chunk_info
|
|
.offsets
|
|
.iter()
|
|
.take(rank)
|
|
.map(|&o| o as usize)
|
|
.collect();
|
|
|
|
if rank == 0 {
|
|
let copy_len = decompressed.len().min(output.len());
|
|
output[..copy_len].copy_from_slice(&decompressed[..copy_len]);
|
|
} else {
|
|
copy_chunk_to_output(
|
|
&decompressed,
|
|
&mut output,
|
|
&chunk_offsets,
|
|
&chunk_dims,
|
|
&ds_dims,
|
|
&ds_strides,
|
|
&chunk_strides,
|
|
elem_size,
|
|
rank,
|
|
);
|
|
}
|
|
}
|
|
|
|
Ok(output)
|
|
}
|
|
|
|
/// Read chunked data using pre-computed chunk layout for fast assembly.
|
|
///
|
|
/// This path builds a `ChunkIndex` and `ChunkLayout` on first access (cached
|
|
/// in the `ChunkCache`), then uses the pre-computed row-copy plan for assembly,
|
|
/// avoiding per-element N-D coordinate math on repeated reads.
|
|
#[cfg(feature = "std")]
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn read_chunked_data_indexed(
|
|
file_data: &[u8],
|
|
layout: &DataLayout,
|
|
dataspace: &Dataspace,
|
|
datatype: &Datatype,
|
|
pipeline: Option<&FilterPipeline>,
|
|
offset_size: u8,
|
|
length_size: u8,
|
|
cache: &ChunkCache,
|
|
) -> Result<Vec<u8>, FormatError> {
|
|
let (chunk_dimensions, addr_opt) = match layout {
|
|
DataLayout::Chunked {
|
|
chunk_dimensions,
|
|
btree_address,
|
|
..
|
|
} => (chunk_dimensions, *btree_address),
|
|
_ => {
|
|
return Err(FormatError::ChunkedReadError(
|
|
"expected chunked layout".into(),
|
|
));
|
|
}
|
|
};
|
|
|
|
let addr = addr_opt
|
|
.ok_or_else(|| FormatError::ChunkedReadError("no address for chunked layout".into()))?;
|
|
|
|
let elem_size = datatype.type_size() as usize;
|
|
let ndims = chunk_dimensions.len();
|
|
let rank = ndims
|
|
.checked_sub(1)
|
|
.ok_or_else(|| FormatError::ChunkedReadError("chunked layout has no dimensions".into()))?;
|
|
let chunk_dims: Vec<usize> = chunk_dimensions[..rank]
|
|
.iter()
|
|
.map(|&d| d as usize)
|
|
.collect();
|
|
|
|
let ds_dims: Vec<usize> = dataspace.dimensions.iter().map(|&d| d as usize).collect();
|
|
if ds_dims.len() != rank {
|
|
return Err(FormatError::ChunkedReadError(format!(
|
|
"rank mismatch: dataspace has {} dims, layout has {} chunk dims (rank={})",
|
|
ds_dims.len(),
|
|
chunk_dimensions.len(),
|
|
rank
|
|
)));
|
|
}
|
|
|
|
// 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()))?;
|
|
|
|
// 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) {
|
|
chunk_buffers.push(cached);
|
|
} else {
|
|
let c_addr = *file_offset as usize;
|
|
let size = *file_size as usize;
|
|
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()
|
|
}
|
|
} else {
|
|
raw_chunk.to_vec()
|
|
};
|
|
let aligned = CacheAlignedBuffer::from_vec(decompressed);
|
|
let arc = cache.put_decompressed_aligned(coord.clone(), aligned);
|
|
chunk_buffers.push(arc);
|
|
}
|
|
}
|
|
|
|
// Assemble using pre-computed layout
|
|
let mut output = vec![0u8; 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);
|
|
});
|
|
|
|
Ok(output)
|
|
}
|
|
|
|
/// Copy chunk data into the output buffer at the correct N-D position.
|
|
#[allow(clippy::too_many_arguments)]
|
|
fn copy_chunk_to_output(
|
|
chunk_data: &[u8],
|
|
output: &mut [u8],
|
|
chunk_offsets: &[usize],
|
|
chunk_dims: &[usize],
|
|
ds_dims: &[usize],
|
|
ds_strides: &[usize],
|
|
chunk_strides: &[usize],
|
|
elem_size: usize,
|
|
rank: usize,
|
|
) {
|
|
// Row-copy approach: iterate over outer dimensions, memcpy the innermost
|
|
// dimension in bulk. For 1-D data this is a single memcpy per chunk.
|
|
// For N-D data this is one memcpy per "row" (innermost dim slice).
|
|
|
|
if rank == 1 {
|
|
// Fast path for 1-D: single contiguous copy per chunk
|
|
let global_start = chunk_offsets[0];
|
|
let copy_len = chunk_dims[0].min(ds_dims[0].saturating_sub(global_start));
|
|
let (Some(src_bytes), Some(dst_start)) = (
|
|
copy_len.checked_mul(elem_size),
|
|
global_start.checked_mul(elem_size),
|
|
) else {
|
|
return;
|
|
};
|
|
if src_bytes > 0
|
|
&& dst_start
|
|
.checked_add(src_bytes)
|
|
.is_some_and(|end| end <= output.len())
|
|
&& src_bytes <= chunk_data.len()
|
|
{
|
|
output[dst_start..dst_start + src_bytes].copy_from_slice(&chunk_data[..src_bytes]);
|
|
}
|
|
return;
|
|
}
|
|
|
|
// General N-D: iterate over outer dimensions, copy innermost rows
|
|
let inner_dim = rank - 1;
|
|
let inner_chunk_len =
|
|
chunk_dims[inner_dim].min(ds_dims[inner_dim].saturating_sub(chunk_offsets[inner_dim]));
|
|
let Some(row_bytes) = inner_chunk_len.checked_mul(elem_size) else {
|
|
return;
|
|
};
|
|
|
|
if row_bytes == 0 {
|
|
return;
|
|
}
|
|
|
|
// Number of rows = product of all outer chunk dimensions
|
|
let Some(outer_count) = chunk_dims[..inner_dim]
|
|
.iter()
|
|
.try_fold(1usize, |acc, &d| acc.checked_mul(d))
|
|
else {
|
|
return;
|
|
};
|
|
|
|
// Outer strides for iterating chunk-local coordinates
|
|
let mut outer_strides = vec![1usize; inner_dim];
|
|
for i in (0..inner_dim.saturating_sub(1)).rev() {
|
|
let Some(stride) = outer_strides[i + 1].checked_mul(chunk_dims[i + 1]) else {
|
|
return;
|
|
};
|
|
outer_strides[i] = stride;
|
|
}
|
|
|
|
for outer_idx in 0..outer_count {
|
|
// Convert outer flat index to N-D chunk-local coords for dims 0..inner_dim
|
|
let mut remaining = outer_idx;
|
|
let mut ds_flat = 0usize;
|
|
let mut src_flat = 0usize;
|
|
let mut out_of_bounds = false;
|
|
|
|
for d in 0..inner_dim {
|
|
let coord_in_chunk = if inner_dim > 1 {
|
|
remaining / outer_strides[d]
|
|
} else {
|
|
remaining
|
|
};
|
|
if inner_dim > 1 {
|
|
remaining %= outer_strides[d];
|
|
}
|
|
|
|
let Some(global_coord) = chunk_offsets[d].checked_add(coord_in_chunk) else {
|
|
out_of_bounds = true;
|
|
break;
|
|
};
|
|
if global_coord >= ds_dims[d] {
|
|
out_of_bounds = true;
|
|
break;
|
|
}
|
|
let (Some(ds_term), Some(src_term)) = (
|
|
global_coord.checked_mul(ds_strides[d]),
|
|
coord_in_chunk.checked_mul(chunk_strides[d]),
|
|
) else {
|
|
out_of_bounds = true;
|
|
break;
|
|
};
|
|
let (Some(new_ds_flat), Some(new_src_flat)) =
|
|
(ds_flat.checked_add(ds_term), src_flat.checked_add(src_term))
|
|
else {
|
|
out_of_bounds = true;
|
|
break;
|
|
};
|
|
ds_flat = new_ds_flat;
|
|
src_flat = new_src_flat;
|
|
}
|
|
|
|
if out_of_bounds {
|
|
continue;
|
|
}
|
|
|
|
// Add innermost dimension offset
|
|
let Some(inner_term) = chunk_offsets[inner_dim].checked_mul(ds_strides[inner_dim]) else {
|
|
continue;
|
|
};
|
|
let Some(ds_flat) = ds_flat.checked_add(inner_term) else {
|
|
continue;
|
|
};
|
|
|
|
let (Some(src_start), Some(dst_start)) = (
|
|
src_flat.checked_mul(elem_size),
|
|
ds_flat.checked_mul(elem_size),
|
|
) else {
|
|
continue;
|
|
};
|
|
|
|
let fits = src_start
|
|
.checked_add(row_bytes)
|
|
.is_some_and(|end| end <= chunk_data.len())
|
|
&& dst_start
|
|
.checked_add(row_bytes)
|
|
.is_some_and(|end| end <= output.len());
|
|
if fits {
|
|
output[dst_start..dst_start + row_bytes]
|
|
.copy_from_slice(&chunk_data[src_start..src_start + row_bytes]);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn simple_space(dimensions: Vec<u64>) -> Dataspace {
|
|
Dataspace {
|
|
space_type: crate::dataspace::DataspaceType::Simple,
|
|
rank: dimensions.len() as u8,
|
|
dimensions,
|
|
max_dimensions: None,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn crafted_dimensions_are_errors_not_wraparound() {
|
|
// 2^63 * 2 wraps to 0 with a plain product; 2^40 * 2^40 wraps too.
|
|
for dims in [
|
|
vec![1u64 << 63, 2],
|
|
vec![1 << 40, 1 << 40],
|
|
vec![u64::MAX, u64::MAX],
|
|
] {
|
|
let space = simple_space(dims.clone());
|
|
assert!(
|
|
matches!(space.checked_num_elements(), Err(FormatError::Overflow(_))),
|
|
"{dims:?}"
|
|
);
|
|
// The infallible accessor saturates instead of wrapping.
|
|
assert_eq!(space.num_elements(), u64::MAX, "{dims:?}");
|
|
}
|
|
assert_eq!(simple_space(vec![3, 4]).checked_num_elements().unwrap(), 12);
|
|
// A zero-sized dimension makes the whole product 0, not an overflow.
|
|
assert_eq!(
|
|
simple_space(vec![0, 1 << 40, 1 << 40])
|
|
.checked_num_elements()
|
|
.unwrap(),
|
|
0
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn byte_length_helpers_check_overflow() {
|
|
assert_eq!(checked_byte_len(10, 8).unwrap(), 80);
|
|
assert!(matches!(
|
|
checked_byte_len(u64::MAX, 8),
|
|
Err(FormatError::Overflow(_))
|
|
));
|
|
assert_eq!(checked_chunk_byte_len(&[10, 10], 4).unwrap(), 400);
|
|
assert!(matches!(
|
|
checked_chunk_byte_len(&[usize::MAX, 2], 4),
|
|
Err(FormatError::Overflow(_))
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn unallocatable_output_is_an_error_not_an_abort() {
|
|
assert_eq!(alloc_output(16).unwrap(), vec![0u8; 16]);
|
|
assert!(matches!(
|
|
alloc_output(usize::MAX / 2),
|
|
Err(FormatError::Overflow(_))
|
|
));
|
|
}
|
|
|
|
fn write_offset(buf: &mut Vec<u8>, val: u64, size: u8) {
|
|
match size {
|
|
4 => buf.extend_from_slice(&(val as u32).to_le_bytes()),
|
|
8 => buf.extend_from_slice(&val.to_le_bytes()),
|
|
_ => panic!("unsupported offset size in test"),
|
|
}
|
|
}
|
|
|
|
/// Build a B-tree v1 type 1 leaf node with given chunk infos.
|
|
fn build_chunk_btree_leaf(chunks: &[ChunkInfo], ndims: usize, offset_size: u8) -> Vec<u8> {
|
|
let _os = offset_size as usize;
|
|
let entries_used = chunks.len() as u16;
|
|
let mut buf = Vec::new();
|
|
|
|
// Header
|
|
buf.extend_from_slice(b"TREE");
|
|
buf.push(1); // node_type = 1 (raw data chunks)
|
|
buf.push(0); // node_level = 0 (leaf)
|
|
buf.extend_from_slice(&entries_used.to_le_bytes());
|
|
|
|
// Left/right sibling = undefined
|
|
let undef: u64 = if offset_size == 4 {
|
|
0xFFFFFFFF
|
|
} else {
|
|
0xFFFFFFFFFFFFFFFF
|
|
};
|
|
write_offset(&mut buf, undef, offset_size);
|
|
write_offset(&mut buf, undef, offset_size);
|
|
|
|
// Entries: key[i], child[i] pairs, then final key
|
|
for chunk in chunks {
|
|
// Key: chunk_size(4) + filter_mask(4) + ndims offsets
|
|
buf.extend_from_slice(&chunk.chunk_size.to_le_bytes());
|
|
buf.extend_from_slice(&chunk.filter_mask.to_le_bytes());
|
|
for d in 0..ndims {
|
|
let off = if d < chunk.offsets.len() {
|
|
chunk.offsets[d]
|
|
} else {
|
|
0
|
|
};
|
|
write_offset(&mut buf, off, offset_size);
|
|
}
|
|
// Child: address
|
|
write_offset(&mut buf, chunk.address, offset_size);
|
|
}
|
|
|
|
// Final key (dummy)
|
|
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);
|
|
}
|
|
|
|
buf
|
|
}
|
|
|
|
// --- ChunkInfo collection tests ---
|
|
|
|
#[test]
|
|
fn collect_two_chunks_from_leaf() {
|
|
let ndims = 2; // rank+1 for 1D dataset
|
|
let os: u8 = 8;
|
|
|
|
let chunks = vec![
|
|
ChunkInfo {
|
|
chunk_size: 80,
|
|
filter_mask: 0,
|
|
offsets: vec![0, 0],
|
|
address: 0x1000,
|
|
},
|
|
ChunkInfo {
|
|
chunk_size: 80,
|
|
filter_mask: 0,
|
|
offsets: vec![10, 0],
|
|
address: 0x2000,
|
|
},
|
|
];
|
|
|
|
let btree = build_chunk_btree_leaf(&chunks, ndims, os);
|
|
let mut file_data = vec![0u8; 0x3000];
|
|
file_data[..btree.len()].copy_from_slice(&btree);
|
|
|
|
let result = collect_chunk_info(&file_data, 0, ndims, os, os).unwrap();
|
|
assert_eq!(result.len(), 2);
|
|
assert_eq!(result[0].address, 0x1000);
|
|
assert_eq!(result[0].offsets, vec![0, 0]);
|
|
assert_eq!(result[0].chunk_size, 80);
|
|
assert_eq!(result[1].address, 0x2000);
|
|
assert_eq!(result[1].offsets, vec![10, 0]);
|
|
}
|
|
|
|
#[test]
|
|
fn collect_three_chunks() {
|
|
let ndims = 2;
|
|
let os: u8 = 8;
|
|
|
|
let chunks = vec![
|
|
ChunkInfo {
|
|
chunk_size: 40,
|
|
filter_mask: 0,
|
|
offsets: vec![0, 0],
|
|
address: 0x100,
|
|
},
|
|
ChunkInfo {
|
|
chunk_size: 40,
|
|
filter_mask: 0,
|
|
offsets: vec![5, 0],
|
|
address: 0x200,
|
|
},
|
|
ChunkInfo {
|
|
chunk_size: 40,
|
|
filter_mask: 0,
|
|
offsets: vec![10, 0],
|
|
address: 0x300,
|
|
},
|
|
];
|
|
|
|
let btree = build_chunk_btree_leaf(&chunks, ndims, os);
|
|
let mut file_data = vec![0u8; 0x1000];
|
|
file_data[..btree.len()].copy_from_slice(&btree);
|
|
|
|
let result = collect_chunk_info(&file_data, 0, ndims, os, os).unwrap();
|
|
assert_eq!(result.len(), 3);
|
|
assert_eq!(result[0].address, 0x100);
|
|
assert_eq!(result[1].address, 0x200);
|
|
assert_eq!(result[2].address, 0x300);
|
|
}
|
|
|
|
#[test]
|
|
fn collect_empty_btree() {
|
|
let ndims = 2;
|
|
let os: u8 = 8;
|
|
let btree = build_chunk_btree_leaf(&[], ndims, os);
|
|
let mut file_data = vec![0u8; 0x1000];
|
|
file_data[..btree.len()].copy_from_slice(&btree);
|
|
|
|
let result = collect_chunk_info(&file_data, 0, ndims, os, os).unwrap();
|
|
assert_eq!(result.len(), 0);
|
|
}
|
|
|
|
// --- Chunked read tests (synthetic) ---
|
|
|
|
use crate::dataspace::{Dataspace, DataspaceType};
|
|
use crate::datatype::{Datatype, DatatypeByteOrder};
|
|
|
|
fn make_f64_type() -> Datatype {
|
|
Datatype::FloatingPoint {
|
|
size: 8,
|
|
byte_order: DatatypeByteOrder::LittleEndian,
|
|
bit_offset: 0,
|
|
bit_precision: 64,
|
|
exponent_location: 52,
|
|
exponent_size: 11,
|
|
mantissa_location: 0,
|
|
mantissa_size: 52,
|
|
exponent_bias: 1023,
|
|
}
|
|
}
|
|
|
|
fn make_f32_type() -> Datatype {
|
|
Datatype::FloatingPoint {
|
|
size: 4,
|
|
byte_order: DatatypeByteOrder::LittleEndian,
|
|
bit_offset: 0,
|
|
bit_precision: 32,
|
|
exponent_location: 23,
|
|
exponent_size: 8,
|
|
mantissa_location: 0,
|
|
mantissa_size: 23,
|
|
exponent_bias: 127,
|
|
}
|
|
}
|
|
|
|
/// Build a synthetic file with a B-tree and chunk data for a 1D uncompressed dataset.
|
|
fn build_1d_chunked_file(
|
|
values: &[f64],
|
|
chunk_size_elems: usize,
|
|
) -> (Vec<u8>, DataLayout, Dataspace) {
|
|
let os: u8 = 8;
|
|
let elem_size = 8usize;
|
|
let ndims = 2; // rank(1) + 1
|
|
let total = values.len();
|
|
|
|
// Place chunk data starting at offset 0x2000
|
|
let mut file_data = vec![0u8; 0x10000];
|
|
let mut chunk_infos = Vec::new();
|
|
let mut data_offset = 0x2000usize;
|
|
|
|
let mut start = 0;
|
|
while start < total {
|
|
let end = (start + chunk_size_elems).min(total);
|
|
let chunk_bytes = chunk_size_elems * elem_size; // full chunk allocation
|
|
|
|
// Write chunk data (full chunk size, padding with zeros)
|
|
for (i, value) in values.iter().enumerate().take(end).skip(start) {
|
|
let byte_offset = data_offset + (i - start) * elem_size;
|
|
file_data[byte_offset..byte_offset + 8].copy_from_slice(&value.to_le_bytes());
|
|
}
|
|
|
|
chunk_infos.push(ChunkInfo {
|
|
chunk_size: chunk_bytes as u32,
|
|
filter_mask: 0,
|
|
offsets: vec![start as u64, 0],
|
|
address: data_offset as u64,
|
|
});
|
|
|
|
data_offset += chunk_bytes;
|
|
start += chunk_size_elems;
|
|
}
|
|
|
|
// Build B-tree at offset 0x100
|
|
let btree = build_chunk_btree_leaf(&chunk_infos, ndims, os);
|
|
let btree_addr = 0x100usize;
|
|
file_data[btree_addr..btree_addr + btree.len()].copy_from_slice(&btree);
|
|
|
|
let layout = DataLayout::Chunked {
|
|
chunk_dimensions: vec![chunk_size_elems as u32, elem_size as u32],
|
|
btree_address: Some(btree_addr as u64),
|
|
version: 3,
|
|
chunk_index_type: None,
|
|
single_chunk_filtered_size: None,
|
|
single_chunk_filter_mask: None,
|
|
};
|
|
|
|
let dataspace = Dataspace {
|
|
space_type: DataspaceType::Simple,
|
|
rank: 1,
|
|
dimensions: vec![total as u64],
|
|
max_dimensions: None,
|
|
};
|
|
|
|
(file_data, layout, dataspace)
|
|
}
|
|
|
|
#[test]
|
|
fn read_chunked_data_rejects_zero_dim_chunk_layout() {
|
|
// Found by fuzzing: chunk_dimensions.len() == 0 caused `ndims - 1` to
|
|
// underflow. A malformed/degenerate chunked layout must error cleanly.
|
|
let layout = DataLayout::Chunked {
|
|
chunk_dimensions: vec![],
|
|
btree_address: Some(0),
|
|
version: 3,
|
|
chunk_index_type: None,
|
|
single_chunk_filtered_size: None,
|
|
single_chunk_filter_mask: None,
|
|
};
|
|
let dataspace = Dataspace {
|
|
space_type: DataspaceType::Simple,
|
|
rank: 1,
|
|
dimensions: vec![10],
|
|
max_dimensions: None,
|
|
};
|
|
let datatype = make_f64_type();
|
|
let file_data = vec![0u8; 64];
|
|
let result = read_chunked_data(&file_data, &layout, &dataspace, &datatype, None, 8, 8);
|
|
assert!(
|
|
matches!(result, Err(FormatError::ChunkedReadError(_))),
|
|
"expected a clean ChunkedReadError, got {result:?}"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn copy_chunk_to_output_1d_rejects_overflowing_offset_without_panicking() {
|
|
// Found by fuzzing: `global_start * elem_size` overflowed for a
|
|
// crafted large chunk offset.
|
|
let chunk_data = vec![1u8; 16];
|
|
let mut output = vec![0u8; 16];
|
|
let chunk_offsets = [usize::MAX - 1];
|
|
let chunk_dims = [1usize];
|
|
let ds_dims = [usize::MAX];
|
|
let ds_strides = [1usize];
|
|
let chunk_strides = [1usize];
|
|
copy_chunk_to_output(
|
|
&chunk_data,
|
|
&mut output,
|
|
&chunk_offsets,
|
|
&chunk_dims,
|
|
&ds_dims,
|
|
&ds_strides,
|
|
&chunk_strides,
|
|
8,
|
|
1,
|
|
);
|
|
// No panic; the out-of-range write was skipped, output left untouched.
|
|
assert_eq!(output, vec![0u8; 16]);
|
|
}
|
|
|
|
#[test]
|
|
fn copy_chunk_to_output_nd_rejects_overflowing_offset_without_panicking() {
|
|
let chunk_data = vec![1u8; 16];
|
|
let mut output = vec![0u8; 16];
|
|
let chunk_offsets = [usize::MAX - 1, 0];
|
|
let chunk_dims = [1usize, 1usize];
|
|
let ds_dims = [usize::MAX, usize::MAX];
|
|
let ds_strides = [1usize, 1usize];
|
|
let chunk_strides = [1usize, 1usize];
|
|
copy_chunk_to_output(
|
|
&chunk_data,
|
|
&mut output,
|
|
&chunk_offsets,
|
|
&chunk_dims,
|
|
&ds_dims,
|
|
&ds_strides,
|
|
&chunk_strides,
|
|
8,
|
|
2,
|
|
);
|
|
assert_eq!(output, vec![0u8; 16]);
|
|
}
|
|
|
|
#[test]
|
|
fn read_1d_two_chunks_no_compression() {
|
|
let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
|
|
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
|
let datatype = make_f64_type();
|
|
|
|
let raw =
|
|
read_chunked_data(&file_data, &layout, &dataspace, &datatype, None, 8, 8).unwrap();
|
|
assert_eq!(raw.len(), 20 * 8);
|
|
|
|
// Verify values
|
|
for i in 0..20 {
|
|
let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
|
|
assert_eq!(val, i as f64);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn read_1d_three_chunks_partial_last() {
|
|
// 25 elements, chunk size 10 => 3 chunks, last has only 5 valid
|
|
let values: Vec<f64> = (0..25).map(|i| i as f64).collect();
|
|
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
|
let datatype = make_f64_type();
|
|
|
|
let raw =
|
|
read_chunked_data(&file_data, &layout, &dataspace, &datatype, None, 8, 8).unwrap();
|
|
assert_eq!(raw.len(), 25 * 8);
|
|
|
|
for i in 0..25 {
|
|
let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
|
|
assert_eq!(val, i as f64, "mismatch at index {i}");
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "deflate")]
|
|
#[test]
|
|
fn read_1d_two_chunks_with_deflate() {
|
|
use crate::filter_pipeline::{FILTER_DEFLATE, FilterDescription, FilterPipeline};
|
|
use crate::filters::compress_chunk;
|
|
|
|
let os: u8 = 8;
|
|
let elem_size = 8usize;
|
|
let ndims = 2;
|
|
let chunk_elems = 10usize;
|
|
let total = 20usize;
|
|
|
|
let pipeline = FilterPipeline {
|
|
version: 2,
|
|
filters: vec![FilterDescription {
|
|
filter_id: FILTER_DEFLATE,
|
|
name: None,
|
|
flags: 0,
|
|
client_data: vec![6],
|
|
}],
|
|
};
|
|
|
|
let values: Vec<f64> = (0..total).map(|i| i as f64).collect();
|
|
let mut file_data = vec![0u8; 0x10000];
|
|
let mut chunk_infos = Vec::new();
|
|
let mut data_offset = 0x2000usize;
|
|
|
|
for chunk_idx in 0..2 {
|
|
let start = chunk_idx * chunk_elems;
|
|
let mut chunk_bytes = Vec::new();
|
|
for value in values.iter().skip(start).take(chunk_elems) {
|
|
chunk_bytes.extend_from_slice(&value.to_le_bytes());
|
|
}
|
|
let compressed = compress_chunk(&chunk_bytes, &pipeline, elem_size as u32).unwrap();
|
|
|
|
file_data[data_offset..data_offset + compressed.len()].copy_from_slice(&compressed);
|
|
|
|
chunk_infos.push(ChunkInfo {
|
|
chunk_size: compressed.len() as u32,
|
|
filter_mask: 0,
|
|
offsets: vec![start as u64, 0],
|
|
address: data_offset as u64,
|
|
});
|
|
|
|
data_offset += compressed.len() + 16; // some padding
|
|
}
|
|
|
|
let btree = build_chunk_btree_leaf(&chunk_infos, ndims, os);
|
|
let btree_addr = 0x100usize;
|
|
file_data[btree_addr..btree_addr + btree.len()].copy_from_slice(&btree);
|
|
|
|
let layout = DataLayout::Chunked {
|
|
chunk_dimensions: vec![chunk_elems as u32, elem_size as u32],
|
|
btree_address: Some(btree_addr as u64),
|
|
version: 3,
|
|
chunk_index_type: None,
|
|
single_chunk_filtered_size: None,
|
|
single_chunk_filter_mask: None,
|
|
};
|
|
let dataspace = Dataspace {
|
|
space_type: DataspaceType::Simple,
|
|
rank: 1,
|
|
dimensions: vec![total as u64],
|
|
max_dimensions: None,
|
|
};
|
|
let datatype = make_f64_type();
|
|
|
|
let raw = read_chunked_data(
|
|
&file_data,
|
|
&layout,
|
|
&dataspace,
|
|
&datatype,
|
|
Some(&pipeline),
|
|
8,
|
|
8,
|
|
)
|
|
.unwrap();
|
|
|
|
for i in 0..total {
|
|
let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
|
|
assert_eq!(val, i as f64, "mismatch at index {i}");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn read_2d_four_chunks() {
|
|
// 4x6 dataset with chunk size 2x3 => 4 chunks
|
|
let os: u8 = 8;
|
|
let elem_size = 4usize; // f32
|
|
let ndims = 3; // rank(2) + 1
|
|
let ds_dims = [4usize, 6];
|
|
let chunk_dims = [2usize, 3];
|
|
|
|
let values: Vec<f32> = (0..24).map(|i| i as f32).collect();
|
|
let mut file_data = vec![0u8; 0x10000];
|
|
let mut chunk_infos = Vec::new();
|
|
let mut data_offset = 0x2000usize;
|
|
|
|
// Generate chunks: (0,0), (0,3), (2,0), (2,3)
|
|
for row_start in (0..ds_dims[0]).step_by(chunk_dims[0]) {
|
|
for col_start in (0..ds_dims[1]).step_by(chunk_dims[1]) {
|
|
let mut chunk_bytes = Vec::new();
|
|
for r in 0..chunk_dims[0] {
|
|
for c in 0..chunk_dims[1] {
|
|
let gr = row_start + r;
|
|
let gc = col_start + c;
|
|
let val = if gr < ds_dims[0] && gc < ds_dims[1] {
|
|
values[gr * ds_dims[1] + gc]
|
|
} else {
|
|
0.0
|
|
};
|
|
chunk_bytes.extend_from_slice(&val.to_le_bytes());
|
|
}
|
|
}
|
|
|
|
let chunk_size = chunk_bytes.len();
|
|
file_data[data_offset..data_offset + chunk_size].copy_from_slice(&chunk_bytes);
|
|
|
|
chunk_infos.push(ChunkInfo {
|
|
chunk_size: chunk_size as u32,
|
|
filter_mask: 0,
|
|
offsets: vec![row_start as u64, col_start as u64, 0],
|
|
address: data_offset as u64,
|
|
});
|
|
|
|
data_offset += chunk_size + 8;
|
|
}
|
|
}
|
|
|
|
let btree = build_chunk_btree_leaf(&chunk_infos, ndims, os);
|
|
let btree_addr = 0x100usize;
|
|
file_data[btree_addr..btree_addr + btree.len()].copy_from_slice(&btree);
|
|
|
|
let layout = DataLayout::Chunked {
|
|
chunk_dimensions: vec![chunk_dims[0] as u32, chunk_dims[1] as u32, elem_size as u32],
|
|
btree_address: Some(btree_addr as u64),
|
|
version: 3,
|
|
chunk_index_type: None,
|
|
single_chunk_filtered_size: None,
|
|
single_chunk_filter_mask: None,
|
|
};
|
|
let dataspace = Dataspace {
|
|
space_type: DataspaceType::Simple,
|
|
rank: 2,
|
|
dimensions: vec![ds_dims[0] as u64, ds_dims[1] as u64],
|
|
max_dimensions: None,
|
|
};
|
|
let datatype = make_f32_type();
|
|
|
|
let raw =
|
|
read_chunked_data(&file_data, &layout, &dataspace, &datatype, None, 8, 8).unwrap();
|
|
assert_eq!(raw.len(), 24 * 4);
|
|
|
|
for i in 0..24 {
|
|
let val = f32::from_le_bytes(raw[i * 4..(i + 1) * 4].try_into().unwrap());
|
|
assert_eq!(val, i as f32, "mismatch at element {i}");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn wrong_node_type_error() {
|
|
// Build a type-0 B-tree and try to collect chunk info
|
|
let mut buf = Vec::new();
|
|
buf.extend_from_slice(b"TREE");
|
|
buf.push(0); // type 0, not 1
|
|
buf.push(0);
|
|
buf.extend_from_slice(&0u16.to_le_bytes());
|
|
buf.extend_from_slice(&[0xFF; 16]); // siblings
|
|
// final key
|
|
buf.extend_from_slice(&[0u8; 24]);
|
|
|
|
let mut file_data = vec![0u8; 512];
|
|
file_data[..buf.len()].copy_from_slice(&buf);
|
|
|
|
let err = collect_chunk_info(&file_data, 0, 2, 8, 8).unwrap_err();
|
|
assert_eq!(err, FormatError::InvalidBTreeNodeType(0));
|
|
}
|
|
|
|
#[test]
|
|
fn collect_chunk_info_rejects_near_usize_max_offset() {
|
|
let file_data = vec![0u8; 64];
|
|
let result = collect_chunk_info(&file_data, u64::MAX - 4, 2, 8, 8);
|
|
assert!(
|
|
matches!(result, Err(FormatError::UnexpectedEof { .. })),
|
|
"expected a clean UnexpectedEof, got {result:?}"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn collect_chunk_info_rejects_self_referencing_internal_node() {
|
|
// A type-1 internal node (level 1) whose single child address points
|
|
// back to itself: an infinite-recursion / cyclic B-tree attack.
|
|
let ndims = 2;
|
|
let os: u8 = 8;
|
|
let mut buf = Vec::new();
|
|
buf.extend_from_slice(b"TREE");
|
|
buf.push(1); // node_type = 1 (raw data chunks)
|
|
buf.push(1); // node_level = 1 (internal)
|
|
buf.extend_from_slice(&1u16.to_le_bytes()); // entries_used = 1
|
|
write_offset(&mut buf, u64::MAX, os); // left sibling undefined
|
|
write_offset(&mut buf, u64::MAX, os); // right sibling undefined
|
|
// key[0]: chunk_size(4) + filter_mask(4) + ndims offsets
|
|
buf.extend_from_slice(&0u32.to_le_bytes());
|
|
buf.extend_from_slice(&0u32.to_le_bytes());
|
|
for _ in 0..ndims {
|
|
write_offset(&mut buf, 0, os);
|
|
}
|
|
// child[0]: points back to offset 0 (this same node) — cyclic.
|
|
write_offset(&mut buf, 0, os);
|
|
// final key
|
|
buf.extend_from_slice(&0u32.to_le_bytes());
|
|
buf.extend_from_slice(&0u32.to_le_bytes());
|
|
for _ in 0..ndims {
|
|
write_offset(&mut buf, u64::MAX, os);
|
|
}
|
|
|
|
let mut file_data = vec![0u8; 256];
|
|
file_data[..buf.len()].copy_from_slice(&buf);
|
|
|
|
let result = collect_chunk_info(&file_data, 0, ndims, os, os);
|
|
assert!(
|
|
matches!(result, Err(FormatError::NestingDepthExceeded)),
|
|
"expected a clean NestingDepthExceeded, got {result:?}"
|
|
);
|
|
}
|
|
|
|
// --- Implicit chunk generation tests ---
|
|
|
|
#[test]
|
|
fn implicit_chunks_1d_five_chunks() {
|
|
let chunks = generate_implicit_chunks(
|
|
0x1000,
|
|
&[100],
|
|
&[20],
|
|
8, // f64
|
|
);
|
|
assert_eq!(chunks.len(), 5);
|
|
let chunk_byte_size = 20 * 8;
|
|
for (i, c) in chunks.iter().enumerate() {
|
|
assert_eq!(c.address, 0x1000 + i as u64 * chunk_byte_size as u64);
|
|
assert_eq!(c.offsets, vec![i as u64 * 20]);
|
|
assert_eq!(c.filter_mask, 0);
|
|
assert_eq!(c.chunk_size, chunk_byte_size as u32);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn implicit_chunks_2d() {
|
|
// 10x6 dataset, 4x3 chunks => ceil(10/4)=3, ceil(6/3)=2 => 6 chunks
|
|
let chunks = generate_implicit_chunks(
|
|
0x2000,
|
|
&[10, 6],
|
|
&[4, 3],
|
|
4, // f32
|
|
);
|
|
assert_eq!(chunks.len(), 6);
|
|
let chunk_byte_size = 4 * 3 * 4;
|
|
// Row-major: (0,0), (0,3), (4,0), (4,3), (8,0), (8,3)
|
|
assert_eq!(chunks[0].offsets, vec![0, 0]);
|
|
assert_eq!(chunks[1].offsets, vec![0, 3]);
|
|
assert_eq!(chunks[2].offsets, vec![4, 0]);
|
|
assert_eq!(chunks[3].offsets, vec![4, 3]);
|
|
assert_eq!(chunks[4].offsets, vec![8, 0]);
|
|
assert_eq!(chunks[5].offsets, vec![8, 3]);
|
|
for (i, c) in chunks.iter().enumerate() {
|
|
assert_eq!(c.address, 0x2000 + i as u64 * chunk_byte_size as u64);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn implicit_chunks_partial_last() {
|
|
// 25 elements, chunk size 10 => 3 chunks (last partial)
|
|
let chunks = generate_implicit_chunks(0x0, &[25], &[10], 8);
|
|
assert_eq!(chunks.len(), 3);
|
|
assert_eq!(chunks[0].offsets, vec![0]);
|
|
assert_eq!(chunks[1].offsets, vec![10]);
|
|
assert_eq!(chunks[2].offsets, vec![20]);
|
|
}
|
|
|
|
// --- V4 single chunk synthetic test ---
|
|
|
|
#[test]
|
|
fn read_v4_single_chunk_synthetic() {
|
|
// Build a synthetic v4 single chunk dataset (no filters)
|
|
let values: Vec<f64> = vec![10.0, 20.0, 30.0];
|
|
let elem_size = 8usize;
|
|
let chunk_elems = 3usize;
|
|
|
|
let mut file_data = vec![0u8; 0x2000];
|
|
let data_addr = 0x1000usize;
|
|
for (i, &v) in values.iter().enumerate() {
|
|
file_data[data_addr + i * elem_size..data_addr + (i + 1) * elem_size]
|
|
.copy_from_slice(&v.to_le_bytes());
|
|
}
|
|
|
|
let layout = DataLayout::Chunked {
|
|
chunk_dimensions: vec![chunk_elems as u32, elem_size as u32],
|
|
btree_address: Some(data_addr as u64),
|
|
version: 4,
|
|
chunk_index_type: Some(1),
|
|
single_chunk_filtered_size: None,
|
|
single_chunk_filter_mask: None,
|
|
};
|
|
let dataspace = Dataspace {
|
|
space_type: DataspaceType::Simple,
|
|
rank: 1,
|
|
dimensions: vec![3],
|
|
max_dimensions: None,
|
|
};
|
|
let datatype = make_f64_type();
|
|
|
|
let raw =
|
|
read_chunked_data(&file_data, &layout, &dataspace, &datatype, None, 8, 8).unwrap();
|
|
assert_eq!(raw.len(), 24);
|
|
for i in 0..3 {
|
|
let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
|
|
assert_eq!(val, values[i]);
|
|
}
|
|
}
|
|
|
|
// --- Cached read tests ---
|
|
|
|
use crate::chunk_cache::ChunkCache;
|
|
|
|
#[test]
|
|
fn cached_read_populates_index_and_returns_correct_data() {
|
|
let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
|
|
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
|
let datatype = make_f64_type();
|
|
let cache = ChunkCache::new();
|
|
|
|
assert!(!cache.has_index());
|
|
let raw = read_chunked_data_cached(
|
|
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
|
|
)
|
|
.unwrap();
|
|
assert!(cache.has_index());
|
|
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());
|
|
assert_eq!(val, i as f64);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn cached_read_second_call_reuses_the_index() {
|
|
let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
|
|
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
|
let datatype = make_f64_type();
|
|
let cache = ChunkCache::new();
|
|
|
|
// First read — populates the chunk index. These chunks are stored
|
|
// unfiltered, so they are copied straight from the file bytes and the
|
|
// decompressed-chunk cache is (deliberately) not involved.
|
|
let raw1 = read_chunked_data_cached(
|
|
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
|
|
)
|
|
.unwrap();
|
|
assert!(cache.has_index());
|
|
assert_eq!(cache.cached_chunk_count(), 0);
|
|
|
|
// Second read — reuses the cached index
|
|
let raw2 = read_chunked_data_cached(
|
|
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
|
|
)
|
|
.unwrap();
|
|
assert_eq!(raw1, raw2);
|
|
}
|
|
|
|
#[test]
|
|
fn cached_read_with_partial_last_chunk() {
|
|
let values: Vec<f64> = (0..25).map(|i| i as f64).collect();
|
|
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
|
let datatype = make_f64_type();
|
|
let cache = ChunkCache::new();
|
|
|
|
let raw = read_chunked_data_cached(
|
|
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
|
|
)
|
|
.unwrap();
|
|
assert_eq!(raw.len(), 25 * 8);
|
|
for i in 0..25 {
|
|
let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
|
|
assert_eq!(val, i as f64, "mismatch at index {i}");
|
|
}
|
|
}
|
|
|
|
// --- Sweep-aware read tests ---
|
|
|
|
#[test]
|
|
fn sweep_read_returns_correct_data() {
|
|
let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
|
|
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
|
let datatype = make_f64_type();
|
|
let cache = ChunkCache::new();
|
|
let mut sweep = SweepContext::with_defaults();
|
|
|
|
let raw = read_chunked_data_sweep(
|
|
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache, &mut sweep,
|
|
)
|
|
.unwrap();
|
|
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());
|
|
assert_eq!(val, i as f64);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn sweep_read_populates_sweep_context() {
|
|
let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
|
|
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
|
let datatype = make_f64_type();
|
|
let cache = ChunkCache::new();
|
|
let mut sweep = SweepContext::with_defaults();
|
|
|
|
read_chunked_data_sweep(
|
|
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache, &mut sweep,
|
|
)
|
|
.unwrap();
|
|
|
|
// After reading 2 chunks (offsets [0] and [10]), history should be populated
|
|
assert!(!sweep.history.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn sweep_context_unit_test() {
|
|
let mut ctx = SweepContext::with_defaults();
|
|
ctx.record(vec![0, 0], 2);
|
|
ctx.record(vec![0, 10], 2);
|
|
ctx.record(vec![0, 20], 2);
|
|
assert_eq!(ctx.direction, "row_major");
|
|
assert!(!ctx.predicted_next.is_empty());
|
|
assert_eq!(ctx.predicted_next[0], vec![0, 30]);
|
|
}
|
|
|
|
#[test]
|
|
fn sweep_context_random() {
|
|
let mut ctx = SweepContext::with_defaults();
|
|
ctx.record(vec![0, 0], 2);
|
|
ctx.record(vec![30, 20], 2);
|
|
ctx.record(vec![10, 0], 2);
|
|
assert_eq!(ctx.direction, "random");
|
|
assert!(ctx.predicted_next.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn sweep_read_access_stats() {
|
|
let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
|
|
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
|
|
let datatype = make_f64_type();
|
|
let cache = ChunkCache::new();
|
|
let mut sweep = SweepContext::with_defaults();
|
|
|
|
read_chunked_data_sweep(
|
|
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache, &mut sweep,
|
|
)
|
|
.unwrap();
|
|
|
|
let stats = cache.access_stats();
|
|
// We accessed 2 chunks; the second should be sequential to the first
|
|
assert!(stats.sequential_count > 0 || stats.random_count > 0);
|
|
}
|
|
}
|