format: decode chunks into reusable scratch buffers
decompress_chunk_exact_with decodes a chunk into a DecodeScratch the caller keeps between chunks, instead of a new Vec per chunk and per filter stage. Deflate inflates into a kept buffer with a reset (not rebuilt) inflater, shuffle interleaves into the other buffer, and Fletcher32 checks and drops its checksum in place (on the stored bytes when it is the first filter undone). Other filters go through the registry as before. Output and errors are those of decompress_chunk_exact; a unit test checks that for every pipeline shape and filter mask with one reused scratch. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
osobh
@@ -148,6 +148,180 @@ pub fn decompress_chunk_exact(
|
||||
Ok(data)
|
||||
}
|
||||
|
||||
/// Buffers a chunk decoder keeps between chunks, so decoding a dataset's
|
||||
/// chunks one after another reuses the same memory instead of allocating
|
||||
/// (and faulting in) fresh buffers for every chunk and every filter stage.
|
||||
///
|
||||
/// Use one per thread with [`decompress_chunk_exact_with`]. Buffers larger
|
||||
/// than [`DecodeScratch::RETAIN_BYTES`] are released by
|
||||
/// [`DecodeScratch::trim`], so a scratch kept for a long time (a
|
||||
/// thread-local, say) does not hold on to a huge chunk's memory.
|
||||
#[derive(Default)]
|
||||
pub struct DecodeScratch {
|
||||
a: Vec<u8>,
|
||||
b: Vec<u8>,
|
||||
#[cfg(feature = "deflate")]
|
||||
inflater: Option<flate2::Decompress>,
|
||||
}
|
||||
|
||||
impl core::fmt::Debug for DecodeScratch {
|
||||
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
|
||||
f.debug_struct("DecodeScratch")
|
||||
.field("a_capacity", &self.a.capacity())
|
||||
.field("b_capacity", &self.b.capacity())
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
/// Which buffer holds the data between two filter stages.
|
||||
#[derive(Clone, Copy)]
|
||||
enum Stage {
|
||||
/// `compressed[..len]`: still (a prefix of) the stored bytes.
|
||||
Stored(usize),
|
||||
A,
|
||||
B,
|
||||
}
|
||||
|
||||
impl DecodeScratch {
|
||||
/// Largest buffer [`trim`](Self::trim) keeps (4 MiB).
|
||||
pub const RETAIN_BYTES: usize = 4 << 20;
|
||||
|
||||
/// An empty scratch; buffers are allocated on first use.
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Release any buffer larger than [`Self::RETAIN_BYTES`].
|
||||
pub fn trim(&mut self) {
|
||||
for buf in [&mut self.a, &mut self.b] {
|
||||
if buf.capacity() > Self::RETAIN_BYTES {
|
||||
*buf = Vec::new();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// [`decompress_chunk_exact`] into reusable buffers: the decoded chunk is
|
||||
/// returned as a slice of `scratch` (or of `compressed`, when every filter
|
||||
/// that was applied only appended a checksum), valid until `scratch` is used
|
||||
/// again.
|
||||
///
|
||||
/// Deflate, shuffle and Fletcher32 — h5py's and libhdf5's usual pipeline —
|
||||
/// decode without allocating once `scratch` has grown to the chunk size: the
|
||||
/// inflater writes into a kept buffer (and its state is reset, not
|
||||
/// rebuilt), shuffle interleaves into the other buffer, and Fletcher32 checks
|
||||
/// the checksum and drops it in place. Every other filter goes through the
|
||||
/// filter registry as [`decompress_chunk_masked`] does, and its output
|
||||
/// replaces a scratch buffer. The result is byte for byte what
|
||||
/// [`decompress_chunk_exact`] returns, with the same errors.
|
||||
pub fn decompress_chunk_exact_with<'s>(
|
||||
compressed: &'s [u8],
|
||||
pipeline: &FilterPipeline,
|
||||
chunk_size: usize,
|
||||
element_size: u32,
|
||||
filter_mask: u32,
|
||||
coords: &[u64],
|
||||
scratch: &'s mut DecodeScratch,
|
||||
) -> Result<&'s [u8], FormatError> {
|
||||
// Same per-stage bounds as `decompress_chunk_masked`.
|
||||
let mut bounds = [0usize; 32];
|
||||
let mut bounds_vec = Vec::new();
|
||||
let bounds: &mut [usize] = if pipeline.filters.len() <= bounds.len() {
|
||||
&mut bounds[..pipeline.filters.len()]
|
||||
} else {
|
||||
bounds_vec.resize(pipeline.filters.len(), 0);
|
||||
&mut bounds_vec
|
||||
};
|
||||
let mut size = chunk_size;
|
||||
for (i, filter) in pipeline.filters.iter().enumerate() {
|
||||
bounds[i] = size;
|
||||
if !filter_skipped(filter_mask, i) {
|
||||
size = filter_output_bound(filter.filter_id, size);
|
||||
}
|
||||
}
|
||||
|
||||
let mut stage = Stage::Stored(compressed.len());
|
||||
for (i, filter) in pipeline.filters.iter().enumerate().rev() {
|
||||
if filter_skipped(filter_mask, i) {
|
||||
continue;
|
||||
}
|
||||
let ctx = FilterContext {
|
||||
filter,
|
||||
element_size: element_size as usize,
|
||||
max_output: bounds[i],
|
||||
};
|
||||
// The stage's input, and the buffer its output goes to (the one
|
||||
// not holding the input).
|
||||
let (input, out): (&[u8], &mut Vec<u8>) = match stage {
|
||||
Stage::Stored(len) => (&compressed[..len], &mut scratch.a),
|
||||
Stage::A => (&scratch.a, &mut scratch.b),
|
||||
Stage::B => (&scratch.b, &mut scratch.a),
|
||||
};
|
||||
let next = match stage {
|
||||
Stage::Stored(_) | Stage::B => Stage::A,
|
||||
Stage::A => Stage::B,
|
||||
};
|
||||
match filter.filter_id {
|
||||
// Built in and never overridable (`register_filter` refuses
|
||||
// built-in IDs), so the registry would pick exactly these.
|
||||
FILTER_FLETCHER32 => {
|
||||
// Check and drop the checksum where the data is.
|
||||
let payload = fletcher32_payload(input)?;
|
||||
stage = match stage {
|
||||
Stage::Stored(_) => Stage::Stored(payload),
|
||||
Stage::A => {
|
||||
scratch.a.truncate(payload);
|
||||
Stage::A
|
||||
}
|
||||
Stage::B => {
|
||||
scratch.b.truncate(payload);
|
||||
Stage::B
|
||||
}
|
||||
};
|
||||
continue;
|
||||
}
|
||||
FILTER_SHUFFLE => shuffle_decompress_into(input, ctx.element_size, out),
|
||||
#[cfg(all(
|
||||
feature = "deflate",
|
||||
not(all(target_os = "macos", feature = "system-zlib-decompress"))
|
||||
))]
|
||||
FILTER_DEFLATE => {
|
||||
let limit = if ctx.max_output != 0 {
|
||||
ctx.max_output
|
||||
} else {
|
||||
MAX_DECOMPRESS_SIZE
|
||||
};
|
||||
let size_hint = if ctx.max_output != 0 {
|
||||
ctx.max_output
|
||||
} else {
|
||||
input.len().saturating_mul(4).min(1 << 20)
|
||||
};
|
||||
let inflater = scratch
|
||||
.inflater
|
||||
.get_or_insert_with(|| flate2::Decompress::new(true));
|
||||
inflater.reset(true);
|
||||
inflate_bounded_into(inflater, input, size_hint, limit, out)
|
||||
.map_err(FormatError::DecompressionError)?;
|
||||
}
|
||||
_ => *out = filter_registry::decode(input, &ctx)?,
|
||||
}
|
||||
stage = next;
|
||||
}
|
||||
|
||||
let data: &[u8] = match stage {
|
||||
Stage::Stored(len) => &compressed[..len],
|
||||
Stage::A => &scratch.a,
|
||||
Stage::B => &scratch.b,
|
||||
};
|
||||
if chunk_size != 0 && data.len() != chunk_size {
|
||||
return Err(FormatError::ChunkedReadError(format!(
|
||||
"chunk at {coords:?} decoded to {} bytes, expected {chunk_size}",
|
||||
data.len()
|
||||
)));
|
||||
}
|
||||
Ok(data)
|
||||
}
|
||||
|
||||
/// Apply a filter pipeline to compress a chunk.
|
||||
/// Filters are applied in FORWARD order for compression.
|
||||
pub fn compress_chunk(
|
||||
@@ -892,33 +1066,55 @@ pub(crate) fn inflate_bounded(
|
||||
size_hint: usize,
|
||||
limit: usize,
|
||||
) -> Result<Vec<u8>, String> {
|
||||
use flate2::{Decompress, FlushDecompress, Status};
|
||||
let mut out = Vec::new();
|
||||
inflate_bounded_into(
|
||||
&mut flate2::Decompress::new(true),
|
||||
data,
|
||||
size_hint,
|
||||
limit,
|
||||
&mut out,
|
||||
)?;
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// [`inflate_bounded`] with a fresh or reset `inflater`, into `out`: its
|
||||
/// contents are replaced and its allocation reused.
|
||||
#[cfg(feature = "deflate")]
|
||||
fn inflate_bounded_into(
|
||||
inflater: &mut flate2::Decompress,
|
||||
data: &[u8],
|
||||
size_hint: usize,
|
||||
limit: usize,
|
||||
out: &mut Vec<u8>,
|
||||
) -> Result<(), String> {
|
||||
use flate2::{FlushDecompress, Status};
|
||||
|
||||
// One byte of headroom past the limit distinguishes an over-size stream
|
||||
// from one that legitimately ends exactly at the limit.
|
||||
let max_capacity = limit.saturating_add(1);
|
||||
let mut out = Vec::new();
|
||||
out.try_reserve_exact(size_hint.clamp(1, max_capacity))
|
||||
// A kept buffer may already be larger than `max_capacity`; the decoder
|
||||
// can then write past the limit, which the check below still refuses.
|
||||
out.clear();
|
||||
let want = size_hint.clamp(1, max_capacity);
|
||||
out.try_reserve_exact(want)
|
||||
.map_err(|e| format!("deflate: cannot allocate output: {e}"))?;
|
||||
|
||||
let mut inflater = Decompress::new(true);
|
||||
loop {
|
||||
let (in_before, out_before) = (inflater.total_in(), inflater.total_out());
|
||||
let status = inflater
|
||||
.decompress_vec(
|
||||
&data[in_before as usize..],
|
||||
&mut out,
|
||||
FlushDecompress::Finish,
|
||||
)
|
||||
.decompress_vec(&data[in_before as usize..], out, FlushDecompress::Finish)
|
||||
.map_err(|e| format!("deflate: {e}"))?;
|
||||
if out.len() > limit {
|
||||
return Err("deflate: output exceeds size limit".into());
|
||||
}
|
||||
match status {
|
||||
Status::StreamEnd => return Ok(out),
|
||||
Status::StreamEnd => return Ok(()),
|
||||
Status::Ok | Status::BufError if out.len() == out.capacity() => {
|
||||
// Out of room: double, up to the limit.
|
||||
let grow = out.capacity().min(max_capacity - out.capacity()).max(1);
|
||||
let grow = out
|
||||
.capacity()
|
||||
.min(max_capacity.saturating_sub(out.capacity()))
|
||||
.max(1);
|
||||
out.try_reserve_exact(grow)
|
||||
.map_err(|e| format!("deflate: cannot allocate output: {e}"))?;
|
||||
}
|
||||
@@ -1217,16 +1413,30 @@ fn zstd_compress(data: &[u8], level: u32) -> Result<Vec<u8>, FormatError> {
|
||||
/// On disk: all byte-0s of each element together, then all byte-1s, etc.
|
||||
/// Output: elements in natural order.
|
||||
fn shuffle_decompress(data: &[u8], element_size: usize) -> Result<Vec<u8>, FormatError> {
|
||||
let mut result = Vec::new();
|
||||
shuffle_decompress_into(data, element_size, &mut result);
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// [`shuffle_decompress`] into `result`, replacing its contents and reusing
|
||||
/// its allocation.
|
||||
fn shuffle_decompress_into(data: &[u8], element_size: usize, result: &mut Vec<u8>) {
|
||||
if element_size <= 1 {
|
||||
return Ok(data.to_vec());
|
||||
result.clear();
|
||||
result.extend_from_slice(data);
|
||||
return;
|
||||
}
|
||||
// Like libhdf5, only whole elements are shuffled; trailing bytes (e.g. a
|
||||
// Fletcher32 checksum appended before the shuffle) are stored as-is.
|
||||
let whole = data.len() - data.len() % element_size;
|
||||
let (data, tail) = data.split_at(whole);
|
||||
let num_elements = data.len() / element_size;
|
||||
let mut result = vec![0u8; whole];
|
||||
result.reserve_exact(tail.len());
|
||||
// Every byte of `result[..whole]` is overwritten below, so a reused
|
||||
// buffer keeps its old bytes instead of being zeroed first; only growth
|
||||
// is zero-filled.
|
||||
result.truncate(whole);
|
||||
result.reserve_exact(whole + tail.len() - result.len());
|
||||
result.resize(whole, 0);
|
||||
|
||||
// The shuffled stream is `element_size` byte planes of `num_elements`
|
||||
// bytes each; un-shuffling interleaves them. This is on the read path of
|
||||
@@ -1245,10 +1455,10 @@ fn shuffle_decompress(data: &[u8], element_size: usize) -> Result<Vec<u8>, Forma
|
||||
}
|
||||
}
|
||||
match element_size {
|
||||
2 => interleave::<2>(data, num_elements, &mut result),
|
||||
4 => interleave::<4>(data, num_elements, &mut result),
|
||||
8 => interleave::<8>(data, num_elements, &mut result),
|
||||
16 => interleave::<16>(data, num_elements, &mut result),
|
||||
2 => interleave::<2>(data, num_elements, result),
|
||||
4 => interleave::<4>(data, num_elements, result),
|
||||
8 => interleave::<8>(data, num_elements, result),
|
||||
16 => interleave::<16>(data, num_elements, result),
|
||||
_ => {
|
||||
for (i, element) in result.chunks_exact_mut(element_size).enumerate() {
|
||||
for (j, byte) in element.iter_mut().enumerate() {
|
||||
@@ -1258,8 +1468,6 @@ fn shuffle_decompress(data: &[u8], element_size: usize) -> Result<Vec<u8>, Forma
|
||||
}
|
||||
}
|
||||
result.extend_from_slice(tail);
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Shuffle (compress direction): group bytes by position within each element.
|
||||
@@ -1411,6 +1619,12 @@ fn fletcher32_compute(data: &[u8]) -> u32 {
|
||||
/// Verify Fletcher32 checksum and strip it from the data.
|
||||
/// The last 4 bytes are the stored checksum.
|
||||
fn fletcher32_verify(data: &[u8]) -> Result<Vec<u8>, FormatError> {
|
||||
fletcher32_payload(data).map(|len| data[..len].to_vec())
|
||||
}
|
||||
|
||||
/// Verify the Fletcher32 checksum that ends `data`; the length of the data
|
||||
/// before it.
|
||||
fn fletcher32_payload(data: &[u8]) -> Result<usize, FormatError> {
|
||||
if data.len() < 4 {
|
||||
return Err(FormatError::FilterError(
|
||||
"fletcher32: data too short for checksum".into(),
|
||||
@@ -1430,7 +1644,7 @@ fn fletcher32_verify(data: &[u8]) -> Result<Vec<u8>, FormatError> {
|
||||
computed,
|
||||
});
|
||||
}
|
||||
Ok(payload.to_vec())
|
||||
Ok(payload.len())
|
||||
}
|
||||
|
||||
/// Append Fletcher32 checksum to data.
|
||||
@@ -1545,6 +1759,113 @@ fn pcodec_decompress(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
/// `decompress_chunk_exact_with` returns exactly what
|
||||
/// `decompress_chunk_exact` returns — data or error — for every pipeline
|
||||
/// shape, filter mask and chunk size, with one scratch reused across all
|
||||
/// of them in an order that grows, shrinks and swaps its buffers.
|
||||
#[test]
|
||||
fn decode_with_scratch_matches_the_allocating_decoder() {
|
||||
let f = |filter_id: u16, client_data: Vec<u32>| FilterDescription {
|
||||
filter_id,
|
||||
name: None,
|
||||
flags: 0,
|
||||
client_data,
|
||||
};
|
||||
let mut pipelines = vec![
|
||||
vec![f(FILTER_SHUFFLE, vec![4])],
|
||||
vec![f(FILTER_FLETCHER32, vec![])],
|
||||
// NetCDF-4's order: the checksum is taken before shuffle.
|
||||
vec![f(FILTER_FLETCHER32, vec![]), f(FILTER_SHUFFLE, vec![4])],
|
||||
];
|
||||
#[cfg(feature = "deflate")]
|
||||
pipelines.extend([
|
||||
vec![f(FILTER_DEFLATE, vec![4])],
|
||||
vec![f(FILTER_SHUFFLE, vec![4]), f(FILTER_DEFLATE, vec![4])],
|
||||
// h5py's order with `fletcher32=True`: checksum last.
|
||||
vec![
|
||||
f(FILTER_SHUFFLE, vec![4]),
|
||||
f(FILTER_DEFLATE, vec![4]),
|
||||
f(FILTER_FLETCHER32, vec![]),
|
||||
],
|
||||
vec![
|
||||
f(FILTER_FLETCHER32, vec![]),
|
||||
f(FILTER_SHUFFLE, vec![4]),
|
||||
f(FILTER_DEFLATE, vec![1]),
|
||||
],
|
||||
]);
|
||||
#[cfg(feature = "lzf")]
|
||||
pipelines.push(vec![
|
||||
f(FILTER_SHUFFLE, vec![4]),
|
||||
f(crate::filter_pipeline::FILTER_LZF, vec![]),
|
||||
]);
|
||||
|
||||
let mut scratch = DecodeScratch::new();
|
||||
for elements in [1usize, 7, 4096, 3, 65536, 100] {
|
||||
let data: Vec<u8> = (0..elements as u32)
|
||||
.flat_map(|i| (i.wrapping_mul(2654435761) >> (i % 13)).to_le_bytes())
|
||||
.collect();
|
||||
for filters in &pipelines {
|
||||
let pipeline = FilterPipeline {
|
||||
version: 2,
|
||||
filters: filters.clone(),
|
||||
};
|
||||
let n = filters.len() as u32;
|
||||
for mask in 0..(1u32 << n) {
|
||||
// Encode only the filters the mask says were applied.
|
||||
let mut stored = data.clone();
|
||||
for (i, filter) in filters.iter().enumerate() {
|
||||
if mask & (1 << i) == 0 {
|
||||
let ctx = FilterContext {
|
||||
filter,
|
||||
element_size: 4,
|
||||
max_output: 0,
|
||||
};
|
||||
stored = filter_registry::encode(&stored, &ctx).unwrap();
|
||||
}
|
||||
}
|
||||
let mut cases = vec![(stored.clone(), data.len())];
|
||||
// Corrupt: last byte flipped, truncated, wrong size.
|
||||
let mut flipped = stored.clone();
|
||||
*flipped.last_mut().unwrap() ^= 0x5a;
|
||||
cases.push((flipped, data.len()));
|
||||
cases.push((stored[..stored.len() / 2].to_vec(), data.len()));
|
||||
cases.push((stored.clone(), data.len() + 4));
|
||||
cases.push((stored.clone(), 0));
|
||||
for (bytes, size) in cases {
|
||||
let want = decompress_chunk_exact(&bytes, &pipeline, size, 4, mask, &[3]);
|
||||
let got = decompress_chunk_exact_with(
|
||||
&bytes,
|
||||
&pipeline,
|
||||
size,
|
||||
4,
|
||||
mask,
|
||||
&[3],
|
||||
&mut scratch,
|
||||
)
|
||||
.map(<[u8]>::to_vec);
|
||||
match (&want, &got) {
|
||||
(Ok(w), Ok(g)) => assert_eq!(w, g, "{filters:?} mask {mask}"),
|
||||
(Err(w), Err(g)) => {
|
||||
assert_eq!(w.to_string(), g.to_string(), "{filters:?}")
|
||||
}
|
||||
_ => panic!("{filters:?} mask {mask} size {size}: {want:?} vs {got:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// Long-lived scratch gives back a huge chunk's buffers.
|
||||
let big = vec![0u8; DecodeScratch::RETAIN_BYTES + 8];
|
||||
let shuffle = FilterPipeline {
|
||||
version: 2,
|
||||
filters: vec![f(FILTER_SHUFFLE, vec![4])],
|
||||
};
|
||||
decompress_chunk_exact_with(&big, &shuffle, big.len(), 4, 0, &[0], &mut scratch).unwrap();
|
||||
scratch.trim();
|
||||
assert!(scratch.a.capacity() <= DecodeScratch::RETAIN_BYTES);
|
||||
assert!(scratch.b.capacity() <= DecodeScratch::RETAIN_BYTES);
|
||||
}
|
||||
|
||||
/// A chunk whose pipeline decodes to fewer bytes than the chunk holds is
|
||||
/// an error naming the chunk, never a short buffer the reader pads.
|
||||
#[test]
|
||||
|
||||
Reference in New Issue
Block a user