#[cfg(not(feature = "std"))]
extern crate alloc;
#[cfg(not(feature = "std"))]
use alloc::{format, vec, vec::Vec};
use crate::checksum::jenkins_lookup3;
const HADDR_UNDEF: u64 = u64::MAX;
use core::num::NonZeroUsize;
use crate::chunk_grid::{ChunkGrid, GridOrder};
use crate::convert::{TryToUsize, nonzero_usize_from};
use crate::error::FormatError;
use crate::extensible_array::{DataBlockGeom, EaGeometry, ExtensibleArrayHeader, SuperBlockGeom};
use crate::fill_value::FillPattern;
#[cfg(feature = "zfp")]
use crate::filter_pipeline::FILTER_ZFP;
use crate::filter_pipeline::{
FILTER_DEFLATE, FILTER_FLETCHER32, FILTER_LZF, FILTER_SCALEOFFSET, FILTER_SHUFFLE,
FilterDescription, FilterPipeline, H5Z_FLAG_OPTIONAL,
};
use crate::filters::{ChunkContext, compress_chunk_with};
use crate::scaleoffset::{FillAvailability, ScaleOffset, build_cd_values};
pub(crate) const FIXED_ARRAY_PAGE_BITS: u8 = 10;
const INDEX_OFFSET_SIZE: u8 = 8;
const INDEX_LENGTH_SIZE: u8 = 8;
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum FilterKind {
#[cfg(feature = "zfp")]
Zfp(f64),
ScaleOffset(ScaleOffset, FillAvailability),
Shuffle,
Lzf,
Deflate(u32),
Fletcher32,
}
impl FilterKind {
pub(crate) fn filter_id(self) -> u16 {
match self {
#[cfg(feature = "zfp")]
Self::Zfp(_) => FILTER_ZFP,
Self::ScaleOffset(..) => FILTER_SCALEOFFSET,
Self::Shuffle => FILTER_SHUFFLE,
Self::Lzf => FILTER_LZF,
Self::Deflate(_) => FILTER_DEFLATE,
Self::Fletcher32 => FILTER_FLETCHER32,
}
}
fn canonical_rank(self) -> u8 {
match self {
#[cfg(feature = "zfp")]
Self::Zfp(_) => 0,
Self::ScaleOffset(..) => 1,
Self::Shuffle => 2,
Self::Lzf => 3,
Self::Deflate(_) => 4,
Self::Fletcher32 => 5,
}
}
fn default_optional(self) -> bool {
matches!(self, Self::Lzf)
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct FilterSpec {
pub kind: FilterKind,
pub optional: bool,
}
#[derive(Debug, Clone, Default)]
pub struct ChunkOptions {
pub chunk_dims: Option<Vec<u64>>,
pub filters: Vec<FilterSpec>,
}
impl ChunkOptions {
pub fn is_chunked(&self) -> bool {
self.chunk_dims.is_some() || !self.filters.is_empty()
}
pub fn set_filter(&mut self, kind: FilterKind) {
let spec = FilterSpec {
optional: kind.default_optional(),
kind,
};
if let Some(slot) = self
.filters
.iter_mut()
.find(|f| f.kind.filter_id() == kind.filter_id())
{
*slot = spec;
return;
}
let at = self
.filters
.iter()
.position(|f| f.kind.canonical_rank() > kind.canonical_rank())
.unwrap_or(self.filters.len());
self.filters.insert(at, spec);
}
pub fn push_filter(&mut self, spec: FilterSpec) {
self.filters.push(spec);
}
fn has(&self, id: u16) -> bool {
self.filters.iter().any(|f| f.kind.filter_id() == id)
}
pub fn refuse_unavailable_filters(&self) -> Result<(), &'static str> {
#[cfg(not(feature = "deflate"))]
if self.has(FILTER_DEFLATE) {
return Err("deflate compression requires the `deflate` crate feature");
}
Ok(())
}
#[cfg(feature = "zfp")]
#[inline]
fn zfp_enabled(&self) -> bool {
self.has(FILTER_ZFP)
}
#[cfg(not(feature = "zfp"))]
#[inline]
fn zfp_enabled(&self) -> bool {
false
}
fn refuse_conflicting_filters(&self) -> Result<(), FormatError> {
let clash = |a: &str, b: &str| {
Err(FormatError::FilterError(format!(
"{a} and {b} cannot be combined on one dataset"
)))
};
if self.zfp_enabled() {
if self.has(FILTER_SCALEOFFSET) {
return clash("scale-offset", "ZFP");
}
if self.has(FILTER_SHUFFLE) {
return clash("shuffle", "ZFP");
}
if self.has(FILTER_LZF) {
return clash("lzf", "ZFP");
}
if self.has(FILTER_DEFLATE) {
return clash("deflate", "ZFP");
}
}
if self.has(FILTER_SCALEOFFSET) && self.has(FILTER_SHUFFLE) {
return clash("shuffle", "scale-offset");
}
if self.has(FILTER_LZF) && self.has(FILTER_DEFLATE) {
return clash("lzf", "deflate");
}
Ok(())
}
pub fn build_pipeline(
&self,
ctx: &ChunkContext<'_>,
fill: FillPattern<'_>,
) -> Result<Option<FilterPipeline>, FormatError> {
self.refuse_conflicting_filters()?;
let element_size = ctx.element_size.get();
let chunk_dims = ctx.chunk_dims;
let scale_offset_type = ctx.scale_offset_type;
let _ = ctx.element_type;
let mut filters = Vec::with_capacity(self.filters.len());
for spec in &self.filters {
let flags = if spec.optional { H5Z_FLAG_OPTIONAL } else { 0 };
filters.push(match spec.kind {
#[cfg(feature = "zfp")]
FilterKind::Zfp(rate) => {
let elem_ty = ctx.element_type.ok_or_else(|| {
FormatError::UnsupportedZfp(
"ZFP compression requires the dataset's datatype to be one \
of f32, f64, i32, or i64"
.into(),
)
})?;
FilterDescription {
filter_id: FILTER_ZFP,
name: Some("zfp".into()),
flags,
client_data: crate::zfp::zfp_cd_values_rate(rate, elem_ty, chunk_dims)?,
}
}
FilterKind::ScaleOffset(mode, fill_avail) => {
let ty = scale_offset_type.ok_or_else(|| {
FormatError::FilterError(
"scale-offset requires an integer or floating-point scalar \
datatype with a definite (little/big endian) byte order"
.into(),
)
})?;
let nelmts =
u32::try_from(chunk_dims.iter().product::<u64>()).map_err(|_| {
FormatError::FilterError(
"scale-offset: chunk has too many elements".into(),
)
})?;
FilterDescription {
filter_id: FILTER_SCALEOFFSET,
name: None,
flags,
client_data: build_cd_values(
mode,
ty,
element_size,
nelmts,
fill_avail.with_value(fill)?,
)?,
}
}
FilterKind::Shuffle => FilterDescription {
filter_id: FILTER_SHUFFLE,
name: None,
flags,
client_data: vec![element_size],
},
FilterKind::Lzf => FilterDescription {
filter_id: FILTER_LZF,
name: Some("lzf".into()),
flags,
client_data: crate::lzf::h5py_cd_values(element_size, chunk_dims).to_vec(),
},
FilterKind::Deflate(level) => FilterDescription {
filter_id: FILTER_DEFLATE,
name: None,
flags,
client_data: vec![level],
},
FilterKind::Fletcher32 => FilterDescription {
filter_id: FILTER_FLETCHER32,
name: None,
flags,
client_data: vec![],
},
});
}
if filters.is_empty() {
Ok(None)
} else {
Ok(Some(FilterPipeline {
version: 2,
filters,
}))
}
}
pub fn resolve_chunk_dims(&self, shape: &[u64]) -> Vec<u64> {
if let Some(ref dims) = self.chunk_dims {
dims.clone()
} else {
shape.to_vec()
}
}
pub fn validate_geometry(
&self,
shape: &[u64],
maxshape: Option<&[u64]>,
) -> Result<(), &'static str> {
if shape.is_empty() {
return Err("a scalar dataset cannot be chunked, filtered, or extensible");
}
if let Some(dims) = self.chunk_dims.as_deref() {
if dims.len() != shape.len() {
return Err("chunk dimensions must have the same rank as the dataset shape");
}
if dims.contains(&0) {
return Err("chunk dimensions must all be non-zero");
}
} else if shape.contains(&0) {
return Err(
"a zero-element dataset must be given explicit chunk dimensions \
(its shape has none to derive)",
);
}
if let Some(ms) = maxshape {
if ms.len() != shape.len() {
return Err("maxshape must have the same rank as the dataset shape");
}
if ms.iter().zip(shape).any(|(&m, &d)| m != u64::MAX && m < d) {
return Err("maxshape must be at least the current shape in every dimension");
}
if ms.iter().filter(|&&m| m == u64::MAX).count() > 1 {
return Err(
"at most one dimension of a maxshape may be unlimited; the chunk index \
for more than one is a version-2 B-tree, which this crate cannot write",
);
}
}
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct WrittenChunk {
pub address: u64,
pub compressed_size: u64,
pub filter_mask: u32,
}
const MAX_UNUSED_INDEX_BYTES: u64 = 32 << 20;
pub(crate) struct IndexSlots<'a> {
chunks: &'a [WrittenChunk],
scattered: Vec<(usize, usize)>,
len: usize,
}
impl<'a> IndexSlots<'a> {
pub(crate) fn new(
chunks: &'a [WrittenChunk],
slot_of: &[u64],
len: u64,
) -> Result<Self, FormatError> {
debug_assert_eq!(
chunks.len(),
slot_of.len(),
"every chunk has exactly one index slot"
);
let len = len.to_usize()?;
if len == chunks.len() && slot_of.iter().enumerate().all(|(i, &s)| s == i as u64) {
return Ok(Self {
chunks,
scattered: Vec::new(),
len,
});
}
let mut scattered = Vec::with_capacity(chunks.len());
for (i, &slot) in slot_of.iter().enumerate() {
scattered.push((slot.to_usize()?, i));
}
scattered.sort_unstable();
Ok(Self {
chunks,
scattered,
len,
})
}
pub(crate) fn dense(chunks: &'a [WrittenChunk]) -> Self {
Self {
len: chunks.len(),
chunks,
scattered: Vec::new(),
}
}
pub(crate) fn len(&self) -> usize {
self.len
}
pub(crate) fn any_occupied(&self, start: u64, count: u64) -> bool {
let Ok(start) = usize::try_from(start) else {
return false;
};
let end = usize::try_from(count)
.ok()
.and_then(|c| start.checked_add(c))
.unwrap_or(usize::MAX);
if self.scattered.is_empty() {
return start < self.chunks.len().min(end);
}
let from = self.scattered.partition_point(|&(s, _)| s < start);
self.scattered.get(from).is_some_and(|&(s, _)| s < end)
}
pub(crate) fn at(&self, slot: usize) -> Option<&WrittenChunk> {
if self.scattered.is_empty() {
return self.chunks.get(slot);
}
self.scattered
.binary_search_by_key(&slot, |&(s, _)| s)
.ok()
.map(|i| &self.chunks[self.scattered[i].1])
}
}
#[derive(Clone, Copy)]
pub(crate) enum SlotOccupancy<'a> {
Dense(u64),
Sparse(&'a IndexSlots<'a>),
Slots(&'a [u64]),
}
impl SlotOccupancy<'_> {
fn any_occupied(&self, start: u64, count: u64) -> bool {
match self {
Self::Dense(len) => start < *len,
Self::Sparse(slots) => slots.any_occupied(start, count),
Self::Slots(sorted) => {
let end = start.saturating_add(count);
let from = sorted.partition_point(|&s| s < start);
sorted.get(from).is_some_and(|&s| s < end)
}
}
}
}
pub struct ChunkedDataResult {
pub data_bytes: Vec<u8>,
pub layout_message: Vec<u8>,
pub pipeline_message: Option<Vec<u8>>,
}
pub fn split_into_chunks(
raw_data: &[u8],
shape: &[u64],
chunk_dims: &[u64],
element_size: NonZeroUsize,
fill: FillPattern<'_>,
) -> Result<Vec<Vec<u8>>, FormatError> {
let rank = shape.len();
if rank == 0 {
return Ok(vec![raw_data.to_vec()]);
}
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();
let mut ds_strides = vec![1usize; rank];
for i in (0..rank.saturating_sub(1)).rev() {
#[expect(
clippy::cast_possible_truncation,
reason = "dataset dimension derived from the in-memory write request; bounded by addressable memory"
)]
let dim = shape[i + 1] as usize;
ds_strides[i] = ds_strides[i + 1] * dim;
}
let mut chunk_strides = vec![1usize; rank];
for i in (0..rank.saturating_sub(1)).rev() {
#[expect(
clippy::cast_possible_truncation,
reason = "chunk dimension derived from the in-memory write request; bounded by addressable memory"
)]
let dim = chunk_dims[i + 1] as usize;
chunk_strides[i] = chunk_strides[i + 1] * dim;
}
#[expect(
clippy::cast_possible_truncation,
reason = "chunk/dataset dimensions derived from the in-memory write request; bounded by addressable memory"
)]
let (chunk_dims_us, shape_us): (Vec<usize>, Vec<usize>) = (
chunk_dims.iter().map(|&d| d as usize).collect(),
shape.iter().map(|&d| d as usize).collect(),
);
let chunk_total_elements: usize = chunk_dims_us.iter().product();
#[expect(
clippy::cast_possible_truncation,
reason = "total_chunks derived from the in-memory write request; bounded by addressable memory"
)]
let mut buffers = Vec::with_capacity(total_chunks as usize);
let inner = rank - 1;
let mut coord = vec![0usize; inner];
let mut offsets_us = vec![0usize; rank];
let mut offsets = vec![0u64; rank];
for linear_idx in 0..total_chunks {
let mut remaining = linear_idx;
for d in (0..rank).rev() {
offsets[d] = (remaining % num_chunks_per_dim[d]) * chunk_dims[d];
remaining /= num_chunks_per_dim[d];
}
#[expect(
clippy::cast_possible_truncation,
reason = "chunk offset derived from the in-memory write request; bounded by addressable memory"
)]
for (slot, &o) in offsets_us.iter_mut().zip(offsets.iter()) {
*slot = o as usize;
}
let whole_chunk_in_bounds =
(0..rank).all(|d| offsets_us[d] + chunk_dims_us[d] <= shape_us[d]);
let mut chunk_bytes = vec![0u8; chunk_total_elements * element_size.get()];
if !whole_chunk_in_bounds {
fill.apply(&mut chunk_bytes)?;
}
let inner_row_len =
chunk_dims_us[inner].min(shape_us[inner].saturating_sub(offsets_us[inner]));
if inner_row_len > 0 {
let row_bytes = inner_row_len * element_size.get();
let inner_src = offsets_us[inner] * ds_strides[inner];
let outer_total: usize = chunk_dims_us[..inner].iter().product();
for c in coord.iter_mut() {
*c = 0;
}
for _ in 0..outer_total {
let mut dst_base = 0usize;
let mut src_base = inner_src;
let mut in_bounds = true;
for d in 0..inner {
dst_base += coord[d] * chunk_strides[d];
let global = offsets_us[d] + coord[d];
if global >= shape_us[d] {
in_bounds = false;
break;
}
src_base += global * ds_strides[d];
}
if in_bounds {
let src = src_base * element_size.get();
let dst = dst_base * element_size.get();
let mut avail = row_bytes.min(raw_data.len().saturating_sub(src));
avail -= avail % element_size;
if avail > 0 {
chunk_bytes[dst..dst + avail].copy_from_slice(&raw_data[src..src + avail]);
}
}
for d in (0..inner).rev() {
coord[d] += 1;
if coord[d] < chunk_dims_us[d] {
break;
}
coord[d] = 0;
}
}
}
buffers.push(chunk_bytes);
}
Ok(buffers)
}
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); buf.push(2);
let flags: u8 = if filtered_size.is_some() { 0x02 } else { 0x00 };
buf.push(flags);
#[expect(
clippy::cast_possible_truncation,
reason = "rank written into the 1-byte dimensionality field selected for this file"
)]
let ndims = chunk_dims.len() as u8 + 1;
buf.push(ndims);
let max_dim = chunk_dims
.iter()
.map(|&d| d as u64)
.chain(core::iter::once(element_size as u64))
.max()
.unwrap_or(1);
let dim_encoded_len: u8 = if max_dim <= 0xFF {
1
} else if max_dim <= 0xFFFF {
2
} else {
4
};
buf.push(dim_encoded_len);
for &d in chunk_dims {
#[expect(
clippy::cast_possible_truncation,
reason = "dimension written into the on-disk encoding width selected for this file"
)]
match dim_encoded_len {
1 => buf.push(d as u8),
2 => buf.extend_from_slice(&(d as u16).to_le_bytes()),
4 => buf.extend_from_slice(&d.to_le_bytes()),
_ => {}
}
}
#[expect(
clippy::cast_possible_truncation,
reason = "element size written into the on-disk encoding width selected for this file"
)]
match dim_encoded_len {
1 => buf.push(element_size as u8),
2 => buf.extend_from_slice(&(element_size as u16).to_le_bytes()),
4 => buf.extend_from_slice(&element_size.to_le_bytes()),
_ => {}
}
buf.push(1);
if let (Some(fs), Some(fm)) = (filtered_size, filter_mask) {
buf.extend_from_slice(&fs.to_le_bytes()); buf.extend_from_slice(&fm.to_le_bytes()); }
#[expect(
clippy::cast_possible_truncation,
reason = "chunk address written into the on-disk offset width selected for this file"
)]
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
}
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 = Vec::new();
buf.push(4); buf.push(2);
let flags: u8 = 0x00;
buf.push(flags);
#[expect(
clippy::cast_possible_truncation,
reason = "rank written into the 1-byte dimensionality field selected for this file"
)]
let ndims = chunk_dims.len() as u8 + 1;
buf.push(ndims);
let max_dim = chunk_dims
.iter()
.map(|&d| d as u64)
.chain(core::iter::once(element_size as u64))
.max()
.unwrap_or(1);
let dim_encoded_len: u8 = if max_dim <= 0xFF {
1
} else if max_dim <= 0xFFFF {
2
} else {
4
};
buf.push(dim_encoded_len);
for &d in chunk_dims {
#[expect(
clippy::cast_possible_truncation,
reason = "dimension written into the on-disk encoding width selected for this file"
)]
match dim_encoded_len {
1 => buf.push(d as u8),
2 => buf.extend_from_slice(&(d as u16).to_le_bytes()),
4 => buf.extend_from_slice(&d.to_le_bytes()),
_ => {}
}
}
#[expect(
clippy::cast_possible_truncation,
reason = "element size written into the on-disk encoding width selected for this file"
)]
match dim_encoded_len {
1 => buf.push(element_size as u8),
2 => buf.extend_from_slice(&(element_size as u16).to_le_bytes()),
4 => buf.extend_from_slice(&element_size.to_le_bytes()),
_ => {}
}
buf.push(3);
buf.push(max_bits);
#[expect(
clippy::cast_possible_truncation,
reason = "fixed array header address written into the on-disk offset width selected for this file"
)]
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
}
#[derive(Debug, Clone, Copy)]
struct ChunkElementEncoding {
chunk_size_bytes: usize,
elem_size: usize,
client_id: u8,
}
pub(crate) fn full_chunk_bytes(
chunk_dims: impl IntoIterator<Item = u64>,
element_size: NonZeroUsize,
) -> u64 {
chunk_dims
.into_iter()
.fold(element_size.get() as u64, |acc, d| acc.saturating_mul(d))
}
fn chunk_element_encoding(
chunk_bytes: u64,
offset_size: u8,
has_filters: bool,
) -> ChunkElementEncoding {
let os = offset_size as usize;
let chunk_size_bytes: usize = if has_filters {
let log2_val = if chunk_bytes <= 1 {
0
} else {
63 - chunk_bytes.leading_zeros()
};
let len = 1 + ((log2_val + 8) / 8) as usize;
len.min(8)
} else {
0
};
ChunkElementEncoding {
chunk_size_bytes,
elem_size: if has_filters {
os + chunk_size_bytes + 4
} else {
os
},
client_id: u8::from(has_filters),
}
}
struct FaLayout {
encoding: ChunkElementEncoding,
page_bits: u8,
page_size: usize,
fahd_size: usize,
total_len: u64,
}
fn fa_layout(
slots: &IndexSlots<'_>,
chunk_bytes: u64,
offset_size: u8,
length_size: u8,
has_filters: bool,
) -> FaLayout {
debug_assert!(
matches!(offset_size, 4 | 8) && matches!(length_size, 4 | 8),
"a fixed array is written at a 4- or 8-byte address and length width"
);
let os = offset_size as usize;
let num_elements = slots.len();
let encoding = chunk_element_encoding(chunk_bytes, offset_size, has_filters);
let fahd_size = 4 + 1 + 1 + 1 + 1 + length_size as usize + os + 4;
let fadb_prefix = 4 + 1 + 1 + os;
let page_bits = FIXED_ARRAY_PAGE_BITS;
let page_size = 1usize << page_bits;
let elements = num_elements * encoding.elem_size;
let fadb_size = if num_elements <= page_size {
fadb_prefix + elements + 4
} else {
let npages = num_elements.div_ceil(page_size);
fadb_prefix + npages.div_ceil(8) + 4 + elements + npages * 4
};
FaLayout {
encoding,
page_bits,
page_size,
fahd_size,
total_len: (fahd_size + fadb_size) as u64,
}
}
pub(crate) fn fixed_array_len(
slots: &IndexSlots<'_>,
chunk_bytes: u64,
offset_size: u8,
length_size: u8,
has_filters: bool,
) -> u64 {
fa_layout(slots, chunk_bytes, offset_size, length_size, has_filters).total_len
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum StorageAllocation {
Allocated,
Unallocated,
}
#[derive(Debug, Clone, Copy)]
pub(crate) enum ChunkIndexKind {
Unallocated,
SingleChunk,
FixedArray,
ExtensibleArray,
}
#[derive(Debug, Clone, Copy)]
pub(crate) enum ChunkArrayKind {
FixedArray,
ExtensibleArray,
}
impl ChunkIndexKind {
pub(crate) fn array_kind(self) -> Option<ChunkArrayKind> {
match self {
Self::Unallocated | Self::SingleChunk => None,
Self::FixedArray => Some(ChunkArrayKind::FixedArray),
Self::ExtensibleArray => Some(ChunkArrayKind::ExtensibleArray),
}
}
}
pub(crate) fn chunk_index_kind(grid: &ChunkGrid, num_chunks: usize) -> ChunkIndexKind {
match grid.slots() {
None => ChunkIndexKind::ExtensibleArray,
_ if num_chunks == 0 => ChunkIndexKind::Unallocated,
Some(1) => ChunkIndexKind::SingleChunk,
Some(_) => ChunkIndexKind::FixedArray,
}
}
pub(crate) fn chunk_index_len(
kind: ChunkArrayKind,
slots: &IndexSlots<'_>,
chunk_bytes: u64,
offset_size: u8,
length_size: u8,
has_filters: bool,
) -> u64 {
match kind {
ChunkArrayKind::ExtensibleArray => {
extensible_array_len(slots, chunk_bytes, offset_size, length_size, has_filters)
}
ChunkArrayKind::FixedArray => {
fixed_array_len(slots, chunk_bytes, offset_size, length_size, has_filters)
}
}
}
pub fn build_fixed_array_at(
slots: &IndexSlots<'_>,
chunk_bytes: u64,
offset_size: u8,
length_size: u8,
has_filters: bool,
fa_address: u64,
) -> Vec<u8> {
let num_elements = slots.len();
let layout = fa_layout(slots, chunk_bytes, offset_size, length_size, has_filters);
let ChunkElementEncoding {
chunk_size_bytes,
elem_size,
client_id,
} = layout.encoding;
let fahd_total_size = layout.fahd_size;
let fadb_address = fa_address + fahd_total_size as u64;
let mut fahd = Vec::with_capacity(fahd_total_size);
fahd.extend_from_slice(b"FAHD");
fahd.push(0); fahd.push(client_id);
#[expect(
clippy::cast_possible_truncation,
reason = "element record size written into the 1-byte FAHD field selected for this file"
)]
fahd.push(elem_size as u8);
fahd.push(layout.page_bits);
#[expect(
clippy::cast_possible_truncation,
reason = "element count written into the on-disk length width selected for this file"
)]
match length_size {
4 => fahd.extend_from_slice(&(num_elements as u32).to_le_bytes()),
8 => fahd.extend_from_slice(&(num_elements as u64).to_le_bytes()),
_ => fahd.extend_from_slice(&(num_elements as u64).to_le_bytes()),
}
#[expect(
clippy::cast_possible_truncation,
reason = "FADB address written into the on-disk offset width selected for this file"
)]
match offset_size {
4 => fahd.extend_from_slice(&(fadb_address as u32).to_le_bytes()),
8 => fahd.extend_from_slice(&fadb_address.to_le_bytes()),
_ => fahd.extend_from_slice(&fadb_address.to_le_bytes()),
}
let checksum = jenkins_lookup3(&fahd);
fahd.extend_from_slice(&checksum.to_le_bytes());
debug_assert_eq!(fahd.len(), fahd_total_size);
let write_element = |buf: &mut Vec<u8>, chunk: Option<&WrittenChunk>| {
let Some(chunk) = chunk else {
write_undefined_element(buf, offset_size, has_filters, chunk_size_bytes);
return;
};
#[expect(
clippy::cast_possible_truncation,
reason = "chunk address written into the on-disk offset width selected for this file"
)]
match offset_size {
4 => buf.extend_from_slice(&(chunk.address as u32).to_le_bytes()),
_ => buf.extend_from_slice(&chunk.address.to_le_bytes()),
}
if has_filters {
let cs_bytes = chunk.compressed_size.to_le_bytes();
buf.extend_from_slice(&cs_bytes[..chunk_size_bytes]);
buf.extend_from_slice(&chunk.filter_mask.to_le_bytes());
}
};
let mut fadb = Vec::new();
fadb.extend_from_slice(b"FADB");
fadb.push(0); fadb.push(client_id);
#[expect(
clippy::cast_possible_truncation,
reason = "fixed array header address written into the on-disk offset width selected for this file"
)]
match offset_size {
4 => fadb.extend_from_slice(&(fa_address as u32).to_le_bytes()),
_ => fadb.extend_from_slice(&fa_address.to_le_bytes()),
}
let page_size = layout.page_size;
if num_elements <= page_size {
for slot in 0..num_elements {
write_element(&mut fadb, slots.at(slot));
}
let fadb_checksum = jenkins_lookup3(&fadb);
fadb.extend_from_slice(&fadb_checksum.to_le_bytes());
} else {
let npages = num_elements.div_ceil(page_size);
let bitmap_size = npages.div_ceil(8);
let mut bitmap = vec![0u8; bitmap_size];
for page in 0..npages {
bitmap[page / 8] |= 1 << (7 - (page % 8));
}
fadb.extend_from_slice(&bitmap);
let prefix_checksum = jenkins_lookup3(&fadb);
fadb.extend_from_slice(&prefix_checksum.to_le_bytes());
for page in 0..npages {
let start = page * page_size;
let end = core::cmp::min(start + page_size, num_elements);
let mut page_buf = Vec::with_capacity((end - start) * elem_size);
for slot in start..end {
write_element(&mut page_buf, slots.at(slot));
}
let page_checksum = jenkins_lookup3(&page_buf);
page_buf.extend_from_slice(&page_checksum.to_le_bytes());
fadb.extend_from_slice(&page_buf);
}
}
let mut combined = fahd;
combined.extend_from_slice(&fadb);
debug_assert_eq!(
combined.len() as u64,
layout.total_len,
"a fixed array must fill the length its layout promised"
);
combined
}
pub(crate) fn serialize_v4_extensible_array(
chunk_dims: &[u32],
ea_address: u64,
offset_size: u8,
element_size: u32,
) -> Vec<u8> {
let mut buf = Vec::new();
buf.push(4); buf.push(2); buf.push(0x00);
#[expect(
clippy::cast_possible_truncation,
reason = "rank written into the 1-byte dimensionality field selected for this file"
)]
let ndims = chunk_dims.len() as u8 + 1;
buf.push(ndims);
let max_dim = chunk_dims
.iter()
.map(|&d| d as u64)
.chain(core::iter::once(element_size as u64))
.max()
.unwrap_or(1);
let dim_encoded_len: u8 = if max_dim <= 0xFF {
1
} else if max_dim <= 0xFFFF {
2
} else {
4
};
buf.push(dim_encoded_len);
for &d in chunk_dims {
#[expect(
clippy::cast_possible_truncation,
reason = "dimension written into the on-disk encoding width selected for this file"
)]
match dim_encoded_len {
1 => buf.push(d as u8),
2 => buf.extend_from_slice(&(d as u16).to_le_bytes()),
4 => buf.extend_from_slice(&d.to_le_bytes()),
_ => {}
}
}
#[expect(
clippy::cast_possible_truncation,
reason = "element size written into the on-disk encoding width selected for this file"
)]
match dim_encoded_len {
1 => buf.push(element_size as u8),
2 => buf.extend_from_slice(&(element_size as u16).to_le_bytes()),
4 => buf.extend_from_slice(&element_size.to_le_bytes()),
_ => {}
}
buf.push(4);
buf.push(32); buf.push(4); buf.push(4); buf.push(16); buf.push(10);
#[expect(
clippy::cast_possible_truncation,
reason = "extensible array header address written into the on-disk offset width selected for this file"
)]
match offset_size {
4 => buf.extend_from_slice(&(ea_address as u32).to_le_bytes()),
8 => buf.extend_from_slice(&ea_address.to_le_bytes()),
_ => {}
}
buf
}
pub(crate) fn write_ea_addr(buf: &mut Vec<u8>, val: u64, offset_size: u8) {
#[expect(
clippy::cast_possible_truncation,
reason = "address written into the on-disk offset width selected for this file"
)]
match offset_size {
4 => buf.extend_from_slice(&(val as u32).to_le_bytes()),
_ => buf.extend_from_slice(&val.to_le_bytes()),
}
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn build_eadb(
slots: &IndexSlots<'_>,
elem_start: usize,
dblk_nelmts: usize,
block_offset_rel: u64,
ea_address: u64,
offset_size: u8,
has_filters: bool,
chunk_size_bytes: usize,
client_id: u8,
page_nelmts: usize,
blk_off_size: usize,
) -> Vec<u8> {
let mut buf = Vec::new();
buf.extend_from_slice(b"EADB");
buf.push(0); buf.push(client_id);
write_ea_addr(&mut buf, ea_address, offset_size);
buf.extend_from_slice(&block_offset_rel.to_le_bytes()[..blk_off_size]);
let blocks = DataBlockGeom {
dblk_nelmts: dblk_nelmts as u64,
page_nelmts: page_nelmts as u64,
};
if !blocks.is_paged() {
for slot in 0..dblk_nelmts {
if let Some(chunk) = slots.at(elem_start + slot) {
write_chunk_element(&mut buf, chunk, offset_size, has_filters, chunk_size_bytes);
} else {
write_undefined_element(&mut buf, offset_size, has_filters, chunk_size_bytes);
}
}
let cks = jenkins_lookup3(&buf);
buf.extend_from_slice(&cks.to_le_bytes());
buf
} else {
let header_cks = jenkins_lookup3(&buf);
buf.extend_from_slice(&header_cks.to_le_bytes());
let npages = dblk_nelmts / page_nelmts;
for page in 0..npages {
let page_start = elem_start + page * page_nelmts;
let mut page_buf = Vec::new();
for slot in 0..page_nelmts {
if let Some(chunk) = slots.at(page_start + slot) {
write_chunk_element(
&mut page_buf,
chunk,
offset_size,
has_filters,
chunk_size_bytes,
);
} else {
write_undefined_element(
&mut page_buf,
offset_size,
has_filters,
chunk_size_bytes,
);
}
}
let page_cks = jenkins_lookup3(&page_buf);
page_buf.extend_from_slice(&page_cks.to_le_bytes());
buf.extend_from_slice(&page_buf);
}
buf
}
}
pub(crate) fn build_aesb(
ea_address: u64,
block_offset_rel: u64,
page_bitmap: &[u8],
dblk_addrs: &[u64],
offset_size: u8,
blk_off_size: usize,
client_id: u8,
) -> Vec<u8> {
let mut buf = Vec::new();
buf.extend_from_slice(b"EASB");
buf.push(0); buf.push(client_id);
write_ea_addr(&mut buf, ea_address, offset_size);
buf.extend_from_slice(&block_offset_rel.to_le_bytes()[..blk_off_size]);
buf.extend_from_slice(page_bitmap);
for &addr in dblk_addrs {
write_ea_addr(&mut buf, addr, offset_size);
}
let cks = jenkins_lookup3(&buf);
buf.extend_from_slice(&cks.to_le_bytes());
buf
}
pub(crate) fn aeib_size(
offset_size: u8,
inline_elmts: usize,
elem_size: usize,
ndblk_addrs: usize,
nsblk_addrs: usize,
) -> usize {
let os = offset_size as usize;
4 + 1 + 1 + os + inline_elmts * elem_size + ndblk_addrs * os + nsblk_addrs * os + 4 }
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct EaStats {
pub nsuper_blks: u64,
pub super_blk_size: u64,
pub ndata_blks: u64,
pub data_blk_size: u64,
pub max_idx_set: u64,
pub nelmts: u64,
}
pub(crate) fn eadb_size(
blocks: DataBlockGeom,
elem_size: usize,
offset_size: u8,
blk_off_size: usize,
) -> u64 {
let prefix = (4 + 1 + 1 + offset_size as u64) + blk_off_size as u64;
if blocks.is_paged() {
prefix + 4 + blocks.npages() * (blocks.page_nelmts * elem_size as u64 + 4)
} else {
prefix + blocks.dblk_nelmts * elem_size as u64 + 4
}
}
pub(crate) fn aesb_size(geom: SuperBlockGeom, offset_size: u8, blk_off_size: usize) -> u64 {
let os = offset_size as u64;
let header = 4 + 1 + 1 + os + blk_off_size as u64;
header + geom.bitmap_size() + geom.ndblks * os + 4
}
pub(crate) fn ea_compute_stats(
geom: &EaGeometry,
idx_blk_elmts: u64,
elem_size: usize,
page_nelmts: u64,
offset_size: u8,
blk_off_size: usize,
num_elements: u64,
occupancy: SlotOccupancy<'_>,
) -> EaStats {
let mut s = EaStats {
nsuper_blks: 0,
super_blk_size: 0,
ndata_blks: 0,
data_blk_size: 0,
max_idx_set: num_elements,
nelmts: idx_blk_elmts,
};
let mut elem = idx_blk_elmts;
for &dn in &geom.direct_dblk_nelmts {
if occupancy.any_occupied(elem, dn) {
let blocks = DataBlockGeom {
dblk_nelmts: dn,
page_nelmts,
};
s.ndata_blks += 1;
s.data_blk_size += eadb_size(blocks, elem_size, offset_size, blk_off_size);
s.nelmts += dn;
}
elem += dn;
}
for j in 0..geom.nsblk_addrs {
let sb = geom.super_block_at(j, page_nelmts);
let (ndblks, dn) = (sb.ndblks, sb.blocks.dblk_nelmts);
let span = ndblks * dn;
if occupancy.any_occupied(elem, span) {
s.nsuper_blks += 1;
s.super_blk_size += aesb_size(sb, offset_size, blk_off_size);
let mut le = elem;
for _ in 0..ndblks {
if occupancy.any_occupied(le, dn) {
s.ndata_blks += 1;
s.data_blk_size += eadb_size(sb.blocks, elem_size, offset_size, blk_off_size);
s.nelmts += dn;
}
le += dn;
}
}
elem += span;
}
s
}
struct EaLayout {
encoding: ChunkElementEncoding,
max_nelmts_bits: u8,
idx_blk_elmts: u8,
min_dblk_nelmts: u8,
super_blk_min_nelmts: u8,
max_dblk_nelmts_bits: u8,
geom: EaGeometry,
page_nelmts: usize,
blk_off_size: usize,
inline: usize,
aehd_size: usize,
aeib_size: usize,
stats: EaStats,
total_len: u64,
}
fn ea_layout(
occupancy: SlotOccupancy<'_>,
num_slots: u64,
chunk_bytes: u64,
offset_size: u8,
length_size: u8,
has_filters: bool,
) -> EaLayout {
let encoding = chunk_element_encoding(chunk_bytes, offset_size, has_filters);
let ChunkElementEncoding {
elem_size,
client_id,
..
} = encoding;
let (
max_nelmts_bits,
idx_blk_elmts,
min_dblk_nelmts,
super_blk_min_nelmts,
max_dblk_nelmts_bits,
) = (
EA_MAX_NELMTS_BITS,
EA_IDX_BLK_ELMTS,
EA_MIN_DBLK_NELMTS,
EA_SUPER_BLK_MIN_NELMTS,
EA_MAX_DBLK_NELMTS_BITS,
);
#[expect(
clippy::cast_possible_truncation,
reason = "element record size written into the 1-byte EA header field selected for this file"
)]
let geom_header = ExtensibleArrayHeader {
client_id,
element_size: elem_size as u8,
max_nelmts_bits,
idx_blk_elmts,
min_dblk_nelmts,
super_blk_min_nelmts,
max_dblk_nelmts_bits,
num_elements: 0,
index_block_address: 0,
};
let geom = EaGeometry::from_header(&geom_header);
let page_nelmts = 1usize << max_dblk_nelmts_bits;
let blk_off_size = (max_nelmts_bits as usize).div_ceil(8);
let inline = idx_blk_elmts as usize;
let aehd_size = ExtensibleArrayHeader::serialized_size(offset_size, length_size);
let aeib_size = aeib_size(
offset_size,
inline,
elem_size,
geom.direct_dblk_nelmts.len(),
geom.nsblk_addrs,
);
let stats = ea_compute_stats(
&geom,
idx_blk_elmts as u64,
elem_size,
page_nelmts as u64,
offset_size,
blk_off_size,
num_slots,
occupancy,
);
let total_len = (aehd_size + aeib_size) as u64 + stats.data_blk_size + stats.super_blk_size;
EaLayout {
encoding,
max_nelmts_bits,
idx_blk_elmts,
min_dblk_nelmts,
super_blk_min_nelmts,
max_dblk_nelmts_bits,
geom,
page_nelmts,
blk_off_size,
inline,
aehd_size,
aeib_size,
stats,
total_len,
}
}
pub(crate) fn extensible_array_len(
slots: &IndexSlots<'_>,
chunk_bytes: u64,
offset_size: u8,
length_size: u8,
has_filters: bool,
) -> u64 {
ea_layout(
SlotOccupancy::Sparse(slots),
slots.len() as u64,
chunk_bytes,
offset_size,
length_size,
has_filters,
)
.total_len
}
pub fn build_extensible_array_at(
slots: &IndexSlots<'_>,
chunk_bytes: u64,
offset_size: u8,
length_size: u8,
has_filters: bool,
ea_address: u64,
) -> Result<Vec<u8>, FormatError> {
let num_elements = slots.len();
let layout = ea_layout(
SlotOccupancy::Sparse(slots),
slots.len() as u64,
chunk_bytes,
offset_size,
length_size,
has_filters,
);
let ChunkElementEncoding {
chunk_size_bytes,
elem_size,
client_id,
} = layout.encoding;
let EaLayout {
max_nelmts_bits,
idx_blk_elmts,
min_dblk_nelmts,
super_blk_min_nelmts,
max_dblk_nelmts_bits,
ref geom,
page_nelmts,
blk_off_size,
inline,
aehd_size,
aeib_size,
..
} = layout;
let aeib_address = ea_address + aehd_size as u64;
let body_base = aeib_address + aeib_size as u64;
let undef_addr: u64 = match offset_size {
4 => 0xFFFF_FFFF,
_ => u64::MAX,
};
let mut body: Vec<u8> =
Vec::with_capacity((layout.stats.data_blk_size + layout.stats.super_blk_size).to_usize()?);
let mut direct_addrs: Vec<u64> = Vec::with_capacity(geom.direct_dblk_nelmts.len());
let mut sblk_addrs: Vec<u64> = Vec::with_capacity(geom.nsblk_addrs);
let mut ndata_blks: u64 = 0;
let mut data_blk_size: u64 = 0;
let mut nsuper_blks: u64 = 0;
let mut super_blk_size: u64 = 0;
let mut alloc_slots: u64 = inline as u64;
let mut elem_cursor: u64 = inline as u64;
let occupancy = SlotOccupancy::Sparse(slots);
for &dblk_nelmts in &geom.direct_dblk_nelmts {
if !occupancy.any_occupied(elem_cursor, dblk_nelmts) {
direct_addrs.push(undef_addr);
elem_cursor += dblk_nelmts;
continue;
}
let addr = body_base + body.len() as u64;
let db_bytes = build_eadb(
slots,
elem_cursor.to_usize()?,
dblk_nelmts.to_usize()?,
elem_cursor - inline as u64,
ea_address,
offset_size,
has_filters,
chunk_size_bytes,
client_id,
page_nelmts,
blk_off_size,
);
ndata_blks += 1;
data_blk_size += db_bytes.len() as u64;
alloc_slots += dblk_nelmts;
body.extend_from_slice(&db_bytes);
direct_addrs.push(addr);
elem_cursor += dblk_nelmts;
}
for j in 0..geom.nsblk_addrs {
let sblk_idx = geom.first_indirect_sblk + j;
let (ndblks, dblk_nelmts) = geom.sblks[sblk_idx];
let sb_span = ndblks * dblk_nelmts;
if !occupancy.any_occupied(elem_cursor, sb_span) {
sblk_addrs.push(undef_addr);
elem_cursor += sb_span;
continue;
}
let sb = geom.super_block_at(j, page_nelmts as u64);
let is_paged = sb.blocks.is_paged();
let npages = sb.blocks.npages();
let sb_block_offset = elem_cursor - inline as u64;
let mut page_bitmap = vec![0u8; sb.bitmap_size().to_usize()?];
let mut sb_dblk_addrs: Vec<u64> = Vec::with_capacity(ndblks.to_usize()?);
let mut local_elem = elem_cursor;
for db_local in 0..ndblks {
if !occupancy.any_occupied(local_elem, dblk_nelmts) {
sb_dblk_addrs.push(undef_addr);
local_elem += dblk_nelmts;
continue;
}
let addr = body_base + body.len() as u64;
let db_bytes = build_eadb(
slots,
local_elem.to_usize()?,
dblk_nelmts.to_usize()?,
local_elem - inline as u64,
ea_address,
offset_size,
has_filters,
chunk_size_bytes,
client_id,
page_nelmts,
blk_off_size,
);
ndata_blks += 1;
data_blk_size += db_bytes.len() as u64;
alloc_slots += dblk_nelmts;
body.extend_from_slice(&db_bytes);
sb_dblk_addrs.push(addr);
if is_paged {
let base = local_elem.to_usize()?;
for p in 0..npages.to_usize()? {
let page_start = base + p * page_nelmts;
if !(0..page_nelmts).any(|s| slots.at(page_start + s).is_some()) {
continue;
}
let global_page = (db_local * npages).to_usize()? + p;
page_bitmap[global_page / 8] |= 0x80 >> (global_page % 8);
}
}
local_elem += dblk_nelmts;
}
let aesb_addr = body_base + body.len() as u64;
let aesb = build_aesb(
ea_address,
sb_block_offset,
&page_bitmap,
&sb_dblk_addrs,
offset_size,
blk_off_size,
client_id,
);
nsuper_blks += 1;
super_blk_size += aesb.len() as u64;
body.extend_from_slice(&aesb);
sblk_addrs.push(aesb_addr);
elem_cursor += sb_span;
}
#[expect(
clippy::cast_possible_truncation,
reason = "statistic written into the on-disk length width selected for this file"
)]
let write_length = |buf: &mut Vec<u8>, val: u64| match length_size {
4 => buf.extend_from_slice(&(val as u32).to_le_bytes()),
_ => buf.extend_from_slice(&val.to_le_bytes()),
};
let mut aehd = Vec::with_capacity(aehd_size);
aehd.extend_from_slice(b"EAHD");
aehd.push(0); aehd.push(client_id);
#[expect(
clippy::cast_possible_truncation,
reason = "element record size written into the 1-byte EA header field selected for this file"
)]
aehd.push(elem_size as u8);
aehd.push(max_nelmts_bits);
aehd.push(idx_blk_elmts);
aehd.push(min_dblk_nelmts);
aehd.push(super_blk_min_nelmts);
aehd.push(max_dblk_nelmts_bits);
write_length(&mut aehd, nsuper_blks);
write_length(&mut aehd, super_blk_size);
write_length(&mut aehd, ndata_blks);
write_length(&mut aehd, data_blk_size);
write_length(&mut aehd, num_elements as u64); write_length(&mut aehd, alloc_slots);
write_ea_addr(&mut aehd, aeib_address, offset_size);
let aehd_checksum = jenkins_lookup3(&aehd);
aehd.extend_from_slice(&aehd_checksum.to_le_bytes());
debug_assert_eq!(aehd.len(), aehd_size);
let mut aeib = Vec::with_capacity(aeib_size);
aeib.extend_from_slice(b"EAIB");
aeib.push(0); aeib.push(client_id);
write_ea_addr(&mut aeib, ea_address, offset_size);
#[allow(clippy::needless_range_loop)]
for i in 0..inline {
if let Some(chunk) = slots.at(i) {
write_chunk_element(&mut aeib, chunk, offset_size, has_filters, chunk_size_bytes);
} else {
write_undefined_element(&mut aeib, offset_size, has_filters, chunk_size_bytes);
}
}
for &addr in &direct_addrs {
write_ea_addr(&mut aeib, addr, offset_size);
}
for &addr in &sblk_addrs {
write_ea_addr(&mut aeib, addr, offset_size);
}
let aeib_checksum = jenkins_lookup3(&aeib);
aeib.extend_from_slice(&aeib_checksum.to_le_bytes());
debug_assert_eq!(aeib.len(), aeib_size);
let mut combined = aehd;
combined.extend_from_slice(&aeib);
combined.extend_from_slice(&body);
debug_assert_eq!(
combined.len() as u64,
layout.total_len,
"an extensible array must fill the length its layout promised"
);
Ok(combined)
}
fn write_chunk_element(
buf: &mut Vec<u8>,
chunk: &WrittenChunk,
offset_size: u8,
has_filters: bool,
chunk_size_bytes: usize,
) {
#[expect(
clippy::cast_possible_truncation,
reason = "chunk address written into the on-disk offset width selected for this file"
)]
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.extend_from_slice(&chunk.address.to_le_bytes()),
}
if has_filters {
let cs_bytes = chunk.compressed_size.to_le_bytes();
buf.extend_from_slice(&cs_bytes[..chunk_size_bytes]);
buf.extend_from_slice(&chunk.filter_mask.to_le_bytes());
}
}
fn write_undefined_element(
buf: &mut Vec<u8>,
offset_size: u8,
has_filters: bool,
chunk_size_bytes: usize,
) {
let os = offset_size as usize;
buf.extend_from_slice(&vec![0xFF; os]);
if has_filters {
buf.extend_from_slice(&vec![0x00; chunk_size_bytes]);
buf.extend_from_slice(&0u32.to_le_bytes());
}
}
pub(crate) struct CompressedChunkSet {
compressed: Vec<Vec<u8>>,
chunk_dims_u32: Vec<u32>,
element_size: NonZeroUsize,
has_filters: bool,
kind: ChunkIndexKind,
slot_of_chunk: Vec<u64>,
index_slots: u64,
pipeline_message: Option<Vec<u8>>,
}
impl CompressedChunkSet {
fn index_slots<'a>(
&self,
written_chunks: &'a [WrittenChunk],
) -> Result<IndexSlots<'a>, FormatError> {
IndexSlots::new(written_chunks, &self.slot_of_chunk, self.index_slots)
}
fn full_chunk_bytes(&self) -> u64 {
full_chunk_bytes(
self.chunk_dims_u32.iter().map(|&d| u64::from(d)),
self.element_size,
)
}
}
pub(crate) fn compress_chunks(
raw_data: &[u8],
shape: &[u64],
ctx: ChunkContext<'_>,
options: &ChunkOptions,
maxshape: Option<&[u64]>,
fill: FillPattern<'_>,
allocation: StorageAllocation,
) -> Result<CompressedChunkSet, FormatError> {
let chunk_dims = ctx.chunk_dims;
let element_size = nonzero_usize_from(ctx.element_size)?;
let pipeline = options.build_pipeline(&ctx, fill)?;
let (kind, slot_of_chunk, index_slots) = plan_index_slots(
shape,
chunk_dims,
maxshape,
full_chunk_bytes(chunk_dims.iter().copied(), element_size),
pipeline.is_some(),
allocation,
)?;
let chunks = match allocation {
StorageAllocation::Allocated => {
split_into_chunks(raw_data, shape, chunk_dims, element_size, fill)?
}
StorageAllocation::Unallocated => Vec::new(),
};
let num_chunks = chunks.len();
let has_filters = pipeline.is_some();
debug_assert_eq!(slot_of_chunk.len(), num_chunks);
let mut compressed = Vec::with_capacity(num_chunks);
let mut scratch = crate::filters::FilterScratch::new();
for chunk_bytes in chunks {
let c = if let Some(ref pl) = pipeline {
compress_chunk_with(&mut scratch, &chunk_bytes, pl, ctx)?
} else {
chunk_bytes
};
compressed.push(c);
}
#[expect(
clippy::cast_possible_truncation,
reason = "chunk dimensions written into the on-disk u32 dimension fields selected for this file"
)]
let chunk_dims_u32: Vec<u32> = chunk_dims.iter().map(|&d| d as u32).collect();
Ok(CompressedChunkSet {
compressed,
chunk_dims_u32,
element_size,
has_filters,
kind,
slot_of_chunk,
index_slots,
pipeline_message: pipeline.as_ref().map(|pl| pl.serialize()),
})
}
pub(crate) fn plan_index_slots(
shape: &[u64],
chunk_dims: &[u64],
maxshape: Option<&[u64]>,
chunk_bytes: u64,
has_filters: bool,
allocation: StorageAllocation,
) -> Result<(ChunkIndexKind, Vec<u64>, u64), FormatError> {
let grid = index_grid(shape, chunk_dims, maxshape)?;
let counts: Vec<u64> = shape
.iter()
.zip(chunk_dims)
.map(|(d, c)| d.div_ceil(*c))
.collect();
let num_chunks = match allocation {
StorageAllocation::Allocated => counts.iter().product::<u64>().to_usize()?,
StorageAllocation::Unallocated => 0,
};
let kind = chunk_index_kind(&grid, num_chunks);
let mut slot_of_chunk = Vec::with_capacity(num_chunks);
let mut coords = vec![0u64; shape.len()];
for dense in 0..num_chunks {
let mut remaining = dense as u64;
for d in (0..shape.len()).rev() {
coords[d] = remaining % counts[d];
remaining /= counts[d];
}
slot_of_chunk.push(grid.slot_of(&coords)?);
}
let index_slots = match kind {
ChunkIndexKind::Unallocated | ChunkIndexKind::SingleChunk => 0,
ChunkIndexKind::FixedArray => grid.slots().ok_or_else(|| {
FormatError::ChunkedReadError(
"a Fixed Array cannot index an unlimited dimension".into(),
)
})?,
_ => slot_of_chunk.iter().max().map_or(0, |s| s + 1),
};
let scratch: Vec<u64> = if slot_of_chunk.windows(2).all(|w| w[0] <= w[1]) {
Vec::new()
} else {
let mut v = slot_of_chunk.clone();
v.sort_unstable();
v
};
let sorted_slots: &[u64] = if scratch.is_empty() {
&slot_of_chunk
} else {
&scratch
};
let encoding = chunk_element_encoding(chunk_bytes, INDEX_OFFSET_SIZE, has_filters);
let allocated = match kind {
ChunkIndexKind::Unallocated | ChunkIndexKind::SingleChunk => 0,
ChunkIndexKind::FixedArray => index_slots,
ChunkIndexKind::ExtensibleArray => {
let capacity = ea_addressable_slots();
if index_slots > capacity {
return Err(FormatError::ChunkedReadError(format!(
"this shape and maximum shape number a chunk at element {} of the chunk \
index, past the {capacity} an extensible array can address; the chunk would \
be dropped. Chunk the dimensions the dataset does not grow along more \
coarsely, or give them a smaller maximum",
index_slots - 1,
)));
}
ea_layout(
SlotOccupancy::Slots(sorted_slots),
index_slots,
chunk_bytes,
INDEX_OFFSET_SIZE,
INDEX_LENGTH_SIZE,
has_filters,
)
.stats
.nelmts
}
};
let unused_bytes = allocated
.saturating_sub(num_chunks as u64)
.saturating_mul(encoding.elem_size as u64);
if unused_bytes > MAX_UNUSED_INDEX_BYTES {
return Err(FormatError::ChunkedReadError(format!(
"this shape and maximum shape need a chunk index holding {allocated} elements for \
{num_chunks} chunk(s), so {unused_bytes} bytes of it describe no chunk, past the \
{MAX_UNUSED_INDEX_BYTES} this writer will emit. The unused elements come from the \
maximum shape exceeding the shape in a dimension other than the one the dataset \
grows along — chunk those dimensions more coarsely, or declare no maximum for them"
)));
}
Ok((kind, slot_of_chunk, index_slots))
}
const EA_MAX_NELMTS_BITS: u8 = 32;
const EA_IDX_BLK_ELMTS: u8 = 4;
const EA_MIN_DBLK_NELMTS: u8 = 16;
const EA_SUPER_BLK_MIN_NELMTS: u8 = 4;
const EA_MAX_DBLK_NELMTS_BITS: u8 = 10;
fn ea_addressable_slots() -> u64 {
let geom_header = ExtensibleArrayHeader {
client_id: 0,
element_size: 8,
max_nelmts_bits: EA_MAX_NELMTS_BITS,
idx_blk_elmts: EA_IDX_BLK_ELMTS,
min_dblk_nelmts: EA_MIN_DBLK_NELMTS,
super_blk_min_nelmts: EA_SUPER_BLK_MIN_NELMTS,
max_dblk_nelmts_bits: EA_MAX_DBLK_NELMTS_BITS,
num_elements: 0,
index_block_address: 0,
};
let geom = EaGeometry::from_header(&geom_header);
let direct: u64 = geom.direct_dblk_nelmts.iter().sum();
let indirect: u64 = (0..geom.nsblk_addrs)
.map(|j| {
let (ndblks, dn) = geom.sblks[geom.first_indirect_sblk + j];
ndblks * dn
})
.sum();
u64::from(EA_IDX_BLK_ELMTS) + direct + indirect
}
fn index_grid(
shape: &[u64],
chunk_dims: &[u64],
maxshape: Option<&[u64]>,
) -> Result<ChunkGrid, FormatError> {
let order = if maxshape.is_some_and(|ms| ms.contains(&u64::MAX)) {
GridOrder::UnlimitedFirst
} else {
GridOrder::RowMajor
};
ChunkGrid::new(chunk_dims, shape, maxshape, order)
}
fn plan_chunk_slots(set: &CompressedChunkSet, data_address: u64) -> (Vec<WrittenChunk>, u64) {
let mut cursor = data_address;
let mut written_chunks = Vec::with_capacity(set.compressed.len());
for chunk in &set.compressed {
written_chunks.push(WrittenChunk {
address: cursor,
compressed_size: chunk.len() as u64,
filter_mask: 0,
});
cursor += chunk.len() as u64;
}
(written_chunks, cursor)
}
fn chunk_index_bytes(
set: &CompressedChunkSet,
written_chunks: &[WrittenChunk],
slots: &IndexSlots<'_>,
index_address: u64,
) -> Result<(Vec<u8>, Vec<u8>), FormatError> {
let index = match set.kind {
ChunkIndexKind::Unallocated => Vec::new(),
ChunkIndexKind::ExtensibleArray => build_extensible_array_at(
slots,
set.full_chunk_bytes(),
INDEX_OFFSET_SIZE,
INDEX_LENGTH_SIZE,
set.has_filters,
index_address,
)?,
ChunkIndexKind::SingleChunk => Vec::new(),
ChunkIndexKind::FixedArray => build_fixed_array_at(
slots,
set.full_chunk_bytes(),
INDEX_OFFSET_SIZE,
INDEX_LENGTH_SIZE,
set.has_filters,
index_address,
),
};
Ok((
index,
chunk_index_layout(set, written_chunks, index_address),
))
}
fn chunk_index_layout(
set: &CompressedChunkSet,
written_chunks: &[WrittenChunk],
index_address: u64,
) -> Vec<u8> {
let has_filters = set.has_filters;
#[expect(
clippy::cast_possible_truncation,
reason = "element size written into the on-disk u32 dimension field selected for this file"
)]
match set.kind {
ChunkIndexKind::Unallocated => serialize_v4_fixed_array(
&set.chunk_dims_u32,
HADDR_UNDEF,
INDEX_OFFSET_SIZE,
set.element_size.get() as u32,
FIXED_ARRAY_PAGE_BITS,
),
ChunkIndexKind::ExtensibleArray => serialize_v4_extensible_array(
&set.chunk_dims_u32,
index_address,
INDEX_OFFSET_SIZE,
set.element_size.get() as u32,
),
ChunkIndexKind::SingleChunk => {
let chunk = &written_chunks[0];
serialize_v4_single_chunk(
&set.chunk_dims_u32,
chunk.address,
has_filters.then_some(chunk.compressed_size),
has_filters.then_some(0u32),
INDEX_OFFSET_SIZE,
set.element_size.get() as u32,
)
}
ChunkIndexKind::FixedArray => serialize_v4_fixed_array(
&set.chunk_dims_u32,
index_address,
INDEX_OFFSET_SIZE,
set.element_size.get() as u32,
FIXED_ARRAY_PAGE_BITS,
),
}
}
pub(crate) fn chunked_data_len(set: &CompressedChunkSet) -> Result<u64, FormatError> {
let (written_chunks, index_address) = plan_chunk_slots(set, 0);
let slots = set.index_slots(&written_chunks)?;
let kind = set.kind;
Ok(index_address
+ kind.array_kind().map_or(0, |array| {
chunk_index_len(
array,
&slots,
set.full_chunk_bytes(),
INDEX_OFFSET_SIZE,
INDEX_LENGTH_SIZE,
set.has_filters,
)
}))
}
fn plan_chunked_at(
set: &CompressedChunkSet,
data_address: u64,
) -> Result<(usize, Vec<u8>, Vec<u8>), FormatError> {
let (written_chunks, index_address) = plan_chunk_slots(set, data_address);
let slots = set.index_slots(&written_chunks)?;
let (index, layout_message) = chunk_index_bytes(set, &written_chunks, &slots, index_address)?;
let chunk_bytes_total: usize = set.compressed.iter().map(Vec::len).sum();
Ok((chunk_bytes_total, index, layout_message))
}
pub(crate) fn measure_chunked_at(
set: &CompressedChunkSet,
data_address: u64,
) -> Result<ChunkedMeasure, FormatError> {
let (written_chunks, index_address) = plan_chunk_slots(set, data_address);
let slots = set.index_slots(&written_chunks)?;
let index_len = set.kind.array_kind().map_or(0, |array| {
chunk_index_len(
array,
&slots,
set.full_chunk_bytes(),
INDEX_OFFSET_SIZE,
INDEX_LENGTH_SIZE,
set.has_filters,
)
});
let data_len = (index_address - data_address) + index_len;
data_len.to_usize()?;
Ok(ChunkedMeasure {
data_len,
layout_message: chunk_index_layout(set, &written_chunks, index_address),
pipeline_message: set.pipeline_message.clone(),
})
}
pub(crate) struct ChunkedMeasure {
pub data_len: u64,
pub layout_message: Vec<u8>,
pub pipeline_message: Option<Vec<u8>>,
}
pub(crate) fn assemble_chunked_at(
set: &CompressedChunkSet,
data_address: u64,
) -> Result<ChunkedDataResult, FormatError> {
let (chunk_bytes_total, index, layout_message) = plan_chunked_at(set, data_address)?;
let mut data_buf = Vec::with_capacity(chunk_bytes_total + index.len());
for chunk in &set.compressed {
data_buf.extend_from_slice(chunk);
}
data_buf.extend_from_slice(&index);
Ok(ChunkedDataResult {
data_bytes: data_buf,
layout_message,
pipeline_message: set.pipeline_message.clone(),
})
}
#[cfg(test)]
pub fn build_chunked_data_at_ext(
raw_data: &[u8],
shape: &[u64],
ctx: ChunkContext<'_>,
options: &ChunkOptions,
data_address: u64,
maxshape: Option<&[u64]>,
fill: FillPattern<'_>,
) -> Result<ChunkedDataResult, FormatError> {
let set = compress_chunks(
raw_data,
shape,
ctx,
options,
maxshape,
fill,
StorageAllocation::Allocated,
)?;
assemble_chunked_at(&set, data_address)
}
#[derive(Debug, Clone)]
pub(crate) struct ChunkMeta {
pub(crate) compressed_size: u64,
pub(crate) filter_mask: u32,
}
pub(crate) trait ChunkProvider: Send + Sync {
fn chunk_bytes(&self, index: usize, out: &mut Vec<u8>) -> Result<(), FormatError>;
}
pub(crate) trait ByteSink {
fn put(&mut self, bytes: &[u8]) -> Result<(), FormatError>;
fn put_zeros(&mut self, n: usize) -> Result<(), FormatError>;
fn position(&self) -> u64;
fn reserve(&mut self, _additional: usize) {}
}
impl ByteSink for Vec<u8> {
fn put(&mut self, bytes: &[u8]) -> Result<(), FormatError> {
self.extend_from_slice(bytes);
Ok(())
}
fn put_zeros(&mut self, n: usize) -> Result<(), FormatError> {
self.resize(self.len() + n, 0u8);
Ok(())
}
fn position(&self) -> u64 {
self.len() as u64
}
fn reserve(&mut self, additional: usize) {
Vec::reserve(self, additional);
}
}
struct VerbatimIndexPlan {
kind: ChunkArrayKind,
address: u64,
has_filters: bool,
chunk_bytes: u64,
len: u64,
}
pub(crate) struct VerbatimPlan {
pub(crate) chunks: Vec<WrittenChunk>,
index: Option<VerbatimIndexPlan>,
slot_of_chunk: Vec<u64>,
index_slots: u64,
pub(crate) total_len: u64,
}
pub(crate) struct VerbatimLayout {
pub(crate) plan: VerbatimPlan,
pub(crate) layout_message: Vec<u8>,
pub(crate) pipeline_message: Option<Vec<u8>>,
}
pub(crate) fn plan_chunked_data_verbatim(
meta: &[ChunkMeta],
shape: &[u64],
chunk_dims: &[u64],
element_size: NonZeroUsize,
pipeline_message: Option<&[u8]>,
data_address: u64,
maxshape: Option<&[u64]>,
) -> Result<VerbatimLayout, FormatError> {
if meta.is_empty() {
return Err(FormatError::ChunkedReadError(
"a verbatim chunked dataset requires at least one chunk".into(),
));
}
let num_chunks = meta.len();
let has_filters = pipeline_message.is_some();
let mut cursor: u64 = 0;
let mut written_chunks = Vec::with_capacity(num_chunks);
for m in meta {
let address = data_address + cursor;
let compressed_size = m.compressed_size;
written_chunks.push(WrittenChunk {
address,
compressed_size,
filter_mask: m.filter_mask,
});
cursor += compressed_size;
}
#[expect(
clippy::cast_possible_truncation,
reason = "chunk dimensions written into the on-disk u32 dimension fields selected for this file"
)]
let chunk_dims_u32: Vec<u32> = chunk_dims.iter().map(|&d| d as u32).collect();
let offset_size = INDEX_OFFSET_SIZE;
let length_size = INDEX_LENGTH_SIZE;
let (kind, slot_of_chunk, index_slots) = plan_index_slots(
shape,
chunk_dims,
maxshape,
full_chunk_bytes(chunk_dims.iter().copied(), element_size),
has_filters,
StorageAllocation::Allocated,
)?;
if slot_of_chunk.len() != num_chunks {
return Err(FormatError::ChunkedReadError(format!(
"a verbatim chunked dataset of shape {shape:?} holds {} chunks, not the \
{num_chunks} it was given",
slot_of_chunk.len(),
)));
}
let slots = IndexSlots::new(&written_chunks, &slot_of_chunk, index_slots)?;
let index_address = data_address + cursor;
let chunk_bytes = full_chunk_bytes(chunk_dims.iter().copied(), element_size);
let index = kind.array_kind().map(|array| VerbatimIndexPlan {
kind: array,
address: index_address,
has_filters,
chunk_bytes,
len: chunk_index_len(
array,
&slots,
chunk_bytes,
offset_size,
length_size,
has_filters,
),
});
cursor += index.as_ref().map_or(0, |i| i.len);
#[expect(
clippy::cast_possible_truncation,
reason = "element size written into the on-disk u32 dimension field selected for this file"
)]
let layout_message = match kind {
ChunkIndexKind::Unallocated => serialize_v4_fixed_array(
&chunk_dims_u32,
HADDR_UNDEF,
offset_size,
element_size.get() as u32,
FIXED_ARRAY_PAGE_BITS,
),
ChunkIndexKind::ExtensibleArray => serialize_v4_extensible_array(
&chunk_dims_u32,
index_address,
offset_size,
element_size.get() as u32,
),
ChunkIndexKind::SingleChunk => {
let chunk_addr = written_chunks[0].address;
let filtered_size = if has_filters {
Some(written_chunks[0].compressed_size)
} else {
None
};
let filter_mask = if has_filters {
Some(written_chunks[0].filter_mask)
} else {
None
};
serialize_v4_single_chunk(
&chunk_dims_u32,
chunk_addr,
filtered_size,
filter_mask,
offset_size,
element_size.get() as u32,
)
}
ChunkIndexKind::FixedArray => serialize_v4_fixed_array(
&chunk_dims_u32,
index_address,
offset_size,
element_size.get() as u32,
FIXED_ARRAY_PAGE_BITS,
),
};
Ok(VerbatimLayout {
plan: VerbatimPlan {
chunks: written_chunks,
slot_of_chunk,
index_slots,
index,
total_len: cursor,
},
layout_message,
pipeline_message: pipeline_message.map(<[u8]>::to_vec),
})
}
pub(crate) fn emit_chunked_data_verbatim<S: ByteSink>(
sink: &mut S,
plan: &VerbatimPlan,
provider: &dyn ChunkProvider,
) -> Result<(), FormatError> {
let mut chunk = Vec::new();
for (i, slot) in plan.chunks.iter().enumerate() {
chunk.clear();
provider.chunk_bytes(i, &mut chunk)?;
if chunk.len() as u64 != slot.compressed_size {
return Err(FormatError::ChunkedReadError(
"verbatim chunk provider returned a chunk whose size differs from the \
planned size"
.into(),
));
}
sink.put(&chunk)?;
}
if let Some(index) = &plan.index {
let slots = IndexSlots::new(&plan.chunks, &plan.slot_of_chunk, plan.index_slots)?;
let bytes = match index.kind {
ChunkArrayKind::ExtensibleArray => build_extensible_array_at(
&slots,
index.chunk_bytes,
INDEX_OFFSET_SIZE,
INDEX_LENGTH_SIZE,
index.has_filters,
index.address,
)?,
ChunkArrayKind::FixedArray => build_fixed_array_at(
&slots,
index.chunk_bytes,
INDEX_OFFSET_SIZE,
INDEX_LENGTH_SIZE,
index.has_filters,
index.address,
),
};
if bytes.len() as u64 != index.len {
return Err(FormatError::SerializationError(format!(
"a chunk index built {} bytes where its plan reserved {}; the data region's \
length was computed from the plan",
bytes.len(),
index.len,
)));
}
sink.put(&bytes)?;
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::chunk_cache::ChunkCache;
use crate::chunked_read::read_chunked_data_cached;
use crate::convert::nz;
use crate::data_layout::DataLayout;
use crate::dataspace::{Dataspace, DataspaceType};
use crate::datatype::{Datatype, DatatypeByteOrder};
use crate::fill_value::FillPattern;
use crate::read_spec::RawReadSpec;
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,
}
}
#[test]
fn an_unknown_fill_refuses_only_the_chunks_that_need_padding() {
let elem = nz(4);
let data = vec![0u8; 8 * 4];
assert!(
split_into_chunks(&data, &[8], &[4], elem, FillPattern::UNKNOWN).is_ok(),
"a write that pads nothing must not consult the fill value"
);
assert!(matches!(
split_into_chunks(&data[..5 * 4], &[5], &[4], elem, FillPattern::UNKNOWN),
Err(FormatError::UnreadableFillValue)
));
assert!(matches!(
split_into_chunks(
&data[..3 * 2 * 4],
&[3, 2],
&[2, 2],
elem,
FillPattern::UNKNOWN
),
Err(FormatError::UnreadableFillValue)
));
for pattern in [FillPattern::ZERO, FillPattern::new(Some(&[7u8; 4]), elem)] {
assert!(split_into_chunks(&data[..5 * 4], &[5], &[4], elem, pattern).is_ok());
}
}
#[test]
fn every_split_chunk_is_a_whole_chunk() {
let cases: &[(&[u64], &[u64])] = &[
(&[8], &[4]), (&[7], &[4]), (&[1], &[512]), (&[4, 6], &[2, 3]), (&[5, 6], &[2, 3]), (&[4, 7], &[2, 3]), (&[5, 7], &[2, 3]), (&[3, 5, 7], &[2, 2, 4]), ];
for &(shape, chunk_dims) in cases {
for elem in [1usize, 4, 8] {
let n: u64 = shape.iter().product();
let raw = vec![0u8; (n as usize) * elem];
let expected = full_chunk_bytes(chunk_dims.iter().copied(), nz(elem));
let chunks =
split_into_chunks(&raw, shape, chunk_dims, nz(elem), FillPattern::ZERO)
.unwrap();
assert!(
!chunks.is_empty(),
"shape {shape:?} chunk {chunk_dims:?} must produce chunks"
);
for (i, c) in chunks.iter().enumerate() {
assert_eq!(
c.len() as u64,
expected,
"chunk {i} of shape {shape:?} chunk {chunk_dims:?} elem {elem} \
is not a whole chunk"
);
}
}
}
}
#[test]
fn an_unallocated_dataset_plans_no_chunks_over_the_shape_it_declares() {
let shape = [1000u64];
let chunk = [100u64];
let (kind, slots, span) = plan_index_slots(
&shape,
&chunk,
None,
400,
false,
StorageAllocation::Allocated,
)
.expect("a dense fixed-shape plan");
assert!(matches!(kind, ChunkIndexKind::FixedArray));
assert_eq!(slots.len(), 10, "ten chunks cover the shape");
assert_eq!(span, 10);
let (kind, slots, span) = plan_index_slots(
&shape,
&chunk,
None,
400,
false,
StorageAllocation::Unallocated,
)
.expect("an unallocated plan");
assert!(
matches!(kind, ChunkIndexKind::Unallocated),
"a fixed-shape dataset that stores nothing carries the undefined \
address, not an index over an empty grid: {kind:?}"
);
assert!(slots.is_empty(), "no chunk has a slot");
assert_eq!(span, 0, "and the index spans none");
let (kind, slots, _) = plan_index_slots(
&shape,
&chunk,
Some(&[u64::MAX]),
400,
false,
StorageAllocation::Unallocated,
)
.expect("an unallocated resizable plan");
assert!(matches!(kind, ChunkIndexKind::ExtensibleArray), "{kind:?}");
assert!(slots.is_empty());
}
#[test]
fn an_unallocated_dataset_encodes_no_chunk_bytes() {
let shape = [1000u64];
let chunk = [100u64];
let fill = [7u8, 0, 0, 0];
let elem = NonZeroUsize::new(4).unwrap();
let set = compress_chunks(
&[],
&shape,
ChunkContext::basic(&chunk, 4),
&ChunkOptions::default(),
None,
FillPattern::new(Some(&fill), elem),
StorageAllocation::Unallocated,
)
.expect("an unallocated set");
assert_eq!(set.compressed.len(), 0, "no chunk was encoded");
assert_eq!(
chunked_data_len(&set).unwrap(),
0,
"and the dataset occupies no data region"
);
let assembled = assemble_chunked_at(&set, 0x1000).unwrap();
assert!(assembled.data_bytes.is_empty(), "nothing to write out");
assert!(
assembled
.layout_message
.windows(8)
.any(|w| w == u64::MAX.to_le_bytes()),
"the layout message must carry the undefined address: {:?}",
assembled.layout_message
);
let raw = vec![0u8; 1000 * 4];
let dense = compress_chunks(
&raw,
&shape,
ChunkContext::basic(&chunk, 4),
&ChunkOptions::default(),
None,
FillPattern::new(Some(&fill), elem),
StorageAllocation::Allocated,
)
.expect("a dense set");
assert_eq!(dense.compressed.len(), 10);
assert!(chunked_data_len(&dense).unwrap() >= 4000);
}
#[test]
fn the_index_bound_counts_unused_bytes_rather_than_slots() {
let many = (MAX_UNUSED_INDEX_BYTES / 8 + 1) as usize;
let (kind, slots, span) = plan_index_slots(
&[many as u64],
&[1],
None,
8,
false,
StorageAllocation::Allocated,
)
.expect("a dense index is not bounded");
assert!(matches!(kind, ChunkIndexKind::FixedArray));
assert_eq!(slots.len(), many);
assert_eq!(span, many as u64);
let chunk_bytes = 4096;
for has_filters in [false, true] {
let elem = u64::from(
chunk_element_encoding(chunk_bytes, INDEX_OFFSET_SIZE, has_filters).elem_size
as u32,
);
let widest = MAX_UNUSED_INDEX_BYTES / elem + 1;
plan_index_slots(
&[1],
&[1],
Some(&[widest]),
chunk_bytes,
has_filters,
StorageAllocation::Allocated,
)
.expect("an index exactly at the budget is written");
let err = plan_index_slots(
&[1],
&[1],
Some(&[widest + 1]),
chunk_bytes,
has_filters,
StorageAllocation::Allocated,
)
.unwrap_err();
assert!(
format!("{err}").contains("describe no chunk"),
"filtered={has_filters}: {err}"
);
}
plan_index_slots(
&[8, 8],
&[4, 4],
Some(&[u64::MAX, 8]),
64,
false,
StorageAllocation::Allocated,
)
.expect("growth along the indexed dimension leaves no gaps");
}
#[test]
fn an_extensible_array_refuses_a_chunk_past_the_slots_it_can_address() {
let capacity = ea_addressable_slots();
assert_eq!(
capacity, 8_589_934_580,
"the C library's default creation parameters address this many slots"
);
let at = |stride: u64| {
plan_index_slots(
&[2, 1],
&[1, 1],
Some(&[u64::MAX, stride]),
4,
false,
StorageAllocation::Allocated,
)
};
at(capacity - 1).expect("the last addressable slot is written");
let err = at(capacity).unwrap_err();
assert!(format!("{err}").contains("can address"), "{err}");
}
#[test]
fn auto_chunking_a_zero_element_shape_is_refused() {
let auto = ChunkOptions {
chunk_dims: None,
..Default::default()
};
for shape in [vec![0u64], vec![4, 0], vec![0, 4]] {
let err = auto
.validate_geometry(&shape, Some(&vec![u64::MAX; shape.len()]))
.unwrap_err();
assert!(
err.contains("explicit chunk dimensions"),
"shape {shape:?}: {err}"
);
assert!(auto.resolve_chunk_dims(&shape).contains(&0));
}
let explicit = ChunkOptions {
chunk_dims: Some(vec![512]),
..Default::default()
};
assert!(explicit.validate_geometry(&[0], Some(&[u64::MAX])).is_ok());
assert!(explicit.validate_geometry(&[0], None).is_ok());
}
fn filtered(filters: &[FilterKind]) -> ChunkOptions {
let mut options = ChunkOptions::default();
for &f in filters {
options.set_filter(f);
}
options
}
fn chunked(chunk_dims: &[u64], filters: &[FilterKind]) -> ChunkOptions {
ChunkOptions {
chunk_dims: Some(chunk_dims.to_vec()),
..filtered(filters)
}
}
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()
}
#[test]
fn measuring_a_chunked_region_agrees_with_assembling_it() {
type Case = (&'static [u64], &'static [u64], Option<&'static [u64]>);
let cases: [Case; 4] = [
(&[512], &[512], None),
(&[4096], &[512], None),
(&[4096], &[64], None),
(&[4096], &[512], Some(&[u64::MAX])),
];
for (shape, chunk_dims, maxshape) in cases {
let elems: usize = shape.iter().product::<u64>().to_usize().unwrap();
let raw = f64_to_bytes(&(0..elems).map(|i| i as f64).collect::<Vec<f64>>());
let ctx = ChunkContext::basic(chunk_dims, 8);
let set = compress_chunks(
&raw,
shape,
ctx,
&ChunkOptions::default(),
maxshape,
FillPattern::ZERO,
StorageAllocation::Allocated,
)
.unwrap();
for base in [0u64, 0x1000, 0x1234_5678] {
let measured = measure_chunked_at(&set, base).unwrap();
let assembled = assemble_chunked_at(&set, base).unwrap();
assert_eq!(
measured.data_len,
assembled.data_bytes.len() as u64,
"measured and assembled lengths differ for shape {shape:?} in \
chunks {chunk_dims:?} at {base:#x}"
);
assert_eq!(
measured.layout_message, assembled.layout_message,
"measured and assembled layout messages differ for shape \
{shape:?} in chunks {chunk_dims:?} at {base:#x}"
);
}
}
}
fn roundtrip_chunked(
values: &[f64],
shape: &[u64],
chunk_dims: &[u64],
options: &ChunkOptions,
) -> Vec<f64> {
let raw = f64_to_bytes(values);
let data_address = 0x1000u64;
let ctx = ChunkContext::basic(chunk_dims, 8);
let result = build_chunked_data_at_ext(
&raw,
shape,
ctx,
options,
data_address,
None,
FillPattern::ZERO,
)
.unwrap();
let file_size = data_address as usize + result.data_bytes.len();
let mut file_data = vec![0u8; file_size];
file_data[data_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: shape.len() as u8,
dimensions: shape.to_vec(),
max_dimensions: None,
};
let datatype = make_f64_type();
let pipeline = result
.pipeline_message
.as_ref()
.map(|pm| crate::filter_pipeline::FilterPipeline::parse(pm).unwrap());
let output = read_chunked_data_cached(
&file_data,
RawReadSpec {
layout: &layout,
dataspace: &dataspace,
datatype: &datatype,
pipeline: pipeline.as_ref(),
fill: FillPattern::ZERO,
},
8,
8,
&ChunkCache::new(),
)
.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], nz(8), FillPattern::ZERO).unwrap();
assert_eq!(result.len(), 1);
assert_eq!(bytes_to_f64(&result[0]), 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], nz(8), FillPattern::ZERO).unwrap();
assert_eq!(result.len(), 3); assert_eq!(bytes_to_f64(&result[0]), vec![0.0, 1.0, 2.0, 3.0]);
assert_eq!(bytes_to_f64(&result[1]), vec![4.0, 5.0, 6.0, 7.0]);
assert_eq!(bytes_to_f64(&result[2]), vec![8.0, 9.0, 0.0, 0.0]);
}
#[test]
fn split_2d_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], nz(8), FillPattern::ZERO).unwrap();
assert_eq!(result.len(), 4);
assert_eq!(bytes_to_f64(&result[0]), vec![0.0, 1.0, 4.0, 5.0]);
assert_eq!(bytes_to_f64(&result[1]), vec![2.0, 3.0, 6.0, 7.0]);
assert_eq!(bytes_to_f64(&result[2]), vec![8.0, 9.0, 12.0, 13.0]);
assert_eq!(bytes_to_f64(&result[3]), vec![10.0, 11.0, 14.0, 15.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 = chunked(&[100], &[FilterKind::Deflate(6)]);
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 = chunked(&[20], &[FilterKind::Deflate(6)]);
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 = chunked(&[50], &[FilterKind::Shuffle, FilterKind::Deflate(6)]);
let result = roundtrip_chunked(&values, &[100], &[50], &options);
assert_eq!(result, values);
}
#[test]
fn roundtrip_2d_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 chunks_are_stored_back_to_back() {
let values: Vec<f64> = (0..21).map(|i| i as f64).collect();
let raw = f64_to_bytes(&values);
let options = ChunkOptions {
chunk_dims: Some(vec![7]),
..Default::default()
};
let dims = [7u64];
let ctx = ChunkContext::basic(&dims, 8);
let result =
build_chunked_data_at_ext(&raw, &[21], ctx, &options, 0x1000, None, FillPattern::ZERO)
.unwrap();
assert_eq!(
&result.data_bytes[..raw.len()],
&raw[..],
"the three chunks must concatenate with nothing between them"
);
assert_eq!(
&result.data_bytes[raw.len()..raw.len() + 4],
b"FAHD",
"the chunk index must begin where the last chunk ends"
);
}
#[test]
fn a_verbatim_plan_reserves_only_the_chunks_and_the_index() {
let meta: Vec<ChunkMeta> = [37u64, 111, 5]
.into_iter()
.map(|compressed_size| ChunkMeta {
compressed_size,
filter_mask: 0,
})
.collect();
let layout =
plan_chunked_data_verbatim(&meta, &[21], &[7], nz(8), Some(&[]), 0x1000, None).unwrap();
let planned: Vec<u64> = layout
.plan
.chunks
.iter()
.map(|c| c.compressed_size)
.collect();
assert_eq!(planned, vec![37, 111, 5]);
struct SizedChunks<'a>(&'a [u64]);
impl ChunkProvider for SizedChunks<'_> {
fn chunk_bytes(&self, index: usize, out: &mut Vec<u8>) -> Result<(), FormatError> {
out.resize(self.0[index] as usize, 0xAB);
Ok(())
}
}
let sizes: Vec<u64> = meta.iter().map(|m| m.compressed_size).collect();
let mut emitted: Vec<u8> = Vec::new();
emit_chunked_data_verbatim(&mut emitted, &layout.plan, &SizedChunks(&sizes)).unwrap();
assert_eq!(emitted.len() as u64, layout.plan.total_len);
let chunk_bytes: u64 = sizes.iter().sum();
assert!(
layout.plan.total_len > chunk_bytes,
"three chunks take a fixed array, so the region is longer than its chunk bytes"
);
for (label, shape, maxshape, chunk_sizes, signature) in [
(
"fixed array",
&[21u64][..],
None,
&[37u64, 111, 5][..],
Some(&b"FAHD"[..]),
),
(
"extensible array",
&[21u64][..],
Some(&[u64::MAX][..]),
&[37u64, 111, 5][..],
Some(&b"EAHD"[..]),
),
("single chunk", &[7u64][..], None, &[37u64][..], None),
] {
let meta: Vec<ChunkMeta> = chunk_sizes
.iter()
.map(|&compressed_size| ChunkMeta {
compressed_size,
filter_mask: 0,
})
.collect();
let layout =
plan_chunked_data_verbatim(&meta, shape, &[7], nz(8), Some(&[]), 0x1000, maxshape)
.unwrap();
let mut emitted: Vec<u8> = Vec::new();
emit_chunked_data_verbatim(&mut emitted, &layout.plan, &SizedChunks(chunk_sizes))
.unwrap();
let chunk_bytes: usize = chunk_sizes.iter().sum::<u64>() as usize;
assert_eq!(
emitted.len() as u64,
layout.plan.total_len,
"{label}: the emit must fill the planned region"
);
match signature {
Some(sig) => assert_eq!(
&emitted[chunk_bytes..chunk_bytes + 4],
sig,
"{label}: the index must begin where the last chunk ends"
),
None => assert_eq!(
emitted.len(),
chunk_bytes,
"{label}: nothing may follow the chunk bytes"
),
}
}
}
#[test]
fn a_verbatim_plan_with_no_chunks_is_refused() {
let result = plan_chunked_data_verbatim(&[], &[21], &[7], nz(8), None, 0x1000, None);
assert!(
matches!(result, Err(FormatError::ChunkedReadError(_))),
"a chunk-less plan must be refused"
);
}
#[test]
fn chunk_options_auto_dims() {
let options = filtered(&[FilterKind::Deflate(6)]);
let dims = options.resolve_chunk_dims(&[100, 50]);
assert_eq!(dims, vec![100, 50]);
}
#[test]
fn the_scale_offset_fill_availability_reaches_the_filter_parameters() {
let ctx = f64_ctx(&[8]);
let elem = NonZeroUsize::new(8).expect("8 is non-zero");
let fill = 2.5f64.to_le_bytes();
let parms = |fill_availability, fill: FillPattern<'_>| {
let options = filtered(&[FilterKind::ScaleOffset(
ScaleOffset::FloatDScale(2),
fill_availability,
)]);
let pl = options.build_pipeline(&ctx, fill).unwrap().unwrap();
let f = pl
.filters
.iter()
.find(|f| f.filter_id == FILTER_SCALEOFFSET)
.expect("the scale-offset filter");
(f.client_data[7], f.client_data[8], f.client_data[9])
};
let bits = 2.5f64.to_bits();
assert_eq!(
parms(
FillAvailability::Defined,
FillPattern::new(Some(&fill), elem)
),
(1, bits as u32, (bits >> 32) as u32)
);
assert_eq!(
parms(FillAvailability::Defined, FillPattern::ZERO),
(1, 0, 0)
);
assert_eq!(
parms(
FillAvailability::Undefined,
FillPattern::new(Some(&fill), elem)
),
(0, 0, 0)
);
let options = filtered(&[FilterKind::ScaleOffset(
ScaleOffset::FloatDScale(2),
FillAvailability::Defined,
)]);
assert!(matches!(
options.build_pipeline(&ctx, FillPattern::UNKNOWN),
Err(FormatError::UnreadableFillValue)
));
let undefined = filtered(&[FilterKind::ScaleOffset(
ScaleOffset::FloatDScale(2),
FillAvailability::Undefined,
)]);
assert!(undefined.build_pipeline(&ctx, FillPattern::UNKNOWN).is_ok());
let mut db = crate::type_builders::DatasetBuilder::new("d");
db.with_scale_offset(ScaleOffset::FloatDScale(2));
assert_eq!(
db.chunk_options.filters,
vec![FilterSpec {
kind: FilterKind::ScaleOffset(
ScaleOffset::FloatDScale(2),
FillAvailability::Defined
),
optional: false,
}]
);
let plain = filtered(&[FilterKind::Deflate(6)]);
assert!(plain.build_pipeline(&ctx, FillPattern::UNKNOWN).is_ok());
}
#[test]
fn chunk_options_pipeline_deflate() {
let options = filtered(&[FilterKind::Deflate(6)]);
let pl = options
.build_pipeline(&ChunkContext::basic(&[], 8), FillPattern::ZERO)
.unwrap()
.unwrap();
assert_eq!(pl.filters.len(), 1);
assert_eq!(pl.filters[0].filter_id, FILTER_DEFLATE);
}
#[test]
fn chunk_options_pipeline_shuffle_deflate_fletcher32() {
let options = filtered(&[
FilterKind::Shuffle,
FilterKind::Deflate(6),
FilterKind::Fletcher32,
]);
let pl = options
.build_pipeline(&ChunkContext::basic(&[], 8), FillPattern::ZERO)
.unwrap()
.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 setting_filters_places_them_in_canonical_order_whatever_order_they_arrive_in() {
let ids = |o: &ChunkOptions| -> Vec<u16> {
o.filters.iter().map(|f| f.kind.filter_id()).collect()
};
let canonical = [FILTER_SHUFFLE, FILTER_DEFLATE, FILTER_FLETCHER32];
assert_eq!(
ids(&filtered(&[
FilterKind::Shuffle,
FilterKind::Deflate(6),
FilterKind::Fletcher32,
])),
canonical
);
assert_eq!(
ids(&filtered(&[
FilterKind::Fletcher32,
FilterKind::Deflate(6),
FilterKind::Shuffle,
])),
canonical
);
let twice = filtered(&[
FilterKind::Deflate(1),
FilterKind::Shuffle,
FilterKind::Deflate(9),
]);
assert_eq!(ids(&twice), [FILTER_SHUFFLE, FILTER_DEFLATE]);
assert_eq!(twice.filters[1].kind, FilterKind::Deflate(9));
}
#[test]
fn a_pushed_pipeline_keeps_its_order_and_its_optional_flags() {
let mut options = ChunkOptions::default();
for (kind, optional) in [
(FilterKind::Shuffle, true),
(FilterKind::Fletcher32, false),
(FilterKind::Deflate(4), true),
] {
options.push_filter(FilterSpec { kind, optional });
}
let pl = options
.build_pipeline(&ChunkContext::basic(&[], 8), FillPattern::ZERO)
.unwrap()
.unwrap();
let stored: Vec<(u16, u16)> = pl.filters.iter().map(|f| (f.filter_id, f.flags)).collect();
assert_eq!(
stored,
[
(FILTER_SHUFFLE, 1),
(FILTER_FLETCHER32, 0),
(FILTER_DEFLATE, 1),
]
);
}
#[test]
fn conflicting_filter_requests_are_refused() {
let so = ScaleOffset::FloatDScale(2);
let fill = FillAvailability::Defined;
let cases: &[(&str, &str, ChunkOptions)] = &[
(
"lzf",
"deflate",
filtered(&[FilterKind::Lzf, FilterKind::Deflate(6)]),
),
(
"shuffle",
"scale-offset",
filtered(&[FilterKind::Shuffle, FilterKind::ScaleOffset(so, fill)]),
),
#[cfg(feature = "zfp")]
(
"scale-offset",
"ZFP",
filtered(&[FilterKind::ScaleOffset(so, fill), FilterKind::Zfp(16.0)]),
),
#[cfg(feature = "zfp")]
(
"shuffle",
"ZFP",
filtered(&[FilterKind::Shuffle, FilterKind::Zfp(16.0)]),
),
#[cfg(feature = "zfp")]
(
"lzf",
"ZFP",
filtered(&[FilterKind::Lzf, FilterKind::Zfp(16.0)]),
),
#[cfg(feature = "zfp")]
(
"deflate",
"ZFP",
filtered(&[FilterKind::Deflate(6), FilterKind::Zfp(16.0)]),
),
];
for (a, b, options) in cases {
let err = options
.build_pipeline(&f64_ctx(&[64]), FillPattern::ZERO)
.expect_err("{a} + {b} was accepted");
let FormatError::FilterError(msg) = &err else {
panic!("{a} + {b}: expected a filter error, got {err}");
};
assert!(msg.contains(a) && msg.contains(b), "{a} + {b}: {msg}");
}
}
fn f64_ctx(chunk_dims: &[u64]) -> ChunkContext<'_> {
ChunkContext {
chunk_dims,
element_size: core::num::NonZeroU32::new(8).expect("8 is non-zero"),
element_type: zfp_f64_type(),
scale_offset_type: crate::scaleoffset::scale_offset_type_from_datatype(&make_f64_type()),
}
}
#[cfg(feature = "zfp")]
fn zfp_f64_type() -> Option<crate::filters::ZfpElementTypeWhenEnabled> {
crate::filters::zfp_element_type_from_datatype(&make_f64_type())
}
#[cfg(not(feature = "zfp"))]
fn zfp_f64_type() -> Option<crate::filters::ZfpElementTypeWhenEnabled> {
None
}
#[test]
fn compatible_filter_requests_still_build() {
let so = FilterKind::ScaleOffset(ScaleOffset::FloatDScale(2), FillAvailability::Defined);
let cases: &[(ChunkOptions, &[u16])] = &[
(
filtered(&[FilterKind::Shuffle, FilterKind::Deflate(6)]),
&[FILTER_SHUFFLE, FILTER_DEFLATE],
),
(
filtered(&[FilterKind::Shuffle, FilterKind::Lzf]),
&[FILTER_SHUFFLE, FILTER_LZF],
),
(
filtered(&[so, FilterKind::Deflate(6)]),
&[FILTER_SCALEOFFSET, FILTER_DEFLATE],
),
(
filtered(&[so, FilterKind::Lzf, FilterKind::Fletcher32]),
&[FILTER_SCALEOFFSET, FILTER_LZF, FILTER_FLETCHER32],
),
];
for (options, expected) in cases {
let pl = options
.build_pipeline(&f64_ctx(&[64]), FillPattern::ZERO)
.unwrap()
.unwrap();
let ids: Vec<u16> = pl.filters.iter().map(|f| f.filter_id).collect();
assert_eq!(&ids, expected);
}
}
#[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,
filter_mask: 0,
},
WrittenChunk {
address: 0x10A0,
compressed_size: 160,
filter_mask: 0,
},
];
let fa = build_fixed_array_at(&IndexSlots::dense(&chunks), 160, 8, 8, false, 0x2000);
assert_eq!(&fa[0..4], b"FAHD");
assert_eq!(&fa[28..32], b"FADB");
}
#[test]
fn serialize_v4_extensible_array_roundtrip() {
let msg = 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,
filter_mask: 0,
},
WrittenChunk {
address: 0x1050,
compressed_size: 80,
filter_mask: 0,
},
];
let ea = build_extensible_array_at(&IndexSlots::dense(&chunks), 80, 8, 8, false, 0x2000)
.unwrap();
assert_eq!(&ea[0..4], b"EAHD");
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");
}
fn roundtrip_ea(
values: &[f64],
shape: &[u64],
chunk_dims: &[u64],
maxshape: &[u64],
) -> Vec<f64> {
let raw = f64_to_bytes(values);
let data_address = 0x1000u64;
let options = ChunkOptions {
chunk_dims: Some(chunk_dims.to_vec()),
..Default::default()
};
let ctx = ChunkContext::basic(chunk_dims, 8);
let result = build_chunked_data_at_ext(
&raw,
shape,
ctx,
&options,
data_address,
Some(maxshape),
FillPattern::ZERO,
)
.unwrap();
let file_size = data_address as usize + result.data_bytes.len();
let mut file_data = vec![0u8; file_size];
file_data[data_address as usize..].copy_from_slice(&result.data_bytes);
let layout = DataLayout::parse(&result.layout_message, 8, 8).unwrap();
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_cached(
&file_data,
RawReadSpec::plain(&layout, &dataspace, &datatype),
8,
8,
&ChunkCache::new(),
)
.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);
}
#[test]
fn ea_roundtrip_super_block_sizes() {
for &n in &[245u64, 300, 2000, 50000] {
let values: Vec<f64> = (0..n).map(|i| i as f64).collect();
let result = roundtrip_ea(&values, &[n], &[1], &[u64::MAX]);
assert_eq!(result.len(), n as usize, "length mismatch at n={n}");
assert_eq!(result, values, "data mismatch at n={n}");
}
}
#[test]
fn ea_roundtrip_paged_data_blocks() {
let n: u64 = 132_000;
let values: Vec<f64> = (0..n).map(|i| i as f64).collect();
let result = roundtrip_ea(&values, &[n], &[1], &[u64::MAX]);
assert_eq!(result.len(), n as usize);
assert_eq!(result, values);
}
#[cfg(feature = "std")]
#[test]
fn ea_compute_stats_matches_builder() {
use crate::extensible_array::{EaGeometry, ExtensibleArrayHeader};
let geom_header = ExtensibleArrayHeader {
client_id: 0,
element_size: 8,
max_nelmts_bits: 32,
idx_blk_elmts: 4,
min_dblk_nelmts: 16,
super_blk_min_nelmts: 4,
max_dblk_nelmts_bits: 10,
num_elements: 0,
index_block_address: 0,
};
let geom = EaGeometry::from_header(&geom_header);
for &n in &[1u64, 4, 20, 100, 244, 300, 2000, 50000, 131056, 140000] {
let chunks: Vec<WrittenChunk> = (0..n)
.map(|i| WrittenChunk {
address: 0x1000 + i * 8,
compressed_size: 8,
filter_mask: 0,
})
.collect();
let ea =
build_extensible_array_at(&IndexSlots::dense(&chunks), 8, 8, 8, false, 0x100000)
.unwrap();
let stat =
|k: usize| u64::from_le_bytes(ea[12 + k * 8..12 + k * 8 + 8].try_into().unwrap());
let built = super::EaStats {
nsuper_blks: stat(0),
super_blk_size: stat(1),
ndata_blks: stat(2),
data_blk_size: stat(3),
max_idx_set: stat(4),
nelmts: stat(5),
};
let computed =
super::ea_compute_stats(&geom, 4, 8, 1024, 8, 4, n, SlotOccupancy::Dense(n));
assert_eq!(computed, built, "stats mismatch at n={n}");
}
}
#[test]
fn chunked_data_len_matches_what_assemble_produces() {
for &(elements, chunk) in &[
(1u64, 8u64), (21, 7), (8_192, 4), ] {
for &deflate in &[false, true] {
for &unlimited in &[false, true] {
let values: Vec<f64> = (0..elements).map(|i| i as f64).collect();
let raw = f64_to_bytes(&values);
let mut options = chunked(&[chunk], &[]);
if deflate {
options.set_filter(FilterKind::Deflate(6));
}
let dims = [chunk];
let maxshape = unlimited.then_some([u64::MAX]);
let set = compress_chunks(
&raw,
&[elements],
ChunkContext::basic(&dims, 8),
&options,
maxshape.as_ref().map(<[u64; 1]>::as_slice),
FillPattern::ZERO,
StorageAllocation::Allocated,
)
.unwrap();
let planned = chunked_data_len(&set).unwrap();
let assembled = assemble_chunked_at(&set, 0x10_0000).unwrap();
assert_eq!(
planned,
assembled.data_bytes.len() as u64,
"planned region must match the assembled one at elements={elements}, \
chunk={chunk}, deflate={deflate}, unlimited={unlimited}"
);
}
}
}
}
const CHUNK_BYTES: [u64; 4] = [8, 300, 100_000, 1 << 32];
#[test]
fn fixed_array_len_matches_what_it_builds() {
fn check(n: u64, chunk_bytes: u64, offset_size: u8, length_size: u8, has_filters: bool) {
let chunks: Vec<WrittenChunk> = (0..n)
.map(|i| WrittenChunk {
address: 0x1000 + i * 8,
compressed_size: 8,
filter_mask: 0,
})
.collect();
let planned = fixed_array_len(
&IndexSlots::dense(&chunks),
chunk_bytes,
offset_size,
length_size,
has_filters,
);
let built = build_fixed_array_at(
&IndexSlots::dense(&chunks),
chunk_bytes,
offset_size,
length_size,
has_filters,
0x10_0000,
);
assert_eq!(
planned,
built.len() as u64,
"planned length must match the emitted array at n={n}, \
chunk_bytes={chunk_bytes}, offset_size={offset_size}, \
has_filters={has_filters}"
);
}
for &(offset_size, length_size) in &[(8u8, 8u8), (4u8, 4u8)] {
for &has_filters in &[false, true] {
for n in 0..=1_100u64 {
check(n, 8, offset_size, length_size, has_filters);
}
for &n in &[4_096u64, 5_000, 100_000] {
check(n, 8, offset_size, length_size, has_filters);
}
}
}
for &chunk_bytes in &CHUNK_BYTES {
for &n in &[1u64, 1_024, 1_025, 5_000] {
check(n, chunk_bytes, 8, 8, true);
}
}
}
#[test]
fn extensible_array_len_matches_what_it_builds() {
fn check(n: u64, chunk_bytes: u64, offset_size: u8, length_size: u8, has_filters: bool) {
let chunks: Vec<WrittenChunk> = (0..n)
.map(|i| WrittenChunk {
address: 0x1000 + i * 8,
compressed_size: 8,
filter_mask: 0,
})
.collect();
let planned = extensible_array_len(
&IndexSlots::dense(&chunks),
chunk_bytes,
offset_size,
length_size,
has_filters,
);
let built = build_extensible_array_at(
&IndexSlots::dense(&chunks),
chunk_bytes,
offset_size,
length_size,
has_filters,
0x10_0000,
)
.unwrap();
assert_eq!(
planned,
built.len() as u64,
"planned length must match the emitted array at n={n}, \
chunk_bytes={chunk_bytes}, offset_size={offset_size}, \
has_filters={has_filters}"
);
}
for &(offset_size, length_size) in &[(8u8, 8u8), (4u8, 4u8)] {
for &has_filters in &[false, true] {
for n in 0..=250u64 {
check(n, 8, offset_size, length_size, has_filters);
}
for &n in &[300u64, 2_000, 50_000, 131_060, 131_061, 140_000] {
check(n, 8, offset_size, length_size, has_filters);
}
}
}
for &chunk_bytes in &CHUNK_BYTES {
for &n in &[1u64, 5, 244, 300, 2_000] {
check(n, chunk_bytes, 8, 8, true);
}
}
let widths: Vec<usize> = CHUNK_BYTES
.iter()
.map(|&chunk_bytes| {
super::ea_layout(SlotOccupancy::Dense(0), 0, chunk_bytes, 8, 8, true)
.encoding
.chunk_size_bytes
})
.collect();
let mut distinct = widths.clone();
distinct.sort_unstable();
distinct.dedup();
assert_eq!(
distinct.len(),
CHUNK_BYTES.len(),
"each chunk size must select a different compressed-size field width, got {widths:?}"
);
}
#[cfg(feature = "std")]
fn h5py_run(path: &std::path::Path, script: &str) -> Option<String> {
let o = std::process::Command::new("python3")
.args(["-c", script, &path.to_string_lossy()])
.output()
.ok()?;
if !o.status.success() {
let err = String::from_utf8_lossy(&o.stderr);
if err.contains("No module named") {
return None; }
panic!("h5py: {err}");
}
Some(String::from_utf8(o.stdout).unwrap().trim().to_string())
}
#[cfg(feature = "std")]
#[test]
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 dir = tempfile::tempdir().unwrap();
let path = dir.path().join("rustyhdf5_chunked_multi.h5");
std::fs::write(&path, &bytes).unwrap();
let script = "import sys,h5py,json; f=h5py.File(sys.argv[1],'r'); print(json.dumps({'a':f['a'][:].tolist(),'b':f['b'][:].tolist()}))";
let Some(out) = h5py_run(&path, script) else {
return;
};
let v: serde_json::Value = serde_json::from_str(&out).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]
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 dir = tempfile::tempdir().unwrap();
let path = dir.path().join("rustyhdf5_chunked_attrs.h5");
std::fs::write(&path, &bytes).unwrap();
let script = "import sys,h5py,json; f=h5py.File(sys.argv[1],'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'])}))";
let Some(out) = h5py_run(&path, script) else {
return;
};
let v: serde_json::Value = serde_json::from_str(&out).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"));
}
}