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:
co-authored by
Claude Fable 5.1
parent
d668e45ab5
commit
0addf328bc
@@ -374,6 +374,11 @@ impl ChunkCache {
|
||||
|
||||
// ----- Index operations -----
|
||||
|
||||
/// The most decompressed bytes this cache will hold.
|
||||
pub fn max_bytes(&self) -> usize {
|
||||
self.inner.lock().map(|g| g.max_bytes).unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Bind the cache to the dataset at chunk-index address `addr`.
|
||||
///
|
||||
/// The cache is shared per file across all of its datasets. If the cache
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -876,6 +876,30 @@ fn get_size(dt: &Datatype) -> usize {
|
||||
dt.type_size() as usize
|
||||
}
|
||||
|
||||
/// Reinterpret little-endian bytes as `count` native values of `T` on a
|
||||
/// little-endian target, in one copy.
|
||||
///
|
||||
/// The buffer is allocated uninitialised and filled by the copy. It used to be
|
||||
/// `vec![0; count]` first, which for a large dataset meant writing every page
|
||||
/// twice (zero it, then overwrite it) — about as expensive as the copy itself.
|
||||
#[cfg(target_endian = "little")]
|
||||
fn native_le_to_vec<T: Copy>(raw: &[u8], count: usize) -> Vec<T> {
|
||||
let bytes = count * core::mem::size_of::<T>();
|
||||
debug_assert!(bytes <= raw.len());
|
||||
let mut result: Vec<T> = Vec::with_capacity(count);
|
||||
// SAFETY: `result` has capacity for `count` values of `T`, i.e. `bytes`
|
||||
// bytes; `raw` holds at least `bytes` bytes (callers derive `count` from
|
||||
// `raw.len() / size_of::<T>()`); the regions cannot overlap because
|
||||
// `result` was just allocated. Every `T` used here (f32/f64/i32/i64) is
|
||||
// valid for any bit pattern, so after the copy all `count` values are
|
||||
// initialised and `set_len` is sound.
|
||||
unsafe {
|
||||
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr().cast::<u8>(), bytes);
|
||||
result.set_len(count);
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
/// Convert raw bytes to `f64` values.
|
||||
pub fn read_as_f64(raw: &[u8], datatype: &Datatype) -> Result<Vec<f64>, FormatError> {
|
||||
// Array datatypes (e.g. an array-typed compound member) are read as a flat
|
||||
@@ -903,14 +927,7 @@ pub fn read_as_f64(raw: &[u8], datatype: &Datatype) -> Result<Vec<f64>, FormatEr
|
||||
..
|
||||
}
|
||||
) {
|
||||
let mut result = vec![0.0f64; count];
|
||||
// SAFETY: On LE platforms, f64 in-memory representation matches LE bytes.
|
||||
// We copy raw bytes directly into the f64 buffer.
|
||||
// SAFETY: The byte slice is properly aligned for this type and the length is divisible by size_of::<T>().
|
||||
unsafe {
|
||||
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr() as *mut u8, raw.len());
|
||||
}
|
||||
return Ok(result);
|
||||
return Ok(native_le_to_vec::<f64>(raw, count));
|
||||
}
|
||||
|
||||
let order = get_byte_order(datatype);
|
||||
@@ -993,12 +1010,7 @@ pub fn read_as_i64(raw: &[u8], datatype: &Datatype) -> Result<Vec<i64>, FormatEr
|
||||
}
|
||||
)
|
||||
{
|
||||
let mut result = vec![0i64; count];
|
||||
// SAFETY: The byte slice is properly aligned for this type and the length is divisible by size_of::<T>().
|
||||
unsafe {
|
||||
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr() as *mut u8, raw.len());
|
||||
}
|
||||
return Ok(result);
|
||||
return Ok(native_le_to_vec::<i64>(raw, count));
|
||||
}
|
||||
|
||||
let order = get_byte_order(datatype);
|
||||
@@ -1062,12 +1074,7 @@ pub fn read_as_f32(raw: &[u8], datatype: &Datatype) -> Result<Vec<f32>, FormatEr
|
||||
..
|
||||
}
|
||||
) {
|
||||
let mut result = vec![0.0f32; count];
|
||||
// SAFETY: The byte slice is properly aligned for this type and the length is divisible by size_of::<T>().
|
||||
unsafe {
|
||||
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr() as *mut u8, raw.len());
|
||||
}
|
||||
return Ok(result);
|
||||
return Ok(native_le_to_vec::<f32>(raw, count));
|
||||
}
|
||||
|
||||
let order = get_byte_order(datatype);
|
||||
@@ -1144,12 +1151,7 @@ pub fn read_as_i32(raw: &[u8], datatype: &Datatype) -> Result<Vec<i32>, FormatEr
|
||||
}
|
||||
)
|
||||
{
|
||||
let mut result = vec![0i32; count];
|
||||
// SAFETY: The byte slice is properly aligned for this type and the length is divisible by size_of::<T>().
|
||||
unsafe {
|
||||
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr() as *mut u8, raw.len());
|
||||
}
|
||||
return Ok(result);
|
||||
return Ok(native_le_to_vec::<i32>(raw, count));
|
||||
}
|
||||
|
||||
let order = get_byte_order(datatype);
|
||||
|
||||
@@ -381,8 +381,13 @@ impl<'f> Dataset<'f> {
|
||||
|
||||
/// Read all data as `f64` values.
|
||||
pub fn read_f64(&self) -> Result<Vec<f64>, Error> {
|
||||
let raw = self.read_raw()?;
|
||||
let dt = self.datatype()?;
|
||||
// A contiguous dataset is converted straight from the file bytes; going
|
||||
// through `read_raw` first copied the whole dataset an extra time.
|
||||
if let Ok(Some(bytes)) = self.read_raw_ref() {
|
||||
return Ok(data_read::read_as_f64(bytes, &dt)?);
|
||||
}
|
||||
let raw = self.read_raw()?;
|
||||
Ok(data_read::read_as_f64(&raw, &dt)?)
|
||||
}
|
||||
|
||||
@@ -393,29 +398,49 @@ impl<'f> Dataset<'f> {
|
||||
///
|
||||
/// Read all data as `f32` values.
|
||||
pub fn read_f32(&self) -> Result<Vec<f32>, Error> {
|
||||
let raw = self.read_raw()?;
|
||||
let dt = self.datatype()?;
|
||||
// A contiguous dataset is converted straight from the file bytes; going
|
||||
// through `read_raw` first copied the whole dataset an extra time.
|
||||
if let Ok(Some(bytes)) = self.read_raw_ref() {
|
||||
return Ok(data_read::read_as_f32(bytes, &dt)?);
|
||||
}
|
||||
let raw = self.read_raw()?;
|
||||
Ok(data_read::read_as_f32(&raw, &dt)?)
|
||||
}
|
||||
|
||||
/// Read all data as `i32` values.
|
||||
pub fn read_i32(&self) -> Result<Vec<i32>, Error> {
|
||||
let raw = self.read_raw()?;
|
||||
let dt = self.datatype()?;
|
||||
// A contiguous dataset is converted straight from the file bytes; going
|
||||
// through `read_raw` first copied the whole dataset an extra time.
|
||||
if let Ok(Some(bytes)) = self.read_raw_ref() {
|
||||
return Ok(data_read::read_as_i32(bytes, &dt)?);
|
||||
}
|
||||
let raw = self.read_raw()?;
|
||||
Ok(data_read::read_as_i32(&raw, &dt)?)
|
||||
}
|
||||
|
||||
/// Read all data as `i64` values.
|
||||
pub fn read_i64(&self) -> Result<Vec<i64>, Error> {
|
||||
let raw = self.read_raw()?;
|
||||
let dt = self.datatype()?;
|
||||
// A contiguous dataset is converted straight from the file bytes; going
|
||||
// through `read_raw` first copied the whole dataset an extra time.
|
||||
if let Ok(Some(bytes)) = self.read_raw_ref() {
|
||||
return Ok(data_read::read_as_i64(bytes, &dt)?);
|
||||
}
|
||||
let raw = self.read_raw()?;
|
||||
Ok(data_read::read_as_i64(&raw, &dt)?)
|
||||
}
|
||||
|
||||
/// Read all data as `u64` values.
|
||||
pub fn read_u64(&self) -> Result<Vec<u64>, Error> {
|
||||
let raw = self.read_raw()?;
|
||||
let dt = self.datatype()?;
|
||||
// A contiguous dataset is converted straight from the file bytes; going
|
||||
// through `read_raw` first copied the whole dataset an extra time.
|
||||
if let Ok(Some(bytes)) = self.read_raw_ref() {
|
||||
return Ok(data_read::read_as_u64(bytes, &dt)?);
|
||||
}
|
||||
let raw = self.read_raw()?;
|
||||
Ok(data_read::read_as_u64(&raw, &dt)?)
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user