diff --git a/crates/clawhdf5-format/src/chunked_read.rs b/crates/clawhdf5-format/src/chunked_read.rs index b0f5141..2dab147 100644 --- a/crates/clawhdf5-format/src/chunked_read.rs +++ b/crates/clawhdf5-format/src/chunked_read.rs @@ -15,7 +15,7 @@ use crate::datatype::Datatype; use crate::error::FormatError; use crate::extensible_array::{ExtensibleArrayHeader, read_extensible_array_chunks}; use crate::filter_pipeline::FilterPipeline; -use crate::filters::decompress_chunk; +use crate::filters::{all_filters_skipped, decompress_chunk_masked}; use crate::fixed_array::{FixedArrayHeader, read_fixed_array_chunks}; #[cfg(feature = "std")] use std::sync::Arc; @@ -65,11 +65,13 @@ fn decompress_all_chunks( let raw_chunk = &file_data[c_addr..c_addr + size]; let decompressed = if let Some(pl) = pipeline { - if chunk_info.filter_mask == 0 { - decompress_chunk(raw_chunk, pl, chunk_total_bytes, element_size)? - } else { - raw_chunk.to_vec() - } + decompress_chunk_masked( + raw_chunk, + pl, + chunk_total_bytes, + element_size, + chunk_info.filter_mask, + )? } else { raw_chunk.to_vec() }; @@ -886,10 +888,11 @@ pub fn read_chunked_data_cached( }; // Chunks stored as-is (no pipeline, or the filter mask says this chunk - // skipped it) are copied straight from the file bytes: they are already in - // memory, so routing them through a Vec and then an aligned cache buffer - // was two extra copies of the whole dataset for nothing. - let stored_raw = |c: &ChunkInfo| pipeline.is_none() || c.filter_mask != 0; + // skipped every filter) are copied straight from the file bytes: they are + // already in memory, so routing them through a Vec and then an aligned + // cache buffer was two extra copies of the whole dataset for nothing. + let stored_raw = + |c: &ChunkInfo| pipeline.is_none_or(|pl| all_filters_skipped(pl, c.filter_mask)); let mut misses: Vec<&ChunkInfo> = Vec::new(); for chunk_info in &chunks { if stored_raw(chunk_info) { @@ -911,7 +914,13 @@ pub fn read_chunked_data_cached( let cache_them = total_bytes <= cache.max_bytes(); if let Some(pl) = pipeline { let decode = |c: &&ChunkInfo| -> Result, FormatError> { - decompress_chunk(raw_bytes(c)?, pl, chunk_total_bytes, elem_size as u32) + decompress_chunk_masked( + raw_bytes(c)?, + pl, + chunk_total_bytes, + elem_size as u32, + c.filter_mask, + ) }; for batch in misses.chunks(DECODE_BATCH) { #[cfg(feature = "parallel")] @@ -1194,11 +1203,13 @@ pub fn read_chunked_data_sweep( ensure_len(file_data, c_addr, size)?; let raw_chunk = &file_data[c_addr..c_addr + size]; let dec = if let Some(pl) = pipeline { - if chunk_info.filter_mask == 0 { - decompress_chunk(raw_chunk, pl, chunk_total_bytes, elem_size as u32)? - } else { - raw_chunk.to_vec() - } + decompress_chunk_masked( + raw_chunk, + pl, + chunk_total_bytes, + elem_size as u32, + chunk_info.filter_mask, + )? } else { raw_chunk.to_vec() }; @@ -1335,11 +1346,13 @@ pub fn read_chunked_data_indexed( ensure_len(file_data, c_addr, size)?; let raw_chunk = &file_data[c_addr..c_addr + size]; let decompressed = if let Some(pl) = pipeline { - if *filter_mask == 0 { - decompress_chunk(raw_chunk, pl, chunk_total_bytes, elem_size as u32)? - } else { - raw_chunk.to_vec() - } + decompress_chunk_masked( + raw_chunk, + pl, + chunk_total_bytes, + elem_size as u32, + *filter_mask, + )? } else { raw_chunk.to_vec() }; diff --git a/crates/clawhdf5-format/src/filters.rs b/crates/clawhdf5-format/src/filters.rs index c10cea9..31c024a 100644 --- a/crates/clawhdf5-format/src/filters.rs +++ b/crates/clawhdf5-format/src/filters.rs @@ -19,33 +19,64 @@ pub(crate) const MAX_DECOMPRESS_SIZE: usize = 256 * 1024 * 1024; /// Apply a filter pipeline to decompress a chunk. /// Filters are applied in REVERSE order for decompression. +/// +/// Equivalent to [`decompress_chunk_masked`] with a filter mask of 0 (every +/// filter was applied when the chunk was written). pub fn decompress_chunk( compressed: &[u8], pipeline: &FilterPipeline, chunk_size: usize, element_size: u32, ) -> Result, FormatError> { - let mut data = compressed.to_vec(); + decompress_chunk_masked(compressed, pipeline, chunk_size, element_size, 0) +} - for filter in pipeline.filters.iter().rev() { +/// Whether bit `index` of a chunk's filter mask says filter `index` was +/// skipped when the chunk was written. +fn filter_skipped(filter_mask: u32, index: usize) -> bool { + index < 32 && filter_mask & (1u32 << index) != 0 +} + +/// Whether `filter_mask` says none of `pipeline`'s filters were applied, so +/// the stored bytes are the chunk itself. +pub fn all_filters_skipped(pipeline: &FilterPipeline, filter_mask: u32) -> bool { + (0..pipeline.filters.len()).all(|i| filter_skipped(filter_mask, i)) +} + +/// Decompress a chunk whose filter mask is `filter_mask`: bit *i* set means +/// filter *i* of the pipeline was not applied when the chunk was written (an +/// optional filter that declined, or a direct chunk write), so only that +/// filter is skipped here; the others are still undone, in reverse order. +/// +/// `chunk_size` is the chunk's decoded size (0 if unknown); it caps each +/// stage's output. +pub fn decompress_chunk_masked( + compressed: &[u8], + pipeline: &FilterPipeline, + chunk_size: usize, + element_size: u32, + filter_mask: u32, +) -> Result, FormatError> { + let mut data = compressed.to_vec(); + for (i, filter) in pipeline.filters.iter().enumerate().rev() { + if filter_skipped(filter_mask, i) { + continue; + } + let bound = chunk_size; data = match filter.filter_id { FILTER_SHUFFLE => shuffle_decompress(&data, element_size as usize)?, - // `chunk_size` is the expected decompressed size (shuffle/fletcher32 - // are size-preserving, so it bounds these too); pass it so these - // decoders can't be forced into unbounded allocation by a hostile - // or corrupted compressed payload. - FILTER_DEFLATE => deflate_decompress(&data, chunk_size)?, - FILTER_LZ4 => lz4_decompress(&data, chunk_size)?, - FILTER_ZSTD => zstd_decompress(&data, chunk_size)?, + // `bound` caps the decoded size so these decoders can't be forced + // into unbounded allocation by a hostile or corrupted payload. + FILTER_DEFLATE => deflate_decompress(&data, bound)?, + FILTER_LZ4 => lz4_decompress(&data, bound)?, + FILTER_ZSTD => zstd_decompress(&data, bound)?, FILTER_FLETCHER32 => fletcher32_verify(&data)?, - FILTER_PCODEC => pcodec_decompress(&data, element_size as usize, chunk_size)?, - // `chunk_size` is the expected decompressed size; pass it so these - // decoders can reject an element count that would over-allocate. - FILTER_SCALEOFFSET => scaleoffset_decompress(&data, &filter.client_data, chunk_size)?, - FILTER_NBIT => nbit_decompress(&data, &filter.client_data, chunk_size)?, - FILTER_SZIP => { - crate::filters_szip::szip_decompress(&data, &filter.client_data, chunk_size)? - } + FILTER_PCODEC => pcodec_decompress(&data, element_size as usize, bound)?, + // These decoders also reject an element count that would + // over-allocate past `bound`. + FILTER_SCALEOFFSET => scaleoffset_decompress(&data, &filter.client_data, bound)?, + FILTER_NBIT => nbit_decompress(&data, &filter.client_data, bound)?, + FILTER_SZIP => crate::filters_szip::szip_decompress(&data, &filter.client_data, bound)?, other => return Err(FormatError::UnsupportedFilter(other)), }; } @@ -1413,6 +1444,65 @@ mod tests { assert_eq!(decompressed, data); } + fn filter(filter_id: u16) -> FilterDescription { + FilterDescription { + filter_id, + name: None, + flags: 0, + client_data: vec![], + } + } + + #[test] + #[cfg(feature = "deflate")] + fn filter_mask_skips_only_the_masked_filters() { + let pipeline = FilterPipeline { + version: 2, + filters: vec![filter(FILTER_SHUFFLE), filter(FILTER_DEFLATE)], + }; + let only_shuffle = FilterPipeline { + version: 2, + filters: vec![filter(FILTER_SHUFFLE)], + }; + let only_deflate = FilterPipeline { + version: 2, + filters: vec![filter(FILTER_DEFLATE)], + }; + let data: Vec = (0..200).map(|i| (i * 7 % 256) as u8).collect(); + let n = data.len(); + + let shuffled = compress_chunk(&data, &only_shuffle, 8).unwrap(); + assert_ne!(shuffled, data); + let deflated = compress_chunk(&data, &only_deflate, 8).unwrap(); + + // Bit 1: deflate skipped, shuffle still undone. + assert_eq!( + decompress_chunk_masked(&shuffled, &pipeline, n, 8, 0b10).unwrap(), + data + ); + // Bit 0: shuffle skipped, deflate still undone. + assert_eq!( + decompress_chunk_masked(&deflated, &pipeline, n, 8, 0b01).unwrap(), + data + ); + // Both bits (and bits past the pipeline): stored as-is. + assert_eq!( + decompress_chunk_masked(&data, &pipeline, n, 8, u32::MAX).unwrap(), + data + ); + assert!(all_filters_skipped(&pipeline, 0b11)); + assert!(!all_filters_skipped(&pipeline, 0b10)); + // An unsupported filter is fine when the chunk skipped it. + let unknown = FilterPipeline { + version: 2, + filters: vec![filter(32000), filter(FILTER_DEFLATE)], + }; + assert_eq!( + decompress_chunk_masked(&deflated, &unknown, n, 8, 0b01).unwrap(), + data + ); + } + #[test] #[cfg(feature = "deflate")] fn pipeline_compress_decompress_roundtrip() { diff --git a/crates/clawhdf5-format/src/parallel_read.rs b/crates/clawhdf5-format/src/parallel_read.rs index 14f5613..0bb2785 100644 --- a/crates/clawhdf5-format/src/parallel_read.rs +++ b/crates/clawhdf5-format/src/parallel_read.rs @@ -10,7 +10,7 @@ use crate::chunked_read::ChunkInfo; use crate::error::FormatError; use crate::filter_pipeline::FilterPipeline; -use crate::filters::decompress_chunk; +use crate::filters::decompress_chunk_masked; use crate::lane_partition::{self, LaneStats, PartitionStats}; /// Threshold: only use parallel decompression when chunk count exceeds this. @@ -84,11 +84,13 @@ pub fn decompress_chunks_lane_partitioned( } let raw_chunk = &file_data[c_addr..c_addr + size]; - let decompressed = if chunk_info.filter_mask == 0 { - decompress_chunk(raw_chunk, pipeline, chunk_total_bytes, element_size)? - } else { - raw_chunk.to_vec() - }; + let decompressed = decompress_chunk_masked( + raw_chunk, + pipeline, + chunk_total_bytes, + element_size, + chunk_info.filter_mask, + )?; stats.chunks_processed += 1; stats.compressed_bytes += size as u64; @@ -158,11 +160,13 @@ pub fn decompress_chunks_parallel( } let raw_chunk = &file_data[c_addr..c_addr + size]; - let decompressed = if chunk_info.filter_mask == 0 { - decompress_chunk(raw_chunk, pipeline, chunk_total_bytes, element_size)? - } else { - raw_chunk.to_vec() - }; + let decompressed = decompress_chunk_masked( + raw_chunk, + pipeline, + chunk_total_bytes, + element_size, + chunk_info.filter_mask, + )?; Ok(DecompressedChunk { index, @@ -200,11 +204,13 @@ pub fn decompress_chunks_sequential( let raw_chunk = &file_data[c_addr..c_addr + size]; let decompressed = if let Some(pl) = pipeline { - if chunk_info.filter_mask == 0 { - decompress_chunk(raw_chunk, pl, chunk_total_bytes, element_size)? - } else { - raw_chunk.to_vec() - } + decompress_chunk_masked( + raw_chunk, + pl, + chunk_total_bytes, + element_size, + chunk_info.filter_mask, + )? } else { raw_chunk.to_vec() }; diff --git a/crates/clawhdf5-format/src/partial_read.rs b/crates/clawhdf5-format/src/partial_read.rs index 7cdd2de..5cc3aaa 100644 --- a/crates/clawhdf5-format/src/partial_read.rs +++ b/crates/clawhdf5-format/src/partial_read.rs @@ -22,7 +22,7 @@ use crate::data_read::extract_selection_from_buffer; use crate::dataspace::Dataspace; use crate::error::FormatError; use crate::filter_pipeline::FilterPipeline; -use crate::filters::decompress_chunk; +use crate::filters::{all_filters_skipped, decompress_chunk_masked}; use crate::selection::Selection; /// The smallest axis-aligned box containing every selected element, as @@ -325,12 +325,18 @@ pub fn read_selection( expected: at.saturating_add(chunk.chunk_size as usize), available: file_data.len(), })?; - // Mirrors the full-read path: a non-zero filter mask means the - // chunk was stored unfiltered. + // Mirrors the full-read path: filter-mask bit i set means + // filter i was not applied to this chunk. let decoded; let data: &[u8] = match pipeline { - Some(pl) if chunk.filter_mask == 0 => { - decoded = decompress_chunk(raw, pl, chunk_bytes, elem_size as u32)?; + Some(pl) if !all_filters_skipped(pl, chunk.filter_mask) => { + decoded = decompress_chunk_masked( + raw, + pl, + chunk_bytes, + elem_size as u32, + chunk.filter_mask, + )?; &decoded } _ => raw, diff --git a/crates/clawhdf5/tests/h5py_chunked_read_tests.rs b/crates/clawhdf5/tests/h5py_chunked_read_tests.rs index d46cff5..58fe290 100644 --- a/crates/clawhdf5/tests/h5py_chunked_read_tests.rs +++ b/crates/clawhdf5/tests/h5py_chunked_read_tests.rs @@ -8,6 +8,7 @@ use std::process::Command; use clawhdf5::File; +use clawhdf5_format::selection::Selection; /// The Python interpreter to drive interop checks with (see /// `h5py_interop_tests.rs`). @@ -103,3 +104,72 @@ for name, lengths in (("{p}", 4), ("{p}.l8", 8)): } } } + +// --------------------------------------------------------------------------- +// Per-chunk filter masks +// --------------------------------------------------------------------------- + +/// A chunk's filter mask has one bit per pipeline filter: bit i set means +/// filter i was not applied to that chunk. Any nonzero mask used to skip the +/// whole pipeline, so a chunk that skipped only gzip was handed back still +/// shuffled. +#[test] +fn h5py_partial_filter_mask_skips_only_masked_filters() { + skip_if_no_python!(); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("mask.h5"); + let p = path.display().to_string(); + run_python(&format!( + r#" +import h5py, numpy as np, zlib +def shuffle(b, es): + a = np.frombuffer(b, dtype=np.uint8).reshape(-1, es) + return a.T.copy().tobytes() +with h5py.File("{p}", "w") as f: + # shuffle (0) + gzip (1). Even chunks: both applied. Odd chunks: mask + # 0b10, gzip skipped, shuffle applied. Chunk 3: mask 0b11, raw. + ds = f.create_dataset("shuf_gzip", shape=(32,), chunks=(8,), dtype=" = (1000..1032).collect(); + assert_eq!(file.dataset("shuf_gzip").unwrap().read_i32().unwrap(), want); + let grid: Vec = (0..64).map(f64::from).collect(); + assert_eq!(file.dataset("grid").unwrap().read_f64().unwrap(), grid); + // The selection path decodes chunks on its own. + let part = file + .dataset("shuf_gzip") + .unwrap() + .read_selection(&Selection::Hyperslab { + start: vec![6], + stride: vec![1], + count: vec![1], + block: vec![20], + }) + .unwrap(); + let want: Vec = (1006..1026i32).flat_map(i32::to_le_bytes).collect(); + assert_eq!(part, want); +}