perf(format): parallel cached decode and fewer copies on full reads

Same-moment A/B on a 64 MB f64 dataset: chunked+deflate 110 -> 69 ms, chunked
72 -> 60 ms, contiguous 56 -> 30 ms.

- read_chunked_data_cached — the path the facade uses — decompressed chunks
  one at a time; only the uncached reader was parallel. Cache misses are now
  decoded in bounded batches (128), in parallel with the `parallel` feature.
- Every chunk was pushed into the 16 MiB chunk cache, which a larger dataset
  just churns (insert, evict moments later). Chunks are cached only when the
  whole dataset fits (new ChunkCache::max_bytes).
- Unfiltered chunks went file -> Vec -> aligned cache buffer -> output. They
  are copied straight from the file bytes.
- The facade's typed reads convert a contiguous dataset straight from the
  borrowed file bytes instead of copying it into a Vec first.
- The native little-endian fast paths allocated vec![0; n] and then overwrote
  it; they now fill an uninitialised buffer in one copy (native_le_to_vec).
  alloc_output requests zeroed memory from the allocator instead of reserving
  and filling.

The unit test that expected unfiltered chunks to land in the decompressed
cache now asserts the new design (index reused, cache not involved).

Co-Authored-By: Claude Fable 5.1 <[email protected]>
This commit is contained in:
osobh
2026-09-19 14:05:15 -07:00
co-authored by Claude Fable 5.1
parent d668e45ab5
commit 0addf328bc
6 changed files with 197 additions and 80 deletions
+106 -49
View File
@@ -165,12 +165,29 @@ pub(crate) fn checked_chunk_byte_len(
/// process when the allocation fails; a size taken from the file must surface
/// as an error instead.
pub(crate) fn alloc_output(len: usize) -> Result<Vec<u8>, FormatError> {
let mut out = Vec::new();
out.try_reserve_exact(len).map_err(|_| {
FormatError::Overflow(format!("cannot allocate {len} bytes for dataset output"))
})?;
out.resize(len, 0);
Ok(out)
if len == 0 {
return Ok(Vec::new());
}
let failed =
|| FormatError::Overflow(format!("cannot allocate {len} bytes for dataset output"));
let layout = core::alloc::Layout::array::<u8>(len).map_err(|_| failed())?;
// Ask the allocator for zeroed memory instead of reserving and then
// writing zeros: for a large buffer the OS hands out already-zero pages
// lazily, where an explicit fill touches every page up front — and most of
// the buffer is about to be overwritten with chunk data anyway.
//
// SAFETY (both arms): `layout` has non-zero size (len > 0) and alignment 1.
#[cfg(feature = "std")]
let ptr = unsafe { std::alloc::alloc_zeroed(layout) };
#[cfg(not(feature = "std"))]
let ptr = unsafe { alloc::alloc::alloc_zeroed(layout) };
if ptr.is_null() {
return Err(failed());
}
// SAFETY: `ptr` came from the global allocator with the layout of
// `[u8; len]`, which is exactly what `Vec<u8>` with capacity `len` frees;
// all `len` bytes are initialised (zero).
Ok(unsafe { Vec::from_raw_parts(ptr, len, len) })
}
fn read_offset(data: &[u8], pos: usize, size: u8) -> Result<u64, FormatError> {
@@ -362,6 +379,10 @@ pub fn generate_implicit_chunks(
}
/// Read a chunked dataset, decompressing chunks as needed.
/// Chunks decompressed together before being copied out, bounding the extra
/// memory a parallel full read holds at once.
const DECODE_BATCH: usize = 128;
/// B-tree v2 record types used for chunk indexing.
const BT2_CHUNK_UNFILTERED: u8 = 10;
const BT2_CHUNK_FILTERED: u8 = 11;
@@ -823,52 +844,86 @@ pub fn read_chunked_data_cached(
let chunk_total_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?;
for chunk_info in &chunks {
let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();
// Try decompressed cache first
let decompressed = if let Some(cached) = cache.get_decompressed_aligned(&coord) {
cached
} else {
// Decompress from file
let c_addr = chunk_info.address as usize;
let size = chunk_info.chunk_size as usize;
ensure_len(file_data, c_addr, size)?;
let raw_chunk = &file_data[c_addr..c_addr + size];
let dec = if let Some(pl) = pipeline {
if chunk_info.filter_mask == 0 {
decompress_chunk(raw_chunk, pl, chunk_total_bytes, elem_size as u32)?
} else {
raw_chunk.to_vec()
}
} else {
raw_chunk.to_vec()
};
cache.put_decompressed(coord, dec)
};
let mut place = |data: &[u8], chunk_info: &ChunkInfo| {
if rank == 0 {
let copy_len = data.len().min(output.len());
output[..copy_len].copy_from_slice(&data[..copy_len]);
return;
}
let chunk_offsets: Vec<usize> = chunk_info
.offsets
.iter()
.take(rank)
.map(|&o| o as usize)
.collect();
copy_chunk_to_output(
data,
&mut output,
&chunk_offsets,
&chunk_dims,
&ds_dims,
&ds_strides,
&chunk_strides,
elem_size,
rank,
);
};
let raw_bytes = |chunk_info: &ChunkInfo| -> Result<&[u8], FormatError> {
let c_addr = chunk_info.address as usize;
let size = chunk_info.chunk_size as usize;
ensure_len(file_data, c_addr, size)?;
Ok(&file_data[c_addr..c_addr + size])
};
if rank == 0 {
let copy_len = decompressed.len().min(output.len());
output[..copy_len].copy_from_slice(&decompressed[..copy_len]);
} else {
copy_chunk_to_output(
&decompressed,
&mut output,
&chunk_offsets,
&chunk_dims,
&ds_dims,
&ds_strides,
&chunk_strides,
elem_size,
rank,
);
// Chunks stored as-is (no pipeline, or the filter mask says this chunk
// skipped it) are copied straight from the file bytes: they are already in
// memory, so routing them through a Vec and then an aligned cache buffer
// was two extra copies of the whole dataset for nothing.
let stored_raw = |c: &ChunkInfo| pipeline.is_none() || c.filter_mask != 0;
let mut misses: Vec<&ChunkInfo> = Vec::new();
for chunk_info in &chunks {
if stored_raw(chunk_info) {
place(raw_bytes(chunk_info)?, chunk_info);
continue;
}
let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();
match cache.get_decompressed_aligned(&coord) {
Some(cached) => place(&cached, chunk_info),
None => misses.push(chunk_info),
}
}
// Decompress what the cache didn't have, a bounded batch at a time — in
// parallel with the `parallel` feature (this path, the one the facade
// uses, was sequential; only the uncached reader was parallel). Chunks are
// cached only when the whole dataset fits: pushing a larger dataset
// through the cache just evicts each chunk moments after inserting it.
let cache_them = total_bytes <= cache.max_bytes();
if let Some(pl) = pipeline {
let decode = |c: &&ChunkInfo| -> Result<Vec<u8>, FormatError> {
decompress_chunk(raw_bytes(c)?, pl, chunk_total_bytes, elem_size as u32)
};
for batch in misses.chunks(DECODE_BATCH) {
#[cfg(feature = "parallel")]
let decoded: Vec<Result<Vec<u8>, FormatError>> = if batch.len() >= 4 {
use rayon::prelude::*;
batch.par_iter().map(decode).collect()
} else {
batch.iter().map(decode).collect()
};
#[cfg(not(feature = "parallel"))]
let decoded: Vec<Result<Vec<u8>, FormatError>> = batch.iter().map(decode).collect();
for (chunk_info, data) in batch.iter().zip(decoded) {
let data = data?;
if cache_them {
let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();
let cached = cache.put_decompressed(coord, data);
place(&cached, chunk_info);
} else {
place(&data, chunk_info);
}
}
}
}
@@ -2186,21 +2241,23 @@ mod tests {
}
#[test]
fn cached_read_second_call_uses_cache() {
fn cached_read_second_call_reuses_the_index() {
let values: Vec<f64> = (0..20).map(|i| i as f64).collect();
let (file_data, layout, dataspace) = build_1d_chunked_file(&values, 10);
let datatype = make_f64_type();
let cache = ChunkCache::new();
// First read — populates index + decompressed cache
// First read — populates the chunk index. These chunks are stored
// unfiltered, so they are copied straight from the file bytes and the
// decompressed-chunk cache is (deliberately) not involved.
let raw1 = read_chunked_data_cached(
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
)
.unwrap();
assert!(cache.has_index());
assert!(cache.cached_chunk_count() > 0);
assert_eq!(cache.cached_chunk_count(), 0);
// Second read — should hit the decompressed cache
// Second read — reuses the cached index
let raw2 = read_chunked_data_cached(
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
)