osobh
591aa71d12
# Conflicts: # crates/clawhdf5/tests/h5py_chunked_read_tests.rs # docs/known-issues.md
2151 lines
75 KiB
Rust
2151 lines
75 KiB
Rust
//! Chunked dataset writing: chunk splitting, compression, index building.
|
||
|
||
#[cfg(not(feature = "std"))]
|
||
extern crate alloc;
|
||
|
||
#[cfg(not(feature = "std"))]
|
||
use alloc::{format, vec, vec::Vec};
|
||
|
||
use crate::checksum::jenkins_lookup3;
|
||
use crate::chunk_cache::{CACHE_LINE_SIZE, align_to_cache_line};
|
||
use crate::chunk_grid::ChunkGrid;
|
||
use crate::ea_writer;
|
||
use crate::error::FormatError;
|
||
use crate::filter_pipeline::{
|
||
FILTER_BITSHUFFLE, FILTER_BLOSC, FILTER_BZIP2, FILTER_DEFLATE, FILTER_FLETCHER32, FILTER_LZ4,
|
||
FILTER_LZF, FILTER_PCODEC, FILTER_PCODEC_NAME, FILTER_SHUFFLE, FILTER_ZSTD, FilterDescription,
|
||
FilterPipeline,
|
||
};
|
||
use crate::filters::compress_chunk;
|
||
/// Round a file offset up to the next cache-line boundary.
|
||
///
|
||
/// This ensures chunk data starts at an address that is a multiple of the
|
||
/// architecture's cache line size, enabling aligned loads in SIMD paths.
|
||
#[inline]
|
||
pub fn align_chunk_offset(offset: u64) -> u64 {
|
||
let align = CACHE_LINE_SIZE as u64;
|
||
(offset + align - 1) & !(align - 1)
|
||
}
|
||
|
||
/// Options for chunked dataset creation.
|
||
#[derive(Debug, Clone, Default)]
|
||
pub struct ChunkOptions {
|
||
/// Chunk dimensions (one per dataset dimension).
|
||
pub chunk_dims: Option<Vec<u64>>,
|
||
/// Deflate compression level (0-9), None = no deflate.
|
||
pub deflate_level: Option<u32>,
|
||
/// Whether to apply shuffle filter before compression.
|
||
/// If `false` AND compression is enabled AND `no_shuffle` is `false`,
|
||
/// shuffle is auto-applied (matches h5py default behavior).
|
||
pub shuffle: bool,
|
||
/// Disable the automatic shuffle pre-filter. Set via `without_shuffle()`.
|
||
pub no_shuffle: bool,
|
||
/// Whether to apply fletcher32 checksum.
|
||
pub fletcher32: bool,
|
||
/// Whether to use LZ4 compression (filter ID 32004).
|
||
pub lz4: bool,
|
||
/// Zstandard compression level (1-22), None = no zstd. Filter ID 32015.
|
||
pub zstd_level: Option<u32>,
|
||
/// Pcodec lossless numerical compression. Private, unregistered filter
|
||
/// ID [`FILTER_PCODEC`] (480): only clawhdf5 can read it.
|
||
pub pcodec: bool,
|
||
/// A plugin compression filter (LZF, ...). Takes priority over the
|
||
/// codecs above. Each needs its cargo feature to be written.
|
||
pub plugin: Option<PluginFilter>,
|
||
}
|
||
|
||
/// A compression filter from the common HDF5 plugin set, written in the
|
||
/// format the libhdf5 plugin (h5py / hdf5plugin) reads.
|
||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||
#[non_exhaustive]
|
||
pub enum PluginFilter {
|
||
/// LZF (filter 32000), h5py's built-in `compression="lzf"`. Needs the
|
||
/// `lzf` feature.
|
||
Lzf,
|
||
/// Bitshuffle (filter 32008): a bit transpose of each block of
|
||
/// `block_size` elements (0 = bitshuffle's default, else a multiple of
|
||
/// 8), optionally compressed. Needs the `bitshuffle` feature.
|
||
Bitshuffle {
|
||
/// Block size in elements; 0 for the default.
|
||
block_size: u32,
|
||
/// Compression after the transpose.
|
||
compression: BitshuffleCompression,
|
||
},
|
||
/// bzip2 (filter 307) at block size `level` (1-9). Needs the `bzip2`
|
||
/// feature.
|
||
Bzip2 {
|
||
/// Block size 1-9 (9 = hdf5plugin's default).
|
||
level: u32,
|
||
},
|
||
/// Blosc 1 (filter 32001): `codec` at `level` (0-9; 0 stores), after
|
||
/// `shuffle`. Needs the `blosc` feature.
|
||
Blosc {
|
||
/// The codec inside the Blosc frame.
|
||
codec: BloscCodec,
|
||
/// Compression level 0-9 (0 stores the data uncompressed).
|
||
level: u32,
|
||
/// The shuffle Blosc applies first.
|
||
shuffle: BloscShuffle,
|
||
},
|
||
}
|
||
|
||
/// The codec inside a Blosc frame that clawhdf5 can write. (It reads
|
||
/// BloscLZ too, but cannot write it.)
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
pub enum BloscCodec {
|
||
/// LZ4.
|
||
Lz4,
|
||
/// Snappy.
|
||
Snappy,
|
||
/// Zlib, at the Blosc level.
|
||
Zlib,
|
||
/// Zstandard (clawhdf5's pure-Rust encoder has one level, about zstd 1).
|
||
Zstd,
|
||
}
|
||
|
||
/// The shuffle Blosc applies before compressing.
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
pub enum BloscShuffle {
|
||
/// None.
|
||
None,
|
||
/// Byte shuffle (Blosc's default).
|
||
Byte,
|
||
/// Bit shuffle.
|
||
Bit,
|
||
}
|
||
|
||
/// What bitshuffle compresses its blocks with.
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
pub enum BitshuffleCompression {
|
||
/// Transpose only.
|
||
None,
|
||
/// LZ4 (bitshuffle's `cname="lz4"`, the common choice).
|
||
Lz4,
|
||
/// Zstandard. clawhdf5's pure-Rust encoder has a single level (about
|
||
/// zstd's level 1); `level` is recorded in the file for other writers.
|
||
Zstd {
|
||
/// Level recorded in `cd_values[5]`.
|
||
level: u32,
|
||
},
|
||
}
|
||
|
||
impl PluginFilter {
|
||
/// Whether the filter reorders bytes itself, so the automatic shuffle
|
||
/// pre-filter would only get in its way.
|
||
fn shuffles_itself(&self) -> bool {
|
||
match self {
|
||
PluginFilter::Lzf => false,
|
||
PluginFilter::Bitshuffle { .. } => true,
|
||
PluginFilter::Bzip2 { .. } => false,
|
||
PluginFilter::Blosc { .. } => true,
|
||
}
|
||
}
|
||
|
||
/// The pipeline entry for this filter. `chunk_bytes` is one chunk's
|
||
/// uncompressed size (0 if unknown).
|
||
fn description(&self, element_size: u32, chunk_bytes: u32) -> FilterDescription {
|
||
match self {
|
||
// h5py's lzf_set_local: filter version, liblzf version, chunk
|
||
// size in bytes. Optional, as h5py flags it: a chunk the filter
|
||
// cannot shrink may then be stored unfiltered.
|
||
PluginFilter::Lzf => FilterDescription {
|
||
filter_id: FILTER_LZF,
|
||
name: Some("lzf".into()),
|
||
flags: 1,
|
||
client_data: vec![4, 0x0105, chunk_bytes],
|
||
},
|
||
// bshuf_h5_set_local: version 0.4, element size, block size,
|
||
// compression (0 none, 2 LZ4, 3 Zstandard), Zstandard level.
|
||
// hdf5-blosc's blosc_set_local: filter revision 2, Blosc format
|
||
// 2, type size, chunk size, then level, shuffle, compressor.
|
||
PluginFilter::Blosc {
|
||
codec,
|
||
level,
|
||
shuffle,
|
||
} => FilterDescription {
|
||
filter_id: FILTER_BLOSC,
|
||
name: Some("blosc".into()),
|
||
flags: 1,
|
||
client_data: vec![
|
||
2,
|
||
2,
|
||
element_size,
|
||
chunk_bytes,
|
||
(*level).min(9),
|
||
match shuffle {
|
||
BloscShuffle::None => 0,
|
||
BloscShuffle::Byte => 1,
|
||
BloscShuffle::Bit => 2,
|
||
},
|
||
match codec {
|
||
BloscCodec::Lz4 => 1,
|
||
BloscCodec::Snappy => 3,
|
||
BloscCodec::Zlib => 4,
|
||
BloscCodec::Zstd => 5,
|
||
},
|
||
],
|
||
},
|
||
PluginFilter::Bzip2 { level } => FilterDescription {
|
||
filter_id: FILTER_BZIP2,
|
||
name: Some("bzip2".into()),
|
||
flags: 1,
|
||
client_data: vec![(*level).clamp(1, 9)],
|
||
},
|
||
PluginFilter::Bitshuffle {
|
||
block_size,
|
||
compression,
|
||
} => {
|
||
let mut cd = vec![0, 4, element_size, *block_size];
|
||
match compression {
|
||
BitshuffleCompression::None => cd.push(0),
|
||
BitshuffleCompression::Lz4 => cd.push(2),
|
||
BitshuffleCompression::Zstd { level } => cd.extend([3, *level]),
|
||
}
|
||
FilterDescription {
|
||
filter_id: FILTER_BITSHUFFLE,
|
||
name: Some("bitshuffle; see https://github.com/kiyo-masui/bitshuffle".into()),
|
||
flags: 1,
|
||
client_data: cd,
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Largest chunk the automatic choice produces, in bytes.
|
||
const AUTO_CHUNK_TARGET_BYTES: u64 = 1 << 20;
|
||
|
||
/// Extent assumed for a dimension that is currently empty (an unlimited
|
||
/// dimension not yet written to) — the same stand-in h5py uses.
|
||
const AUTO_CHUNK_EMPTY_DIM: u64 = 1024;
|
||
|
||
/// Choose chunk dimensions for a dataset nobody specified them for.
|
||
///
|
||
/// Asking for compression (or any filter) without chunk dimensions used to
|
||
/// make the whole dataset one chunk. That defeats the point of chunking: any
|
||
/// read — even a single row — must decompress everything, and a large dataset
|
||
/// cannot be decompressed in parallel. Datasets up to the target size stay a
|
||
/// single chunk, exactly as before; larger ones are split by halving the
|
||
/// dimensions in turn (so chunks keep roughly the dataset's proportions, the
|
||
/// approach h5py takes) until a chunk fits the target.
|
||
pub fn auto_chunk_dims(shape: &[u64], elem_size: usize) -> Vec<u64> {
|
||
let mut dims: Vec<u64> = shape
|
||
.iter()
|
||
.map(|&d| if d == 0 { AUTO_CHUNK_EMPTY_DIM } else { d })
|
||
.collect();
|
||
let elem = elem_size.max(1) as u64;
|
||
let bytes = |dims: &[u64]| dims.iter().fold(elem, |acc, &d| acc.saturating_mul(d));
|
||
let mut axis = 0;
|
||
while bytes(&dims) > AUTO_CHUNK_TARGET_BYTES && dims.iter().any(|&d| d > 1) {
|
||
let i = axis % dims.len();
|
||
dims[i] = dims[i].div_ceil(2);
|
||
axis += 1;
|
||
}
|
||
dims
|
||
}
|
||
|
||
impl ChunkOptions {
|
||
/// Whether any chunking option is enabled.
|
||
pub fn is_chunked(&self) -> bool {
|
||
self.chunk_dims.is_some()
|
||
|| self.deflate_level.is_some()
|
||
|| self.shuffle
|
||
|| self.fletcher32
|
||
|| self.lz4
|
||
|| self.zstd_level.is_some()
|
||
|| self.pcodec
|
||
|| self.plugin.is_some()
|
||
}
|
||
|
||
/// Build a FilterPipeline from the options.
|
||
pub fn build_pipeline(&self, element_size: u32) -> Option<FilterPipeline> {
|
||
self.build_pipeline_for_chunk(element_size, 0)
|
||
}
|
||
|
||
/// Build a FilterPipeline for chunks of `chunk_bytes` uncompressed bytes
|
||
/// (0 if unknown). Some plugin filters record the chunk size in their
|
||
/// client data.
|
||
pub fn build_pipeline_for_chunk(
|
||
&self,
|
||
element_size: u32,
|
||
chunk_bytes: u32,
|
||
) -> Option<FilterPipeline> {
|
||
let mut filters = Vec::new();
|
||
|
||
let plugin_shuffles = self
|
||
.plugin
|
||
.as_ref()
|
||
.is_some_and(PluginFilter::shuffles_itself);
|
||
let has_compression = self.deflate_level.is_some()
|
||
|| self.zstd_level.is_some()
|
||
|| self.lz4
|
||
|| self.pcodec
|
||
|| (self.plugin.is_some() && !plugin_shuffles);
|
||
|
||
// Shuffle before compression. Applied if explicitly requested OR if compression
|
||
// is active and the caller hasn't disabled it — matches h5py default behavior
|
||
// and implements TDT byte-grouping (arXiv:2506.18062) for free.
|
||
if self.shuffle || (has_compression && !self.no_shuffle) {
|
||
filters.push(FilterDescription {
|
||
filter_id: FILTER_SHUFFLE,
|
||
name: None,
|
||
flags: 0,
|
||
client_data: vec![element_size],
|
||
});
|
||
}
|
||
|
||
// Compression filters (mutually exclusive, priority: plugin > pcodec >
|
||
// zstd > lz4 > deflate)
|
||
if let Some(plugin) = &self.plugin {
|
||
filters.push(plugin.description(element_size, chunk_bytes));
|
||
} else if self.pcodec {
|
||
filters.push(FilterDescription {
|
||
filter_id: FILTER_PCODEC,
|
||
name: Some(FILTER_PCODEC_NAME.into()),
|
||
flags: 0,
|
||
client_data: vec![element_size],
|
||
});
|
||
} else if let Some(level) = self.zstd_level {
|
||
filters.push(FilterDescription {
|
||
filter_id: FILTER_ZSTD,
|
||
name: Some("zstd".into()),
|
||
flags: 0,
|
||
client_data: vec![level],
|
||
});
|
||
} else if self.lz4 {
|
||
filters.push(FilterDescription {
|
||
filter_id: FILTER_LZ4,
|
||
name: Some("lz4".into()),
|
||
flags: 0,
|
||
client_data: vec![],
|
||
});
|
||
} else if let Some(level) = self.deflate_level {
|
||
filters.push(FilterDescription {
|
||
filter_id: FILTER_DEFLATE,
|
||
name: None,
|
||
flags: 0,
|
||
client_data: vec![level],
|
||
});
|
||
}
|
||
|
||
if self.fletcher32 {
|
||
filters.push(FilterDescription {
|
||
filter_id: FILTER_FLETCHER32,
|
||
name: None,
|
||
flags: 0,
|
||
client_data: vec![],
|
||
});
|
||
}
|
||
|
||
// Note: h5py sets flags=0x0001 (optional) on filters, but this is not required
|
||
// for read compatibility.
|
||
|
||
if filters.is_empty() {
|
||
None
|
||
} else {
|
||
Some(FilterPipeline {
|
||
version: 2,
|
||
filters,
|
||
})
|
||
}
|
||
}
|
||
|
||
/// Determine chunk dimensions, using user-specified or auto-computing.
|
||
pub fn resolve_chunk_dims(&self, shape: &[u64]) -> Vec<u64> {
|
||
// Without the element size, assume 8 bytes (the widest common scalar);
|
||
// the writer uses `resolve_chunk_dims_for`.
|
||
self.resolve_chunk_dims_for(shape, 8)
|
||
}
|
||
|
||
/// Chunk dimensions for a dataset of `shape` whose elements are `elem_size`
|
||
/// bytes: the caller's if given, otherwise chosen automatically.
|
||
pub fn resolve_chunk_dims_for(&self, shape: &[u64], elem_size: usize) -> Vec<u64> {
|
||
match self.chunk_dims {
|
||
Some(ref dims) => dims.clone(),
|
||
None => auto_chunk_dims(shape, elem_size),
|
||
}
|
||
}
|
||
}
|
||
|
||
/// A chunk that has been written to the file buffer.
|
||
#[derive(Debug, Clone)]
|
||
pub struct WrittenChunk {
|
||
/// Address within the file where chunk data starts.
|
||
pub address: u64,
|
||
/// Size of the (possibly compressed) chunk data in bytes.
|
||
pub compressed_size: u64,
|
||
/// Original uncompressed size in bytes.
|
||
pub raw_size: u64,
|
||
/// Filter mask (0 = all filters applied).
|
||
pub filter_mask: u32,
|
||
}
|
||
|
||
/// Result of building a chunked dataset.
|
||
pub struct ChunkedDataResult {
|
||
/// Raw bytes containing all chunk data + index structures.
|
||
pub data_bytes: Vec<u8>,
|
||
/// The DataLayout v4 message bytes.
|
||
pub layout_message: Vec<u8>,
|
||
/// The FilterPipeline message bytes, if any.
|
||
pub pipeline_message: Option<Vec<u8>>,
|
||
}
|
||
|
||
/// Split raw data into chunk-sized pieces based on shape and chunk dimensions.
|
||
/// Returns a Vec of (chunk_offset_per_dim, chunk_raw_bytes).
|
||
pub fn split_into_chunks(
|
||
raw_data: &[u8],
|
||
shape: &[u64],
|
||
chunk_dims: &[u64],
|
||
element_size: usize,
|
||
) -> Vec<(Vec<u64>, Vec<u8>)> {
|
||
let rank = shape.len();
|
||
if rank == 0 {
|
||
return vec![(vec![], raw_data.to_vec())];
|
||
}
|
||
|
||
// Compute number of chunks per dimension
|
||
let mut num_chunks_per_dim = Vec::with_capacity(rank);
|
||
for d in 0..rank {
|
||
num_chunks_per_dim.push(shape[d].div_ceil(chunk_dims[d]));
|
||
}
|
||
let total_chunks: u64 = num_chunks_per_dim.iter().product();
|
||
|
||
// 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] * shape[i + 1] as usize;
|
||
}
|
||
|
||
// Chunk strides
|
||
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] as usize;
|
||
}
|
||
|
||
let chunk_total_elements: usize = chunk_dims.iter().map(|&d| d as usize).product();
|
||
|
||
let mut result = Vec::with_capacity(total_chunks as usize);
|
||
|
||
for linear_idx in 0..total_chunks {
|
||
// Convert linear index to chunk grid coordinates
|
||
let mut chunk_grid_coords = vec![0u64; rank];
|
||
let mut remaining = linear_idx;
|
||
for d in (0..rank).rev() {
|
||
chunk_grid_coords[d] = remaining % num_chunks_per_dim[d];
|
||
remaining /= num_chunks_per_dim[d];
|
||
}
|
||
|
||
// Chunk offset in dataset space
|
||
let offsets: Vec<u64> = (0..rank)
|
||
.map(|d| chunk_grid_coords[d] * chunk_dims[d])
|
||
.collect();
|
||
|
||
// Extract chunk data
|
||
let mut chunk_bytes = vec![0u8; chunk_total_elements * element_size];
|
||
|
||
for flat_idx in 0..chunk_total_elements {
|
||
let mut remaining_idx = flat_idx;
|
||
let mut ds_flat = 0usize;
|
||
let mut out_of_bounds = false;
|
||
|
||
for d in 0..rank {
|
||
let coord_in_chunk = remaining_idx / chunk_strides[d];
|
||
remaining_idx %= chunk_strides[d];
|
||
|
||
let global_coord = offsets[d] as usize + coord_in_chunk;
|
||
if global_coord >= shape[d] as usize {
|
||
out_of_bounds = true;
|
||
break;
|
||
}
|
||
ds_flat += global_coord * ds_strides[d];
|
||
}
|
||
|
||
if out_of_bounds {
|
||
// Zero-filled (already initialized)
|
||
continue;
|
||
}
|
||
|
||
let src_start = ds_flat * element_size;
|
||
let dst_start = flat_idx * element_size;
|
||
|
||
if src_start + element_size <= raw_data.len() {
|
||
chunk_bytes[dst_start..dst_start + element_size]
|
||
.copy_from_slice(&raw_data[src_start..src_start + element_size]);
|
||
}
|
||
}
|
||
|
||
result.push((offsets, chunk_bytes));
|
||
}
|
||
|
||
result
|
||
}
|
||
|
||
/// Parallel compression threshold: use rayon when chunk count exceeds this.
|
||
///
|
||
/// Lowered to 2 to enable parallel compression for typical 4-chunk workloads
|
||
/// (e.g., 128×128 matrix with 32-row chunks = 4 chunks). Rayon's overhead is
|
||
/// ~2 µs, worthwhile at ≥2 chunks with any real compression (arXiv:2206.14761).
|
||
#[cfg(feature = "parallel")]
|
||
const PARALLEL_COMPRESS_THRESHOLD: usize = 2;
|
||
|
||
/// Compress all chunks, using parallel compression when beneficial.
|
||
///
|
||
/// With the `parallel` feature and more than [`PARALLEL_COMPRESS_THRESHOLD`]
|
||
/// filtered chunks, compression runs across rayon threads; otherwise it is
|
||
/// sequential. Output order matches input order, so per-chunk bytes are
|
||
/// identical to the sequential path.
|
||
fn compress_all_chunks(
|
||
chunks: &[(Vec<u64>, Vec<u8>)],
|
||
pipeline: &Option<FilterPipeline>,
|
||
element_size: u32,
|
||
) -> Result<Vec<Vec<u8>>, FormatError> {
|
||
#[cfg(feature = "parallel")]
|
||
{
|
||
if let Some(pl) = pipeline
|
||
&& chunks.len() > PARALLEL_COMPRESS_THRESHOLD
|
||
{
|
||
use rayon::prelude::*;
|
||
return chunks
|
||
.par_iter()
|
||
.map(|(_offsets, chunk_bytes)| compress_chunk(chunk_bytes, pl, element_size))
|
||
.collect();
|
||
}
|
||
}
|
||
|
||
// Sequential fallback
|
||
chunks
|
||
.iter()
|
||
.map(|(_offsets, chunk_bytes)| {
|
||
if let Some(pl) = pipeline {
|
||
compress_chunk(chunk_bytes, pl, element_size)
|
||
} else {
|
||
Ok(chunk_bytes.clone())
|
||
}
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
/// Build the complete chunked dataset blob (chunk data + index) and return
|
||
/// layout/pipeline messages. `base_address` is where the blob will be placed in the file.
|
||
/// Serialize a v4 single chunk layout message (public for OH size estimation).
|
||
pub fn serialize_v4_single_chunk_pub(
|
||
chunk_dims: &[u32],
|
||
chunk_address: u64,
|
||
filtered_size: Option<u64>,
|
||
filter_mask: Option<u32>,
|
||
offset_size: u8,
|
||
element_size: u32,
|
||
) -> Vec<u8> {
|
||
serialize_v4_single_chunk(
|
||
chunk_dims,
|
||
chunk_address,
|
||
filtered_size,
|
||
filter_mask,
|
||
offset_size,
|
||
element_size,
|
||
)
|
||
}
|
||
|
||
/// Serialize a v4 single chunk layout message.
|
||
fn serialize_v4_single_chunk(
|
||
chunk_dims: &[u32],
|
||
chunk_address: u64,
|
||
filtered_size: Option<u64>,
|
||
filter_mask: Option<u32>,
|
||
offset_size: u8,
|
||
element_size: u32,
|
||
) -> Vec<u8> {
|
||
let mut buf = Vec::new();
|
||
buf.push(4); // version
|
||
buf.push(2); // class = chunked
|
||
|
||
// flags: bit 0 = unknown meaning in some files, bit 1 = filters for single chunk
|
||
let flags: u8 = if filtered_size.is_some() { 0x02 } else { 0x00 };
|
||
buf.push(flags);
|
||
|
||
// dimensionality = rank + 1 (chunk dims + element size dim)
|
||
let ndims = chunk_dims.len() as u8 + 1;
|
||
buf.push(ndims);
|
||
|
||
push_v4_chunk_dims(&mut buf, chunk_dims, element_size);
|
||
|
||
// chunk index type = 1 (single chunk)
|
||
buf.push(1);
|
||
|
||
// Index-specific fields
|
||
if let (Some(fs), Some(fm)) = (filtered_size, filter_mask) {
|
||
// filtered_size (length_size bytes)
|
||
buf.extend_from_slice(&fs.to_le_bytes()); // 8 bytes for length_size=8
|
||
buf.extend_from_slice(&fm.to_le_bytes()); // 4 bytes
|
||
}
|
||
|
||
// chunk address
|
||
match offset_size {
|
||
4 => buf.extend_from_slice(&(chunk_address as u32).to_le_bytes()),
|
||
8 => buf.extend_from_slice(&chunk_address.to_le_bytes()),
|
||
_ => {}
|
||
}
|
||
|
||
buf
|
||
}
|
||
|
||
/// Serialize a v4 Fixed Array layout message.
|
||
fn serialize_v4_fixed_array(
|
||
chunk_dims: &[u32],
|
||
fixed_array_address: u64,
|
||
offset_size: u8,
|
||
element_size: u32,
|
||
max_bits: u8,
|
||
) -> Vec<u8> {
|
||
let mut buf = layout_v4_chunked_prefix(chunk_dims, element_size);
|
||
|
||
// chunk index type = 3 (Fixed Array)
|
||
buf.push(3);
|
||
|
||
// max_dblk_page_nelmts_bits — must match FAHD max_nelmts_bits
|
||
buf.push(max_bits);
|
||
|
||
// Fixed Array header address
|
||
match offset_size {
|
||
4 => buf.extend_from_slice(&(fixed_array_address as u32).to_le_bytes()),
|
||
8 => buf.extend_from_slice(&fixed_array_address.to_le_bytes()),
|
||
_ => {}
|
||
}
|
||
|
||
buf
|
||
}
|
||
|
||
/// The part of a v4 chunked layout message before the chunk index type:
|
||
/// version, class, flags and the chunk dimensions (plus the element size).
|
||
/// Append a v4 layout's dimension width and its dimensions (the chunk
|
||
/// dimensions, then the element size). Each takes the fewest bytes that hold
|
||
/// the largest, as libhdf5 computes it (`H5D__chunk_set_sizes`:
|
||
/// `(log2(dim) + 8) / 8`); HDF5 2.0.0 refuses any other width.
|
||
pub(crate) fn push_v4_chunk_dims(buf: &mut Vec<u8>, chunk_dims: &[u32], element_size: u32) {
|
||
let max_dim = chunk_dims
|
||
.iter()
|
||
.copied()
|
||
.chain(core::iter::once(element_size))
|
||
.max()
|
||
.unwrap_or(1)
|
||
.max(1);
|
||
let width = (32 - max_dim.leading_zeros()).div_ceil(8) as usize;
|
||
buf.push(width as u8);
|
||
for &d in chunk_dims.iter().chain(core::iter::once(&element_size)) {
|
||
buf.extend_from_slice(&d.to_le_bytes()[..width]);
|
||
}
|
||
}
|
||
|
||
fn layout_v4_chunked_prefix(chunk_dims: &[u32], element_size: u32) -> Vec<u8> {
|
||
let mut buf = Vec::new();
|
||
buf.push(4); // version
|
||
buf.push(2); // class = chunked
|
||
|
||
let flags: u8 = 0x00;
|
||
buf.push(flags);
|
||
|
||
let ndims = chunk_dims.len() as u8 + 1;
|
||
buf.push(ndims);
|
||
|
||
push_v4_chunk_dims(&mut buf, chunk_dims, element_size);
|
||
buf
|
||
}
|
||
|
||
/// log2 of the elements per Fixed Array data block page (the library's
|
||
/// default, `H5D_FARRAY_MAX_DBLK_PAGE_NELMTS_BITS`).
|
||
const FA_PAGE_BITS: u8 = 10;
|
||
|
||
pub(crate) fn push_addr(buf: &mut Vec<u8>, addr: u64, offset_size: u8) {
|
||
match offset_size {
|
||
4 => buf.extend_from_slice(&(addr as u32).to_le_bytes()),
|
||
_ => buf.extend_from_slice(&addr.to_le_bytes()),
|
||
}
|
||
}
|
||
|
||
/// Width of the chunk-size field of a filtered chunk index element. Must
|
||
/// match the library's `H5D_FARRAY_FILT_COMPUTE_CHUNK_SIZE_LEN` (the EA and
|
||
/// B-tree v2 indexes use the same formula):
|
||
/// `1 + ((log2(unfiltered chunk bytes) + 8) / 8)`, capped at 8.
|
||
pub(crate) fn filtered_chunk_size_len(slots: &[Option<WrittenChunk>]) -> usize {
|
||
let max_raw = slots
|
||
.iter()
|
||
.flatten()
|
||
.map(|c| c.raw_size)
|
||
.max()
|
||
.unwrap_or(1);
|
||
let log2_val = if max_raw <= 1 {
|
||
0
|
||
} else {
|
||
63 - max_raw.leading_zeros()
|
||
};
|
||
(1 + ((log2_val + 8) / 8) as usize).min(8)
|
||
}
|
||
|
||
/// Append one chunk index element: the chunk's address, plus its stored size
|
||
/// and filter mask when the dataset is filtered. `None` is an unallocated
|
||
/// chunk (undefined address, zero size and mask).
|
||
pub(crate) fn push_index_element(
|
||
buf: &mut Vec<u8>,
|
||
slot: Option<&WrittenChunk>,
|
||
offset_size: u8,
|
||
chunk_size_bytes: Option<usize>,
|
||
) {
|
||
match slot {
|
||
Some(c) => {
|
||
push_addr(buf, c.address, offset_size);
|
||
if let Some(n) = chunk_size_bytes {
|
||
buf.extend_from_slice(&c.compressed_size.to_le_bytes()[..n]);
|
||
buf.extend_from_slice(&c.filter_mask.to_le_bytes());
|
||
}
|
||
}
|
||
None => {
|
||
buf.extend(core::iter::repeat_n(0xFF, offset_size as usize));
|
||
if let Some(n) = chunk_size_bytes {
|
||
buf.extend(core::iter::repeat_n(0x00, n + 4));
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Build a complete Fixed Array at a known absolute address.
|
||
///
|
||
/// `slots` holds one entry per element of the array, i.e. per chunk of the
|
||
/// dataset's *maximum* extent in the order [`crate::chunk_grid`] defines;
|
||
/// `None` marks a chunk that is not allocated. An array with more elements
|
||
/// than fit in one page (`2^FA_PAGE_BITS`) gets a paged data block: a
|
||
/// page-init bitmap after the prefix, then one checksummed page per
|
||
/// `2^FA_PAGE_BITS` elements, the last one short (`H5FA__dblock_create`).
|
||
pub fn build_fixed_array_at(
|
||
slots: &[Option<WrittenChunk>],
|
||
offset_size: u8,
|
||
length_size: u8,
|
||
has_filters: bool,
|
||
fa_base_address: u64,
|
||
) -> Vec<u8> {
|
||
let os = offset_size as usize;
|
||
let num_elements = slots.len();
|
||
|
||
let chunk_size_bytes = has_filters.then(|| filtered_chunk_size_len(slots));
|
||
let elem_size = os + chunk_size_bytes.map_or(0, |n| n + 4);
|
||
let client_id: u8 = if has_filters { 1 } else { 0 };
|
||
|
||
// FAHD total size
|
||
let fahd_total_size = 4 + 1 + 1 + 1 + 1 + length_size as usize + os + 4;
|
||
let fadb_address = fa_base_address + fahd_total_size as u64;
|
||
|
||
let mut fahd = Vec::with_capacity(fahd_total_size);
|
||
fahd.extend_from_slice(b"FAHD");
|
||
fahd.push(0); // version
|
||
fahd.push(client_id);
|
||
fahd.push(elem_size as u8);
|
||
fahd.push(FA_PAGE_BITS);
|
||
match length_size {
|
||
4 => fahd.extend_from_slice(&(num_elements as u32).to_le_bytes()),
|
||
_ => fahd.extend_from_slice(&(num_elements as u64).to_le_bytes()),
|
||
}
|
||
push_addr(&mut fahd, fadb_address, offset_size);
|
||
let checksum = jenkins_lookup3(&fahd);
|
||
fahd.extend_from_slice(&checksum.to_le_bytes());
|
||
assert_eq!(fahd.len(), fahd_total_size);
|
||
|
||
// FADB prefix
|
||
let mut fadb = Vec::new();
|
||
fadb.extend_from_slice(b"FADB");
|
||
fadb.push(0); // version
|
||
fadb.push(client_id);
|
||
push_addr(&mut fadb, fa_base_address, offset_size);
|
||
|
||
let page_nelmts = 1usize << FA_PAGE_BITS;
|
||
if num_elements <= page_nelmts {
|
||
// Unpaged: the elements follow the prefix, one checksum over both.
|
||
for slot in slots {
|
||
push_index_element(&mut fadb, slot.as_ref(), offset_size, chunk_size_bytes);
|
||
}
|
||
let fadb_checksum = jenkins_lookup3(&fadb);
|
||
fadb.extend_from_slice(&fadb_checksum.to_le_bytes());
|
||
} else {
|
||
// Paged: every page is written, so every page-init bit is set
|
||
// (MSB-first, as `H5VM_bit_set` packs them). The prefix and bitmap
|
||
// share a checksum; each page carries its own.
|
||
let npages = num_elements.div_ceil(page_nelmts);
|
||
let mut bitmap = vec![0u8; npages.div_ceil(8)];
|
||
for p in 0..npages {
|
||
bitmap[p / 8] |= 0x80 >> (p % 8);
|
||
}
|
||
fadb.extend_from_slice(&bitmap);
|
||
let prefix_checksum = jenkins_lookup3(&fadb);
|
||
fadb.extend_from_slice(&prefix_checksum.to_le_bytes());
|
||
for page in slots.chunks(page_nelmts) {
|
||
let start = fadb.len();
|
||
for slot in page {
|
||
push_index_element(&mut fadb, slot.as_ref(), offset_size, chunk_size_bytes);
|
||
}
|
||
let page_checksum = jenkins_lookup3(&fadb[start..]);
|
||
fadb.extend_from_slice(&page_checksum.to_le_bytes());
|
||
}
|
||
}
|
||
|
||
let mut combined = fahd;
|
||
combined.extend_from_slice(&fadb);
|
||
combined
|
||
}
|
||
|
||
/// Compressed chunks ready to be laid out at any file address.
|
||
///
|
||
/// Created by [`precompress_chunks`] and consumed by
|
||
/// [`build_chunked_data_from_precompressed`]. Caching this between the two
|
||
/// writer passes eliminates the double-compression that the two-pass layout
|
||
/// algorithm previously performed.
|
||
pub struct PrecompressedChunks {
|
||
/// Per-chunk: (raw_size_bytes, compressed_bytes).
|
||
pub chunks: Vec<(u64, Vec<u8>)>,
|
||
pub has_filters: bool,
|
||
pub element_size: usize,
|
||
pub shape: Vec<u64>,
|
||
pub chunk_dims: Vec<u64>,
|
||
pub pipeline_message: Option<Vec<u8>>,
|
||
}
|
||
|
||
/// Compress all chunks of a dataset without laying them out at a file address.
|
||
///
|
||
/// Call this once per dataset in Pass 1, cache the result, then call
|
||
/// [`build_chunked_data_from_precompressed`] in both Pass 1 (dummy address
|
||
/// for sizing) and Pass 2 (real address) to avoid re-compressing.
|
||
pub fn precompress_chunks(
|
||
raw_data: &[u8],
|
||
shape: &[u64],
|
||
chunk_dims: &[u64],
|
||
element_size: usize,
|
||
options: &ChunkOptions,
|
||
) -> Result<PrecompressedChunks, FormatError> {
|
||
let chunk_bytes = chunk_dims
|
||
.iter()
|
||
.try_fold(element_size as u64, |acc, &d| acc.checked_mul(d))
|
||
.and_then(|b| u32::try_from(b).ok())
|
||
.unwrap_or(0);
|
||
let pipeline = options.build_pipeline_for_chunk(element_size as u32, chunk_bytes);
|
||
let has_filters = pipeline.is_some();
|
||
let pipeline_message = pipeline.as_ref().map(|pl| pl.serialize());
|
||
|
||
let raw_chunks = split_into_chunks(raw_data, shape, chunk_dims, element_size);
|
||
let compressed = compress_all_chunks(&raw_chunks, &pipeline, element_size as u32)?;
|
||
|
||
let chunks = raw_chunks
|
||
.into_iter()
|
||
.zip(compressed)
|
||
.map(|((_offsets, raw_bytes), c)| (raw_bytes.len() as u64, c))
|
||
.collect();
|
||
|
||
Ok(PrecompressedChunks {
|
||
chunks,
|
||
has_filters,
|
||
element_size,
|
||
shape: shape.to_vec(),
|
||
chunk_dims: chunk_dims.to_vec(),
|
||
pipeline_message,
|
||
})
|
||
}
|
||
|
||
/// Lay out precompressed chunks at `base_address` and build index structures.
|
||
///
|
||
/// This is the address-dependent half of chunk writing. Call it in Pass 1
|
||
/// with a dummy address (to get the blob size), and again in Pass 2 with the
|
||
/// real address — both times reusing the same [`PrecompressedChunks`] so
|
||
/// compression happens only once.
|
||
pub fn build_chunked_data_from_precompressed(
|
||
pre: &PrecompressedChunks,
|
||
base_address: u64,
|
||
maxshape: Option<&[u64]>,
|
||
) -> Result<ChunkedDataResult, FormatError> {
|
||
let index = ChunkIndexPlan::new(&pre.shape, maxshape, &pre.chunk_dims)?;
|
||
let offset_size: u8 = 8;
|
||
let length_size: u8 = 8;
|
||
let num_chunks = pre.chunks.len();
|
||
let element_size = pre.element_size;
|
||
|
||
let mut data_buf = Vec::new();
|
||
let mut written_chunks = Vec::with_capacity(num_chunks);
|
||
|
||
for (raw_size, compressed) in &pre.chunks {
|
||
let aligned_offset = align_to_cache_line(data_buf.len());
|
||
if aligned_offset > data_buf.len() {
|
||
data_buf.resize(aligned_offset, 0u8);
|
||
}
|
||
let address = base_address + data_buf.len() as u64;
|
||
let compressed_size = compressed.len() as u64;
|
||
data_buf.extend_from_slice(compressed);
|
||
written_chunks.push(WrittenChunk {
|
||
address,
|
||
compressed_size,
|
||
raw_size: *raw_size,
|
||
filter_mask: 0,
|
||
});
|
||
}
|
||
|
||
let chunk_dims_u32: Vec<u32> = pre.chunk_dims.iter().map(|&d| d as u32).collect();
|
||
|
||
let aligned_idx = align_to_cache_line(data_buf.len());
|
||
if aligned_idx > data_buf.len() {
|
||
data_buf.resize(aligned_idx, 0u8);
|
||
}
|
||
|
||
let layout_message = match &index {
|
||
ChunkIndexPlan::ExtensibleArray(grid) => {
|
||
let ea_address = base_address + data_buf.len() as u64;
|
||
let slots = index_slots(grid, &pre.shape, &pre.chunk_dims, &written_chunks, None)?;
|
||
let ea_bytes = ea_writer::build_extensible_array_at(
|
||
&slots,
|
||
offset_size,
|
||
length_size,
|
||
pre.has_filters,
|
||
ea_address,
|
||
);
|
||
data_buf.extend_from_slice(&ea_bytes);
|
||
ea_writer::serialize_v4_extensible_array(
|
||
&chunk_dims_u32,
|
||
ea_address,
|
||
offset_size,
|
||
element_size as u32,
|
||
)
|
||
}
|
||
ChunkIndexPlan::SingleChunk => {
|
||
let chunk_addr = written_chunks[0].address;
|
||
let filtered_size = if pre.has_filters {
|
||
Some(written_chunks[0].compressed_size)
|
||
} else {
|
||
None
|
||
};
|
||
let filter_mask = if pre.has_filters { Some(0u32) } else { None };
|
||
serialize_v4_single_chunk(
|
||
&chunk_dims_u32,
|
||
chunk_addr,
|
||
filtered_size,
|
||
filter_mask,
|
||
offset_size,
|
||
element_size as u32,
|
||
)
|
||
}
|
||
ChunkIndexPlan::FixedArray(grid, nslots) => {
|
||
let fa_address = base_address + data_buf.len() as u64;
|
||
let slots = index_slots(
|
||
grid,
|
||
&pre.shape,
|
||
&pre.chunk_dims,
|
||
&written_chunks,
|
||
Some(*nslots),
|
||
)?;
|
||
let fa_bytes = build_fixed_array_at(
|
||
&slots,
|
||
offset_size,
|
||
length_size,
|
||
pre.has_filters,
|
||
fa_address,
|
||
);
|
||
data_buf.extend_from_slice(&fa_bytes);
|
||
serialize_v4_fixed_array(
|
||
&chunk_dims_u32,
|
||
fa_address,
|
||
offset_size,
|
||
element_size as u32,
|
||
FA_PAGE_BITS,
|
||
)
|
||
}
|
||
ChunkIndexPlan::BTreeV2 => {
|
||
let bt_address = base_address + data_buf.len() as u64;
|
||
let records: Vec<(Vec<u64>, &WrittenChunk)> = written_chunks
|
||
.iter()
|
||
.enumerate()
|
||
.map(|(i, c)| (scaled_coords(&pre.shape, &pre.chunk_dims, i), c))
|
||
.collect();
|
||
let (bt_bytes, node_size) = build_btree_v2_chunk_index_at(
|
||
pre.shape.len(),
|
||
&records,
|
||
offset_size,
|
||
length_size,
|
||
pre.has_filters,
|
||
bt_address,
|
||
)?;
|
||
data_buf.extend_from_slice(&bt_bytes);
|
||
serialize_v4_btree_v2(
|
||
&chunk_dims_u32,
|
||
bt_address,
|
||
offset_size,
|
||
element_size as u32,
|
||
node_size,
|
||
)
|
||
}
|
||
};
|
||
|
||
Ok(ChunkedDataResult {
|
||
data_bytes: data_buf,
|
||
layout_message,
|
||
pipeline_message: pre.pipeline_message.clone(),
|
||
})
|
||
}
|
||
|
||
/// Most slots a Fixed Array index may have before we refuse to build it: its
|
||
/// data block holds one element per chunk of the *maximum* extent, so a huge
|
||
/// finite maxshape with small chunks would otherwise exhaust memory.
|
||
const MAX_FIXED_ARRAY_SLOTS: u64 = 1 << 26;
|
||
|
||
/// Which chunk index a dataset gets, following the library's choice in
|
||
/// `H5D__layout_set_latest_indexing`: version-2 B-tree for more than one
|
||
/// unlimited dimension, Extensible Array for exactly one, Fixed Array for a
|
||
/// finite maxshape, Single Chunk when the whole maximum extent is one chunk.
|
||
enum ChunkIndexPlan {
|
||
SingleChunk,
|
||
/// The grid and the number of array elements (chunks of the max extent).
|
||
FixedArray(ChunkGrid, usize),
|
||
ExtensibleArray(ChunkGrid),
|
||
BTreeV2,
|
||
}
|
||
|
||
impl ChunkIndexPlan {
|
||
fn new(
|
||
shape: &[u64],
|
||
maxshape: Option<&[u64]>,
|
||
chunk_dims: &[u64],
|
||
) -> Result<Self, FormatError> {
|
||
let bad = |what: &str| FormatError::ChunkedReadError(format!("maxshape: {what}"));
|
||
if let Some(ms) = maxshape {
|
||
if ms.len() != shape.len() {
|
||
return Err(bad("rank differs from the shape"));
|
||
}
|
||
if ms.iter().zip(shape).any(|(&m, &s)| m < s) {
|
||
return Err(bad("smaller than the shape"));
|
||
}
|
||
}
|
||
let max = maxshape.unwrap_or(shape);
|
||
let nunlim = max.iter().filter(|&&d| d == u64::MAX).count();
|
||
match nunlim {
|
||
0 => {
|
||
let nslots = max
|
||
.iter()
|
||
.zip(chunk_dims)
|
||
.try_fold(1u64, |acc, (&m, &c)| acc.checked_mul(m.div_ceil(c.max(1))))
|
||
.filter(|&n| n <= MAX_FIXED_ARRAY_SLOTS)
|
||
.ok_or_else(|| {
|
||
bad("too many chunks for a Fixed Array index; \
|
||
use larger chunks or an unlimited dimension")
|
||
})?;
|
||
// A Single Chunk index needs that one chunk to exist; an
|
||
// empty dataset gets an all-unallocated Fixed Array instead.
|
||
let empty = shape.contains(&0);
|
||
if nslots == 1 && !empty {
|
||
Ok(Self::SingleChunk)
|
||
} else {
|
||
let grid = ChunkGrid::fixed_array(shape, Some(max), chunk_dims)?;
|
||
Ok(Self::FixedArray(grid, nslots as usize))
|
||
}
|
||
}
|
||
1 => Ok(Self::ExtensibleArray(ChunkGrid::extensible_array(
|
||
shape,
|
||
Some(max),
|
||
chunk_dims,
|
||
)?)),
|
||
_ => Ok(Self::BTreeV2),
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Place each written chunk at its linear index in `grid`. `chunks` are in
|
||
/// row-major order over the chunks of the current extent (`split_into_chunks`).
|
||
/// `len` fixes the slot count (Fixed Array); otherwise it is one past the
|
||
/// highest index used.
|
||
fn index_slots(
|
||
grid: &ChunkGrid,
|
||
shape: &[u64],
|
||
chunk_dims: &[u64],
|
||
chunks: &[WrittenChunk],
|
||
len: Option<usize>,
|
||
) -> Result<Vec<Option<WrittenChunk>>, FormatError> {
|
||
let mut placed: Vec<(usize, &WrittenChunk)> = Vec::with_capacity(chunks.len());
|
||
for (i, chunk) in chunks.iter().enumerate() {
|
||
let scaled = scaled_coords(shape, chunk_dims, i);
|
||
let idx = usize::try_from(grid.linear_index(&scaled))
|
||
.map_err(|_| FormatError::Overflow("chunk index slot".into()))?;
|
||
placed.push((idx, chunk));
|
||
}
|
||
let n = len.unwrap_or_else(|| placed.iter().map(|&(i, _)| i + 1).max().unwrap_or(0));
|
||
let mut slots = vec![None; n];
|
||
for (idx, chunk) in placed {
|
||
*slots
|
||
.get_mut(idx)
|
||
.ok_or_else(|| FormatError::Overflow("chunk index slot".into()))? = Some(chunk.clone());
|
||
}
|
||
Ok(slots)
|
||
}
|
||
|
||
/// Scaled coordinates (`offset / chunk_dim`) of the `i`-th chunk in the
|
||
/// row-major order `split_into_chunks` produces over the current extent.
|
||
fn scaled_coords(shape: &[u64], chunk_dims: &[u64], i: usize) -> Vec<u64> {
|
||
let rank = shape.len();
|
||
let mut scaled = vec![0u64; rank];
|
||
let mut rem = i as u64;
|
||
for d in (0..rank).rev() {
|
||
let n = shape[d].div_ceil(chunk_dims[d]);
|
||
scaled[d] = rem % n;
|
||
rem /= n;
|
||
}
|
||
scaled
|
||
}
|
||
|
||
/// Node size the library gives a chunk index B-tree (`H5D_BT2_NODE_SIZE`),
|
||
/// with its split and merge percentages.
|
||
const BT2_NODE_SIZE: u32 = 2048;
|
||
const BT2_SPLIT_PERCENT: u8 = 100;
|
||
const BT2_MERGE_PERCENT: u8 = 40;
|
||
/// B-tree v2 record types for chunk indexes (`H5B2_CDSET_ID`,
|
||
/// `H5B2_CDSET_FILT_ID`).
|
||
const BT2_CHUNK_UNFILTERED: u8 = 10;
|
||
const BT2_CHUNK_FILTERED: u8 = 11;
|
||
|
||
/// Build a version-2 B-tree chunk index (the library's index for datasets
|
||
/// with more than one unlimited dimension) at a known absolute address.
|
||
///
|
||
/// `records` are `(scaled coordinates, chunk)` in lexicographic order of the
|
||
/// coordinates, which is the order the library's comparator
|
||
/// (`H5VM_vector_cmp_u`) keeps them in. The tree is a single leaf: the
|
||
/// library's 2048-byte node when the records fit, otherwise a leaf node
|
||
/// sized to hold them all (the root's record count is 16-bit, so at most
|
||
/// 65535 chunks). Returns the bytes and the node size the layout message
|
||
/// must record.
|
||
fn build_btree_v2_chunk_index_at(
|
||
rank: usize,
|
||
records: &[(Vec<u64>, &WrittenChunk)],
|
||
offset_size: u8,
|
||
length_size: u8,
|
||
has_filters: bool,
|
||
base_address: u64,
|
||
) -> Result<(Vec<u8>, u32), FormatError> {
|
||
let os = offset_size as usize;
|
||
let nrec = u16::try_from(records.len()).map_err(|_| {
|
||
FormatError::ChunkedReadError(
|
||
"more than 65535 chunks with more than one unlimited dimension: \
|
||
use larger chunks"
|
||
.into(),
|
||
)
|
||
})?;
|
||
let chunk_size_bytes = has_filters.then(|| {
|
||
let slots: Vec<Option<WrittenChunk>> =
|
||
records.iter().map(|(_, c)| Some((*c).clone())).collect();
|
||
filtered_chunk_size_len(&slots)
|
||
});
|
||
let record_size = os + chunk_size_bytes.map_or(0, |n| n + 4) + 8 * rank;
|
||
// Leaf: signature, version, type, records, checksum.
|
||
let leaf_len = 4 + 1 + 1 + records.len() * record_size + 4;
|
||
let node_size = u32::try_from(leaf_len)
|
||
.map_err(|_| FormatError::Overflow("B-tree v2 leaf size".into()))?
|
||
.max(BT2_NODE_SIZE);
|
||
let tree_type = if has_filters {
|
||
BT2_CHUNK_FILTERED
|
||
} else {
|
||
BT2_CHUNK_UNFILTERED
|
||
};
|
||
|
||
let hdr_len = 4 + 1 + 1 + 4 + 2 + 2 + 1 + 1 + os + 2 + length_size as usize + 4;
|
||
let leaf_address = base_address + hdr_len as u64;
|
||
|
||
let mut out = Vec::with_capacity(hdr_len + node_size as usize);
|
||
out.extend_from_slice(b"BTHD");
|
||
out.push(0); // version
|
||
out.push(tree_type);
|
||
out.extend_from_slice(&node_size.to_le_bytes());
|
||
out.extend_from_slice(&(record_size as u16).to_le_bytes());
|
||
out.extend_from_slice(&0u16.to_le_bytes()); // depth
|
||
out.push(BT2_SPLIT_PERCENT);
|
||
out.push(BT2_MERGE_PERCENT);
|
||
if records.is_empty() {
|
||
out.extend(core::iter::repeat_n(0xFF, os));
|
||
} else {
|
||
push_addr(&mut out, leaf_address, offset_size);
|
||
}
|
||
out.extend_from_slice(&nrec.to_le_bytes());
|
||
match length_size {
|
||
4 => out.extend_from_slice(&(records.len() as u32).to_le_bytes()),
|
||
_ => out.extend_from_slice(&(records.len() as u64).to_le_bytes()),
|
||
}
|
||
let sum = jenkins_lookup3(&out);
|
||
out.extend_from_slice(&sum.to_le_bytes());
|
||
debug_assert_eq!(out.len(), hdr_len);
|
||
if records.is_empty() {
|
||
return Ok((out, node_size));
|
||
}
|
||
|
||
let leaf_start = out.len();
|
||
out.extend_from_slice(b"BTLF");
|
||
out.push(0); // version
|
||
out.push(tree_type);
|
||
for (scaled, chunk) in records {
|
||
push_index_element(&mut out, Some(chunk), offset_size, chunk_size_bytes);
|
||
for &c in scaled {
|
||
out.extend_from_slice(&c.to_le_bytes());
|
||
}
|
||
}
|
||
let sum = jenkins_lookup3(&out[leaf_start..]);
|
||
out.extend_from_slice(&sum.to_le_bytes());
|
||
// The library reads whole nodes; pad the leaf out to the node size.
|
||
out.resize(leaf_start + node_size as usize, 0);
|
||
Ok((out, node_size))
|
||
}
|
||
|
||
/// Serialize a v4 layout message for a version-2 B-tree chunk index.
|
||
fn serialize_v4_btree_v2(
|
||
chunk_dims: &[u32],
|
||
btree_address: u64,
|
||
offset_size: u8,
|
||
element_size: u32,
|
||
node_size: u32,
|
||
) -> Vec<u8> {
|
||
let mut buf = layout_v4_chunked_prefix(chunk_dims, element_size);
|
||
buf.push(5); // chunk index type = 5 (version-2 B-tree)
|
||
buf.extend_from_slice(&node_size.to_le_bytes());
|
||
buf.push(BT2_SPLIT_PERCENT);
|
||
buf.push(BT2_MERGE_PERCENT);
|
||
push_addr(&mut buf, btree_address, offset_size);
|
||
buf
|
||
}
|
||
|
||
/// Build chunked data with absolute addresses.
|
||
/// If `maxshape` has unlimited dims, uses Extensible Array index.
|
||
pub fn build_chunked_data_at(
|
||
raw_data: &[u8],
|
||
shape: &[u64],
|
||
chunk_dims: &[u64],
|
||
element_size: usize,
|
||
options: &ChunkOptions,
|
||
base_address: u64,
|
||
) -> Result<ChunkedDataResult, FormatError> {
|
||
build_chunked_data_at_ext(
|
||
raw_data,
|
||
shape,
|
||
chunk_dims,
|
||
element_size,
|
||
options,
|
||
base_address,
|
||
None,
|
||
)
|
||
}
|
||
|
||
/// Build chunked data with absolute addresses and optional maxshape.
|
||
pub fn build_chunked_data_at_ext(
|
||
raw_data: &[u8],
|
||
shape: &[u64],
|
||
chunk_dims: &[u64],
|
||
element_size: usize,
|
||
options: &ChunkOptions,
|
||
base_address: u64,
|
||
maxshape: Option<&[u64]>,
|
||
) -> Result<ChunkedDataResult, FormatError> {
|
||
let pre = precompress_chunks(raw_data, shape, chunk_dims, element_size, options)?;
|
||
build_chunked_data_from_precompressed(&pre, base_address, maxshape)
|
||
}
|
||
|
||
/// Write selected elements into an existing in-memory dataset buffer.
|
||
///
|
||
/// This performs a read-modify-write on the full buffer: elements matching
|
||
/// the selection are overwritten with the provided `new_data`. The buffer
|
||
/// must be a complete, uncompressed dataset of the given shape.
|
||
///
|
||
/// This is the in-memory equivalent of partial hyperslab writes. The caller
|
||
/// is responsible for re-chunking and re-compressing the buffer afterward.
|
||
pub fn write_selection_to_buffer(
|
||
buffer: &mut [u8],
|
||
dims: &[u64],
|
||
elem_size: usize,
|
||
selection: &crate::selection::Selection,
|
||
new_data: &[u8],
|
||
) {
|
||
use crate::selection::Selection;
|
||
|
||
match selection {
|
||
Selection::All => {
|
||
let len = buffer.len().min(new_data.len());
|
||
buffer[..len].copy_from_slice(&new_data[..len]);
|
||
}
|
||
Selection::None => {}
|
||
Selection::Hyperslab {
|
||
start,
|
||
stride,
|
||
count,
|
||
block,
|
||
} => {
|
||
let rank = dims.len();
|
||
let mut ds_strides = vec![1usize; rank];
|
||
for i in (0..rank.saturating_sub(1)).rev() {
|
||
ds_strides[i] = ds_strides[i + 1] * dims[i + 1] as usize;
|
||
}
|
||
|
||
let mut src_offset = 0usize;
|
||
|
||
#[allow(clippy::too_many_arguments)]
|
||
fn write_hyperslab(
|
||
d: usize,
|
||
rank: usize,
|
||
start: &[u64],
|
||
stride: &[u64],
|
||
count: &[u64],
|
||
block: &[u64],
|
||
dims: &[u64],
|
||
ds_strides: &[usize],
|
||
elem_size: usize,
|
||
buffer: &mut [u8],
|
||
new_data: &[u8],
|
||
src_offset: &mut usize,
|
||
current_ds_offset: usize,
|
||
) {
|
||
if d == rank {
|
||
let dst = current_ds_offset * elem_size;
|
||
let src = *src_offset * elem_size;
|
||
if dst + elem_size <= buffer.len() && src + elem_size <= new_data.len() {
|
||
buffer[dst..dst + elem_size]
|
||
.copy_from_slice(&new_data[src..src + elem_size]);
|
||
}
|
||
*src_offset += 1;
|
||
return;
|
||
}
|
||
|
||
for bi in 0..count[d] {
|
||
let block_start = start[d] + bi * stride[d];
|
||
for bj in 0..block[d] {
|
||
let coord = block_start + bj;
|
||
if coord < dims[d] {
|
||
write_hyperslab(
|
||
d + 1,
|
||
rank,
|
||
start,
|
||
stride,
|
||
count,
|
||
block,
|
||
dims,
|
||
ds_strides,
|
||
elem_size,
|
||
buffer,
|
||
new_data,
|
||
src_offset,
|
||
current_ds_offset + coord as usize * ds_strides[d],
|
||
);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
write_hyperslab(
|
||
0,
|
||
rank,
|
||
start,
|
||
stride,
|
||
count,
|
||
block,
|
||
dims,
|
||
&ds_strides,
|
||
elem_size,
|
||
buffer,
|
||
new_data,
|
||
&mut src_offset,
|
||
0,
|
||
);
|
||
}
|
||
Selection::Points(pts) => {
|
||
let rank = dims.len();
|
||
let mut ds_strides = vec![1usize; rank];
|
||
for i in (0..rank.saturating_sub(1)).rev() {
|
||
ds_strides[i] = ds_strides[i + 1] * dims[i + 1] as usize;
|
||
}
|
||
|
||
for (pi, pt) in pts.iter().enumerate() {
|
||
let flat: usize = pt
|
||
.iter()
|
||
.zip(ds_strides.iter())
|
||
.map(|(&p, &s)| p as usize * s)
|
||
.sum();
|
||
let dst = flat * elem_size;
|
||
let src = pi * elem_size;
|
||
if dst + elem_size <= buffer.len() && src + elem_size <= new_data.len() {
|
||
buffer[dst..dst + elem_size].copy_from_slice(&new_data[src..src + elem_size]);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
|
||
use super::*;
|
||
use crate::chunked_read::read_chunked_data;
|
||
use crate::data_layout::DataLayout;
|
||
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 f64_to_bytes(data: &[f64]) -> Vec<u8> {
|
||
let mut b = Vec::with_capacity(data.len() * 8);
|
||
for &v in data {
|
||
b.extend_from_slice(&v.to_le_bytes());
|
||
}
|
||
b
|
||
}
|
||
|
||
fn bytes_to_f64(data: &[u8]) -> Vec<f64> {
|
||
data.chunks(8)
|
||
.map(|c| f64::from_le_bytes(c.try_into().unwrap()))
|
||
.collect()
|
||
}
|
||
|
||
/// Helper: build a chunked file blob and read it back using read_chunked_data
|
||
fn roundtrip_chunked(
|
||
values: &[f64],
|
||
shape: &[u64],
|
||
chunk_dims: &[u64],
|
||
options: &ChunkOptions,
|
||
) -> Vec<f64> {
|
||
let raw = f64_to_bytes(values);
|
||
let base_address = 0x1000u64;
|
||
let result =
|
||
build_chunked_data_at(&raw, shape, chunk_dims, 8, options, base_address).unwrap();
|
||
|
||
// Build a fake file buffer
|
||
let file_size = base_address as usize + result.data_bytes.len();
|
||
let mut file_data = vec![0u8; file_size];
|
||
file_data[base_address as usize..].copy_from_slice(&result.data_bytes);
|
||
|
||
// Parse layout
|
||
let layout = DataLayout::parse(&result.layout_message, 8, 8).unwrap();
|
||
let dataspace = Dataspace {
|
||
space_type: DataspaceType::Simple,
|
||
rank: shape.len() as u8,
|
||
dimensions: shape.to_vec(),
|
||
max_dimensions: None,
|
||
};
|
||
let datatype = make_f64_type();
|
||
|
||
// Parse pipeline if present
|
||
let pipeline = result
|
||
.pipeline_message
|
||
.as_ref()
|
||
.map(|pm| crate::filter_pipeline::FilterPipeline::parse(pm).unwrap());
|
||
|
||
let output = read_chunked_data(
|
||
&file_data,
|
||
&layout,
|
||
&dataspace,
|
||
&datatype,
|
||
pipeline.as_ref(),
|
||
8,
|
||
8,
|
||
)
|
||
.unwrap();
|
||
|
||
bytes_to_f64(&output)
|
||
}
|
||
|
||
#[test]
|
||
fn split_1d_single_chunk() {
|
||
let data = f64_to_bytes(&[1.0, 2.0, 3.0]);
|
||
let result = split_into_chunks(&data, &[3], &[3], 8);
|
||
assert_eq!(result.len(), 1);
|
||
assert_eq!(result[0].0, vec![0]);
|
||
assert_eq!(bytes_to_f64(&result[0].1), vec![1.0, 2.0, 3.0]);
|
||
}
|
||
|
||
#[test]
|
||
fn split_1d_multiple_chunks() {
|
||
let values: Vec<f64> = (0..10).map(|i| i as f64).collect();
|
||
let data = f64_to_bytes(&values);
|
||
let result = split_into_chunks(&data, &[10], &[4], 8);
|
||
assert_eq!(result.len(), 3); // ceil(10/4) = 3
|
||
assert_eq!(result[0].0, vec![0]);
|
||
assert_eq!(result[1].0, vec![4]);
|
||
assert_eq!(result[2].0, vec![8]);
|
||
assert_eq!(bytes_to_f64(&result[0].1), vec![0.0, 1.0, 2.0, 3.0]);
|
||
assert_eq!(bytes_to_f64(&result[1].1), vec![4.0, 5.0, 6.0, 7.0]);
|
||
// Last chunk: 2 valid + 2 padding zeros
|
||
assert_eq!(bytes_to_f64(&result[2].1), vec![8.0, 9.0, 0.0, 0.0]);
|
||
}
|
||
|
||
#[test]
|
||
fn split_2d_chunks() {
|
||
// 4x4 dataset, 2x2 chunks -> 4 chunks
|
||
let values: Vec<f64> = (0..16).map(|i| i as f64).collect();
|
||
let data = f64_to_bytes(&values);
|
||
let result = split_into_chunks(&data, &[4, 4], &[2, 2], 8);
|
||
assert_eq!(result.len(), 4);
|
||
assert_eq!(result[0].0, vec![0, 0]);
|
||
assert_eq!(result[1].0, vec![0, 2]);
|
||
assert_eq!(result[2].0, vec![2, 0]);
|
||
assert_eq!(result[3].0, vec![2, 2]);
|
||
// chunk (0,0): elements [0,1,4,5]
|
||
assert_eq!(bytes_to_f64(&result[0].1), vec![0.0, 1.0, 4.0, 5.0]);
|
||
// chunk (0,2): elements [2,3,6,7]
|
||
assert_eq!(bytes_to_f64(&result[1].1), vec![2.0, 3.0, 6.0, 7.0]);
|
||
}
|
||
|
||
#[test]
|
||
fn roundtrip_1d_single_chunk_no_compression() {
|
||
let values: Vec<f64> = (0..10).map(|i| i as f64).collect();
|
||
let options = ChunkOptions {
|
||
chunk_dims: Some(vec![10]),
|
||
..Default::default()
|
||
};
|
||
let result = roundtrip_chunked(&values, &[10], &[10], &options);
|
||
assert_eq!(result, values);
|
||
}
|
||
|
||
#[cfg(feature = "deflate")]
|
||
#[test]
|
||
fn roundtrip_1d_single_chunk_deflate() {
|
||
let values: Vec<f64> = (0..100).map(|i| i as f64).collect();
|
||
let options = ChunkOptions {
|
||
chunk_dims: Some(vec![100]),
|
||
deflate_level: Some(6),
|
||
..Default::default()
|
||
};
|
||
let result = roundtrip_chunked(&values, &[100], &[100], &options);
|
||
assert_eq!(result, values);
|
||
}
|
||
|
||
#[test]
|
||
fn roundtrip_1d_multi_chunk_no_compression() {
|
||
let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
|
||
let options = ChunkOptions {
|
||
chunk_dims: Some(vec![8]),
|
||
..Default::default()
|
||
};
|
||
let result = roundtrip_chunked(&values, &[20], &[8], &options);
|
||
assert_eq!(result, values);
|
||
}
|
||
|
||
#[cfg(feature = "deflate")]
|
||
#[test]
|
||
fn roundtrip_1d_multi_chunk_deflate() {
|
||
let values: Vec<f64> = (0..100).map(|i| i as f64).collect();
|
||
let options = ChunkOptions {
|
||
chunk_dims: Some(vec![20]),
|
||
deflate_level: Some(6),
|
||
..Default::default()
|
||
};
|
||
let result = roundtrip_chunked(&values, &[100], &[20], &options);
|
||
assert_eq!(result, values);
|
||
}
|
||
|
||
#[cfg(feature = "deflate")]
|
||
#[test]
|
||
fn roundtrip_1d_shuffle_deflate() {
|
||
let values: Vec<f64> = (0..100).map(|i| i as f64).collect();
|
||
let options = ChunkOptions {
|
||
chunk_dims: Some(vec![50]),
|
||
deflate_level: Some(6),
|
||
shuffle: true,
|
||
..Default::default()
|
||
};
|
||
let result = roundtrip_chunked(&values, &[100], &[50], &options);
|
||
assert_eq!(result, values);
|
||
}
|
||
|
||
#[test]
|
||
fn roundtrip_2d_chunks() {
|
||
// 6x4 dataset, 3x2 chunks
|
||
let values: Vec<f64> = (0..24).map(|i| i as f64).collect();
|
||
let options = ChunkOptions {
|
||
chunk_dims: Some(vec![3, 2]),
|
||
..Default::default()
|
||
};
|
||
let result = roundtrip_chunked(&values, &[6, 4], &[3, 2], &options);
|
||
assert_eq!(result, values);
|
||
}
|
||
|
||
#[test]
|
||
fn align_chunk_offset_values() {
|
||
use super::CACHE_LINE_SIZE;
|
||
use super::align_chunk_offset;
|
||
let cl = CACHE_LINE_SIZE as u64;
|
||
assert_eq!(align_chunk_offset(0), 0);
|
||
assert_eq!(align_chunk_offset(1), cl);
|
||
assert_eq!(align_chunk_offset(cl), cl);
|
||
assert_eq!(align_chunk_offset(cl + 1), cl * 2);
|
||
assert_eq!(align_chunk_offset(cl * 10), cl * 10);
|
||
}
|
||
|
||
#[test]
|
||
fn chunk_addresses_are_cache_aligned() {
|
||
use super::align_chunk_offset;
|
||
let values: Vec<f64> = (0..100).map(|i| i as f64).collect();
|
||
let raw = f64_to_bytes(&values);
|
||
let base_address = 0x1000u64;
|
||
// Ensure base is aligned for this test
|
||
let base_address = align_chunk_offset(base_address);
|
||
let options = ChunkOptions {
|
||
chunk_dims: Some(vec![20]),
|
||
..Default::default()
|
||
};
|
||
let result = build_chunked_data_at(&raw, &[100], &[20], 8, &options, base_address).unwrap();
|
||
|
||
// Parse layout to get chunk addresses (via roundtrip read)
|
||
let file_size = base_address as usize + result.data_bytes.len();
|
||
let mut file_data = vec![0u8; file_size];
|
||
file_data[base_address as usize..].copy_from_slice(&result.data_bytes);
|
||
|
||
let layout = DataLayout::parse(&result.layout_message, 8, 8).unwrap();
|
||
let dataspace = Dataspace {
|
||
space_type: DataspaceType::Simple,
|
||
rank: 1,
|
||
dimensions: vec![100],
|
||
max_dimensions: None,
|
||
};
|
||
let datatype = make_f64_type();
|
||
|
||
// Verify data roundtrips correctly
|
||
let output =
|
||
read_chunked_data(&file_data, &layout, &dataspace, &datatype, None, 8, 8).unwrap();
|
||
assert_eq!(bytes_to_f64(&output), values);
|
||
}
|
||
|
||
#[test]
|
||
fn chunk_options_auto_dims() {
|
||
let options = ChunkOptions {
|
||
chunk_dims: None,
|
||
deflate_level: Some(6),
|
||
..Default::default()
|
||
};
|
||
let dims = options.resolve_chunk_dims(&[100, 50]);
|
||
assert_eq!(dims, vec![100, 50]);
|
||
}
|
||
|
||
#[test]
|
||
fn auto_chunking_splits_only_large_datasets() {
|
||
let bytes = |dims: &[u64], elem: u64| dims.iter().product::<u64>() * elem;
|
||
// Up to the target: one chunk, as before.
|
||
assert_eq!(auto_chunk_dims(&[100, 50], 8), [100, 50]);
|
||
assert_eq!(auto_chunk_dims(&[131_072], 8), [131_072]); // exactly 1 MiB
|
||
// Larger: split, keeping proportions, never above the target.
|
||
let big = auto_chunk_dims(&[4096, 2048], 8);
|
||
assert!(bytes(&big, 8) <= AUTO_CHUNK_TARGET_BYTES, "{big:?}");
|
||
assert!(bytes(&big, 8) > AUTO_CHUNK_TARGET_BYTES / 4, "{big:?}");
|
||
assert_eq!(big[0] / big[1], 2, "proportions kept: {big:?}");
|
||
// Every dimension stays within the dataset and at least 1.
|
||
for shape in [
|
||
vec![10_000_000u64],
|
||
vec![3, 5_000_000],
|
||
vec![1, 1, 9_000_000],
|
||
vec![7; 9],
|
||
] {
|
||
let dims = auto_chunk_dims(&shape, 4);
|
||
assert!(
|
||
dims.iter().zip(&shape).all(|(c, s)| *c >= 1 && c <= s),
|
||
"{shape:?} -> {dims:?}"
|
||
);
|
||
assert!(
|
||
bytes(&dims, 4) <= AUTO_CHUNK_TARGET_BYTES,
|
||
"{shape:?} -> {dims:?}"
|
||
);
|
||
}
|
||
// An empty (unlimited, unwritten) dimension still gets a usable chunk.
|
||
let growable = auto_chunk_dims(&[0, 128], 8);
|
||
assert!(growable[0] >= 1 && bytes(&growable, 8) <= AUTO_CHUNK_TARGET_BYTES);
|
||
// Explicit dimensions always win.
|
||
let explicit = ChunkOptions {
|
||
chunk_dims: Some(vec![10, 10]),
|
||
..Default::default()
|
||
};
|
||
assert_eq!(explicit.resolve_chunk_dims_for(&[4096, 2048], 8), [10, 10]);
|
||
}
|
||
|
||
#[test]
|
||
fn chunk_options_pipeline_deflate() {
|
||
// Auto-shuffle is applied before compression by default (matches h5py).
|
||
let options = ChunkOptions {
|
||
deflate_level: Some(6),
|
||
..Default::default()
|
||
};
|
||
let pl = options.build_pipeline(8).unwrap();
|
||
assert_eq!(pl.filters.len(), 2);
|
||
assert_eq!(pl.filters[0].filter_id, FILTER_SHUFFLE);
|
||
assert_eq!(pl.filters[1].filter_id, FILTER_DEFLATE);
|
||
}
|
||
|
||
#[test]
|
||
fn chunk_options_pipeline_deflate_no_shuffle() {
|
||
// Users can opt out of auto-shuffle with no_shuffle = true.
|
||
let options = ChunkOptions {
|
||
deflate_level: Some(6),
|
||
no_shuffle: true,
|
||
..Default::default()
|
||
};
|
||
let pl = options.build_pipeline(8).unwrap();
|
||
assert_eq!(pl.filters.len(), 1);
|
||
assert_eq!(pl.filters[0].filter_id, FILTER_DEFLATE);
|
||
}
|
||
|
||
#[test]
|
||
fn chunk_options_pipeline_lz4() {
|
||
// Auto-shuffle before LZ4.
|
||
let options = ChunkOptions {
|
||
lz4: true,
|
||
..Default::default()
|
||
};
|
||
let pl = options.build_pipeline(8).unwrap();
|
||
assert_eq!(pl.filters.len(), 2);
|
||
assert_eq!(pl.filters[0].filter_id, FILTER_SHUFFLE);
|
||
assert_eq!(pl.filters[1].filter_id, FILTER_LZ4);
|
||
}
|
||
|
||
#[test]
|
||
fn chunk_options_pipeline_zstd() {
|
||
// Auto-shuffle before Zstd.
|
||
let options = ChunkOptions {
|
||
zstd_level: Some(3),
|
||
..Default::default()
|
||
};
|
||
let pl = options.build_pipeline(8).unwrap();
|
||
assert_eq!(pl.filters.len(), 2);
|
||
assert_eq!(pl.filters[0].filter_id, FILTER_SHUFFLE);
|
||
assert_eq!(pl.filters[1].filter_id, FILTER_ZSTD);
|
||
assert_eq!(pl.filters[1].client_data, vec![3]);
|
||
}
|
||
|
||
#[test]
|
||
fn chunk_options_pipeline_lzf() {
|
||
let options = ChunkOptions {
|
||
plugin: Some(PluginFilter::Lzf),
|
||
..Default::default()
|
||
};
|
||
assert!(options.is_chunked());
|
||
let pl = options.build_pipeline_for_chunk(8, 800).unwrap();
|
||
assert_eq!(pl.filters.len(), 2);
|
||
assert_eq!(pl.filters[0].filter_id, FILTER_SHUFFLE);
|
||
assert_eq!(pl.filters[1].filter_id, FILTER_LZF);
|
||
assert_eq!(pl.filters[1].client_data, vec![4, 0x0105, 800]);
|
||
}
|
||
|
||
#[test]
|
||
fn chunk_options_pipeline_bitshuffle_has_no_auto_shuffle() {
|
||
let options = ChunkOptions {
|
||
plugin: Some(PluginFilter::Bitshuffle {
|
||
block_size: 0,
|
||
compression: BitshuffleCompression::Zstd { level: 5 },
|
||
}),
|
||
..Default::default()
|
||
};
|
||
let pl = options.build_pipeline(4).unwrap();
|
||
assert_eq!(pl.filters.len(), 1);
|
||
assert_eq!(pl.filters[0].filter_id, FILTER_BITSHUFFLE);
|
||
assert_eq!(pl.filters[0].client_data, vec![0, 4, 4, 0, 3, 5]);
|
||
}
|
||
|
||
#[test]
|
||
fn chunk_options_zstd_priority_over_deflate() {
|
||
let options = ChunkOptions {
|
||
deflate_level: Some(6),
|
||
zstd_level: Some(3),
|
||
..Default::default()
|
||
};
|
||
let pl = options.build_pipeline(8).unwrap();
|
||
// shuffle + zstd (deflate is ignored when zstd wins priority)
|
||
assert_eq!(pl.filters.len(), 2);
|
||
assert_eq!(pl.filters[0].filter_id, FILTER_SHUFFLE);
|
||
assert_eq!(pl.filters[1].filter_id, FILTER_ZSTD);
|
||
}
|
||
|
||
#[test]
|
||
fn chunk_options_pipeline_shuffle_deflate_fletcher32() {
|
||
let options = ChunkOptions {
|
||
deflate_level: Some(6),
|
||
shuffle: true,
|
||
fletcher32: true,
|
||
..Default::default()
|
||
};
|
||
let pl = options.build_pipeline(8).unwrap();
|
||
assert_eq!(pl.filters.len(), 3);
|
||
assert_eq!(pl.filters[0].filter_id, FILTER_SHUFFLE);
|
||
assert_eq!(pl.filters[1].filter_id, FILTER_DEFLATE);
|
||
assert_eq!(pl.filters[2].filter_id, FILTER_FLETCHER32);
|
||
}
|
||
|
||
#[test]
|
||
fn serialize_v4_single_chunk_no_filters_roundtrip() {
|
||
let msg = serialize_v4_single_chunk(&[20], 0x1000, None, None, 8, 8);
|
||
let layout = DataLayout::parse(&msg, 8, 8).unwrap();
|
||
match layout {
|
||
DataLayout::Chunked {
|
||
chunk_dimensions,
|
||
btree_address,
|
||
version,
|
||
chunk_index_type,
|
||
single_chunk_filtered_size,
|
||
single_chunk_filter_mask,
|
||
..
|
||
} => {
|
||
assert_eq!(version, 4);
|
||
assert_eq!(chunk_index_type, Some(1));
|
||
assert_eq!(chunk_dimensions, vec![20, 8]);
|
||
assert_eq!(btree_address, Some(0x1000));
|
||
assert_eq!(single_chunk_filtered_size, None);
|
||
assert_eq!(single_chunk_filter_mask, None);
|
||
}
|
||
_ => panic!("expected chunked layout"),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn serialize_v4_single_chunk_with_filters_roundtrip() {
|
||
let msg = serialize_v4_single_chunk(&[100], 0x2000, Some(500), Some(0), 8, 8);
|
||
let layout = DataLayout::parse(&msg, 8, 8).unwrap();
|
||
match layout {
|
||
DataLayout::Chunked {
|
||
btree_address,
|
||
single_chunk_filtered_size,
|
||
single_chunk_filter_mask,
|
||
..
|
||
} => {
|
||
assert_eq!(btree_address, Some(0x2000));
|
||
assert_eq!(single_chunk_filtered_size, Some(500));
|
||
assert_eq!(single_chunk_filter_mask, Some(0));
|
||
}
|
||
_ => panic!("expected chunked layout"),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn serialize_v4_fixed_array_roundtrip() {
|
||
let msg = serialize_v4_fixed_array(&[20], 0x3000, 8, 8, 4);
|
||
let layout = DataLayout::parse(&msg, 8, 8).unwrap();
|
||
match layout {
|
||
DataLayout::Chunked {
|
||
version,
|
||
chunk_index_type,
|
||
btree_address,
|
||
chunk_dimensions,
|
||
..
|
||
} => {
|
||
assert_eq!(version, 4);
|
||
assert_eq!(chunk_index_type, Some(3));
|
||
assert_eq!(btree_address, Some(0x3000));
|
||
assert_eq!(chunk_dimensions, vec![20, 8]);
|
||
}
|
||
_ => panic!("expected chunked layout"),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn build_fixed_array_valid_structure() {
|
||
let chunks = vec![
|
||
WrittenChunk {
|
||
address: 0x1000,
|
||
compressed_size: 160,
|
||
raw_size: 160,
|
||
filter_mask: 0,
|
||
},
|
||
WrittenChunk {
|
||
address: 0x10A0,
|
||
compressed_size: 160,
|
||
raw_size: 160,
|
||
filter_mask: 0,
|
||
},
|
||
];
|
||
let slots: Vec<_> = chunks.into_iter().map(Some).collect();
|
||
let fa = build_fixed_array_at(&slots, 8, 8, false, 0x2000);
|
||
// Should start with FAHD
|
||
assert_eq!(&fa[0..4], b"FAHD");
|
||
// FAHD size = 4+1+1+1+1+8+8+4 = 28
|
||
// FADB starts at offset 28
|
||
assert_eq!(&fa[28..32], b"FADB");
|
||
}
|
||
|
||
// ---- Extensible Array tests ----
|
||
|
||
#[test]
|
||
fn serialize_v4_extensible_array_roundtrip() {
|
||
let msg = ea_writer::serialize_v4_extensible_array(&[10], 0x4000, 8, 8);
|
||
let layout = DataLayout::parse(&msg, 8, 8).unwrap();
|
||
match layout {
|
||
DataLayout::Chunked {
|
||
version,
|
||
chunk_index_type,
|
||
btree_address,
|
||
chunk_dimensions,
|
||
..
|
||
} => {
|
||
assert_eq!(version, 4);
|
||
assert_eq!(chunk_index_type, Some(4));
|
||
assert_eq!(btree_address, Some(0x4000));
|
||
assert_eq!(chunk_dimensions, vec![10, 8]);
|
||
}
|
||
_ => panic!("expected chunked layout"),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn build_extensible_array_valid_structure() {
|
||
let chunks = vec![
|
||
WrittenChunk {
|
||
address: 0x1000,
|
||
compressed_size: 80,
|
||
raw_size: 80,
|
||
filter_mask: 0,
|
||
},
|
||
WrittenChunk {
|
||
address: 0x1050,
|
||
compressed_size: 80,
|
||
raw_size: 80,
|
||
filter_mask: 0,
|
||
},
|
||
];
|
||
let slots: Vec<_> = chunks.into_iter().map(Some).collect();
|
||
let ea = ea_writer::build_extensible_array_at(&slots, 8, 8, false, 0x2000);
|
||
assert_eq!(&ea[0..4], b"EAHD");
|
||
// Find EAIB after EAHD: 12 fixed + 6*8 stats + 8 addr + 4 checksum = 72
|
||
let aehd_size = 4 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 6 * 8 + 8 + 4;
|
||
assert_eq!(&ea[aehd_size..aehd_size + 4], b"EAIB");
|
||
}
|
||
|
||
/// Helper: roundtrip with EA (maxshape)
|
||
fn roundtrip_ea(
|
||
values: &[f64],
|
||
shape: &[u64],
|
||
chunk_dims: &[u64],
|
||
maxshape: &[u64],
|
||
) -> Vec<f64> {
|
||
let raw = f64_to_bytes(values);
|
||
let base_address = 0x1000u64;
|
||
let options = ChunkOptions {
|
||
chunk_dims: Some(chunk_dims.to_vec()),
|
||
..Default::default()
|
||
};
|
||
let result = build_chunked_data_at_ext(
|
||
&raw,
|
||
shape,
|
||
chunk_dims,
|
||
8,
|
||
&options,
|
||
base_address,
|
||
Some(maxshape),
|
||
)
|
||
.unwrap();
|
||
|
||
let file_size = base_address as usize + result.data_bytes.len();
|
||
let mut file_data = vec![0u8; file_size];
|
||
file_data[base_address as usize..].copy_from_slice(&result.data_bytes);
|
||
|
||
let layout = DataLayout::parse(&result.layout_message, 8, 8).unwrap();
|
||
// Verify it uses EA index
|
||
match &layout {
|
||
DataLayout::Chunked {
|
||
chunk_index_type, ..
|
||
} => {
|
||
assert_eq!(*chunk_index_type, Some(4), "expected EA index type");
|
||
}
|
||
_ => panic!("expected chunked layout"),
|
||
}
|
||
|
||
let dataspace = Dataspace {
|
||
space_type: DataspaceType::Simple,
|
||
rank: shape.len() as u8,
|
||
dimensions: shape.to_vec(),
|
||
max_dimensions: Some(maxshape.to_vec()),
|
||
};
|
||
let datatype = make_f64_type();
|
||
|
||
let output =
|
||
read_chunked_data(&file_data, &layout, &dataspace, &datatype, None, 8, 8).unwrap();
|
||
|
||
bytes_to_f64(&output)
|
||
}
|
||
|
||
#[test]
|
||
fn ea_roundtrip_1d_inline_only() {
|
||
let values: Vec<f64> = (0..10).map(|i| i as f64).collect();
|
||
let result = roundtrip_ea(&values, &[10], &[10], &[u64::MAX]);
|
||
assert_eq!(result, values);
|
||
}
|
||
|
||
#[test]
|
||
fn ea_roundtrip_1d_multi_chunks() {
|
||
let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
|
||
let result = roundtrip_ea(&values, &[20], &[5], &[u64::MAX]);
|
||
assert_eq!(result, values);
|
||
}
|
||
|
||
#[test]
|
||
fn ea_roundtrip_1d_many_chunks() {
|
||
let values: Vec<f64> = (0..100).map(|i| i as f64).collect();
|
||
let result = roundtrip_ea(&values, &[100], &[10], &[u64::MAX]);
|
||
assert_eq!(result, values);
|
||
}
|
||
|
||
// ---- h5py round-trip tests for chunked writes ----
|
||
|
||
/// The Python interpreter to drive interop checks with.
|
||
///
|
||
/// `CLAWHDF5_PYTHON` lets these run against a virtualenv holding h5py,
|
||
/// which on a PEP 668 "externally managed" system is the only place it
|
||
/// can be installed. Without it the suite silently skips, and a silent
|
||
/// skip here is how a datatype bug once reached a release.
|
||
#[cfg(feature = "std")]
|
||
fn python() -> String {
|
||
std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string())
|
||
}
|
||
|
||
#[cfg(feature = "std")]
|
||
fn h5py_available() -> bool {
|
||
std::process::Command::new(python())
|
||
.args(["-c", "import h5py"])
|
||
.output()
|
||
.map(|o| o.status.success())
|
||
.unwrap_or(false)
|
||
}
|
||
|
||
#[cfg(feature = "std")]
|
||
fn h5py_run(script: &str) -> String {
|
||
if !h5py_available() {
|
||
panic!("h5py not installed — skipping interop test");
|
||
}
|
||
let o = std::process::Command::new(python())
|
||
.args(["-c", script])
|
||
.output()
|
||
.expect("python interpreter");
|
||
if !o.status.success() {
|
||
panic!("h5py: {}", String::from_utf8_lossy(&o.stderr));
|
||
}
|
||
String::from_utf8(o.stdout).unwrap().trim().to_string()
|
||
}
|
||
|
||
#[cfg(feature = "std")]
|
||
#[test]
|
||
#[ignore = "requires Python h5py module"]
|
||
fn h5py_reads_multiple_chunked_datasets() {
|
||
use crate::file_writer::FileWriter;
|
||
let mut fw = FileWriter::new();
|
||
let data1: Vec<f64> = (0..50).map(|i| i as f64).collect();
|
||
let data2: Vec<f64> = (0..30).map(|i| (i * 10) as f64).collect();
|
||
fw.create_dataset("a")
|
||
.with_f64_data(&data1)
|
||
.with_shape(&[50])
|
||
.with_chunks(&[25]);
|
||
fw.create_dataset("b")
|
||
.with_f64_data(&data2)
|
||
.with_shape(&[30])
|
||
.with_chunks(&[10]);
|
||
let bytes = fw.finish().unwrap();
|
||
let path = std::env::temp_dir().join("clawhdf5_chunked_multi.h5");
|
||
std::fs::write(&path, &bytes).unwrap();
|
||
let script = format!(
|
||
"import h5py,json; f=h5py.File('{}','r'); print(json.dumps({{'a':f['a'][:].tolist(),'b':f['b'][:].tolist()}}))",
|
||
path.display()
|
||
);
|
||
let v: serde_json::Value = serde_json::from_str(&h5py_run(&script)).unwrap();
|
||
let va: Vec<f64> = serde_json::from_value(v["a"].clone()).unwrap();
|
||
let vb: Vec<f64> = serde_json::from_value(v["b"].clone()).unwrap();
|
||
assert_eq!(va, data1);
|
||
assert_eq!(vb, data2);
|
||
}
|
||
|
||
#[cfg(feature = "std")]
|
||
#[test]
|
||
#[ignore = "requires Python h5py module"]
|
||
fn h5py_reads_chunked_with_attrs() {
|
||
use crate::file_writer::{AttrValue, FileWriter};
|
||
let mut fw = FileWriter::new();
|
||
let data: Vec<f64> = (0..50).map(|i| i as f64).collect();
|
||
fw.create_dataset("data")
|
||
.with_f64_data(&data)
|
||
.with_shape(&[50])
|
||
.with_chunks(&[25])
|
||
.set_attr("units", AttrValue::String("meters".to_string()));
|
||
let bytes = fw.finish().unwrap();
|
||
let path = std::env::temp_dir().join("clawhdf5_chunked_attrs.h5");
|
||
std::fs::write(&path, &bytes).unwrap();
|
||
let script = format!(
|
||
"import h5py,json; f=h5py.File('{}','r'); d=f['data']; print(json.dumps({{'values':d[:].tolist(),'units':d.attrs['units'].decode() if isinstance(d.attrs['units'],bytes) else str(d.attrs['units'])}}))",
|
||
path.display()
|
||
);
|
||
let v: serde_json::Value = serde_json::from_str(&h5py_run(&script)).unwrap();
|
||
let values: Vec<f64> = serde_json::from_value(v["values"].clone()).unwrap();
|
||
assert_eq!(values, data);
|
||
assert_eq!(v["units"], serde_json::json!("meters"));
|
||
}
|
||
|
||
// --- LZ4 chunked roundtrip tests ---
|
||
|
||
#[cfg(feature = "lz4")]
|
||
#[test]
|
||
fn roundtrip_1d_single_chunk_lz4() {
|
||
let values: Vec<f64> = (0..100).map(|i| i as f64).collect();
|
||
let options = ChunkOptions {
|
||
chunk_dims: Some(vec![100]),
|
||
lz4: true,
|
||
..Default::default()
|
||
};
|
||
let result = roundtrip_chunked(&values, &[100], &[100], &options);
|
||
assert_eq!(result, values);
|
||
}
|
||
|
||
#[cfg(feature = "lz4")]
|
||
#[test]
|
||
fn roundtrip_1d_multi_chunk_lz4() {
|
||
let values: Vec<f64> = (0..100).map(|i| i as f64).collect();
|
||
let options = ChunkOptions {
|
||
chunk_dims: Some(vec![20]),
|
||
lz4: true,
|
||
..Default::default()
|
||
};
|
||
let result = roundtrip_chunked(&values, &[100], &[20], &options);
|
||
assert_eq!(result, values);
|
||
}
|
||
|
||
// --- Zstd chunked roundtrip tests ---
|
||
|
||
#[cfg(feature = "zstd")]
|
||
#[test]
|
||
fn roundtrip_1d_single_chunk_zstd() {
|
||
let values: Vec<f64> = (0..100).map(|i| i as f64).collect();
|
||
let options = ChunkOptions {
|
||
chunk_dims: Some(vec![100]),
|
||
zstd_level: Some(3),
|
||
..Default::default()
|
||
};
|
||
let result = roundtrip_chunked(&values, &[100], &[100], &options);
|
||
assert_eq!(result, values);
|
||
}
|
||
|
||
#[cfg(feature = "zstd")]
|
||
#[test]
|
||
fn roundtrip_1d_multi_chunk_zstd() {
|
||
let values: Vec<f64> = (0..100).map(|i| i as f64).collect();
|
||
let options = ChunkOptions {
|
||
chunk_dims: Some(vec![20]),
|
||
zstd_level: Some(1),
|
||
..Default::default()
|
||
};
|
||
let result = roundtrip_chunked(&values, &[100], &[20], &options);
|
||
assert_eq!(result, values);
|
||
}
|
||
|
||
#[cfg(feature = "zstd")]
|
||
#[test]
|
||
fn roundtrip_1d_shuffle_zstd() {
|
||
let values: Vec<f64> = (0..100).map(|i| i as f64).collect();
|
||
let options = ChunkOptions {
|
||
chunk_dims: Some(vec![50]),
|
||
zstd_level: Some(3),
|
||
shuffle: true,
|
||
..Default::default()
|
||
};
|
||
let result = roundtrip_chunked(&values, &[100], &[50], &options);
|
||
assert_eq!(result, values);
|
||
}
|
||
}
|