From e4de76f26c57e4603f7b648055e051a1c7412f41 Mon Sep 17 00:00:00 2001 From: "w.pomp" Date: Tue, 4 Aug 2026 18:29:56 +0200 Subject: [PATCH] - as_array_dyn lru caching --- .gitignore | 1 + src/view.rs | 333 ++++++++++++++++++++++++++++++++++++++++++++++------ 2 files changed, 296 insertions(+), 38 deletions(-) diff --git a/.gitignore b/.gitignore index c5bd81c..52e6cb8 100644 --- a/.gitignore +++ b/.gitignore @@ -75,6 +75,7 @@ docs/_build/ /tests/files/* AGENTS.md .agentbridge +.agent-work py/ndbioimage/jassets py/ndbioimage/deps \ No newline at end of file diff --git a/src/view.rs b/src/view.rs index 044bf6f..7e035b1 100644 --- a/src/view.rs +++ b/src/view.rs @@ -13,7 +13,7 @@ use num::traits::ToBytes; use num::{Bounded, FromPrimitive, ToPrimitive, Zero}; use serde::{Deserialize, Serialize}; use serde_with::serde_as; -use std::any::type_name; +use std::any::{Any, type_name}; use std::collections::{HashMap, HashSet}; use std::fmt::{Debug, Display, Formatter}; use std::hash::{Hash, Hasher}; @@ -48,11 +48,22 @@ fn slc_bnd(idx: isize, bnd: isize) -> Result { } pub trait Number: - 'static + AddAssign + Bounded + Clone + Div + FromPrimitive + PartialOrd + Zero + 'static + + Send + + Sync + + AddAssign + + Bounded + + Clone + + Div + + FromPrimitive + + PartialOrd + + Zero { } impl Number for T where T: 'static + + Send + + Sync + AddAssign + Bounded + Clone @@ -66,6 +77,9 @@ impl Number for T where /// maximum number of frames held in the cache const DEFAULT_FRAME_CACHE_SIZE: usize = 128; +/// maximum number of materialized arrays held in the cache +const DEFAULT_ARRAY_CACHE_SIZE: usize = 2; + /// identity of the reader a frame was read from #[derive(Debug, Clone, PartialEq, Eq, Hash)] struct ReaderKey { @@ -75,17 +89,52 @@ struct ReaderKey { position: usize, } -type FrameKey = (ReaderKey, usize, usize, usize); - -/// borrowed view of [`FrameKey`] for cache lookups without allocation. -/// hashes byte-identically to [`FrameKey`] (str/String and Path/PathBuf hash +/// borrowed view of [`ReaderKey`] for cache lookups without allocation. +/// hashes byte-identically to [`ReaderKey`] (str/String and Path/PathBuf hash /// the same), so it only matches the same reader identity. #[derive(Debug, PartialEq, Eq)] -struct FrameKeyRef<'a> { +struct ReaderKeyRef<'a> { name: &'a str, path: &'a Path, series: usize, position: usize, +} + +impl Hash for ReaderKeyRef<'_> { + fn hash(&self, state: &mut H) { + self.name.hash(state); + self.path.hash(state); + self.series.hash(state); + self.position.hash(state); + } +} + +impl Equivalent for ReaderKeyRef<'_> { + fn equivalent(&self, key: &ReaderKey) -> bool { + self.name == key.name + && self.path == key.path.as_path() + && self.series == key.series + && self.position == key.position + } +} + +impl ReaderKeyRef<'_> { + fn to_owned(&self) -> ReaderKey { + ReaderKey { + name: self.name.to_string(), + path: self.path.to_path_buf(), + series: self.series, + position: self.position, + } + } +} + +type FrameKey = (ReaderKey, usize, usize, usize); + +/// borrowed view of [`FrameKey`] for cache lookups without allocation. +#[derive(Debug, PartialEq, Eq)] +struct FrameKeyRef<'a> { + reader: ReaderKeyRef<'a>, c: usize, z: usize, t: usize, @@ -93,10 +142,7 @@ struct FrameKeyRef<'a> { impl Hash for FrameKeyRef<'_> { fn hash(&self, state: &mut H) { - self.name.hash(state); - self.path.hash(state); - self.series.hash(state); - self.position.hash(state); + self.reader.hash(state); self.c.hash(state); self.z.hash(state); self.t.hash(state); @@ -106,29 +152,78 @@ impl Hash for FrameKeyRef<'_> { impl Equivalent for FrameKeyRef<'_> { fn equivalent(&self, key: &FrameKey) -> bool { let (rk, c, z, t) = key; - self.name == rk.name - && self.path == rk.path.as_path() - && self.series == rk.series - && self.position == rk.position - && self.c == *c - && self.z == *z - && self.t == *t + self.reader.equivalent(rk) && self.c == *c && self.z == *z && self.t == *t } } impl FrameKeyRef<'_> { fn to_owned(&self) -> FrameKey { - ( - ReaderKey { - name: self.name.to_string(), - path: self.path.to_path_buf(), - series: self.series, - position: self.position, - }, - self.c, - self.z, - self.t, - ) + (self.reader.to_owned(), self.c, self.z, self.t) + } +} + +/// identity of a materialized array in the process-wide cache +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct ArrayKey { + reader: ReaderKey, + dtype: &'static str, + slice: Vec, + axes: Vec, + operations: Vec<(Axis, Operation)>, +} + +/// borrowed view of [`ArrayKey`] for cache lookups without allocation. +#[derive(Debug, PartialEq, Eq)] +struct ArrayKeyRef<'a> { + reader: ReaderKeyRef<'a>, + dtype: &'static str, + slice: &'a [SliceInfoElem], + axes: &'a [Axis], + operations: &'a IndexMap, +} + +impl Hash for ArrayKeyRef<'_> { + fn hash(&self, state: &mut H) { + self.reader.hash(state); + self.dtype.hash(state); + self.slice.hash(state); + self.axes.hash(state); + self.operations.len().hash(state); + for (ax, op) in self.operations.iter() { + ax.hash(state); + op.hash(state); + } + } +} + +impl Equivalent for ArrayKeyRef<'_> { + fn equivalent(&self, key: &ArrayKey) -> bool { + self.reader.equivalent(&key.reader) + && self.dtype == key.dtype + && self.slice == key.slice + && self.axes == key.axes + && self.operations.len() == key.operations.len() + && self + .operations + .iter() + .zip(&key.operations) + .all(|((ax, op), (key_ax, key_op))| ax == key_ax && op == key_op) + } +} + +impl ArrayKeyRef<'_> { + fn to_owned(&self) -> ArrayKey { + ArrayKey { + reader: self.reader.to_owned(), + dtype: self.dtype, + slice: self.slice.to_vec(), + axes: self.axes.to_vec(), + operations: self + .operations + .iter() + .map(|(ax, op)| (*ax, op.clone())) + .collect(), + } } } @@ -211,6 +306,86 @@ impl Debug for FrameCache { } } +/// process-wide LRU cache of materialized arrays, shared between all views and threads +static GLOBAL_ARRAY_CACHE: OnceLock = OnceLock::new(); + +/// thread-safe LRU cache of materialized arrays produced by `as_array_dyn` +struct ArrayCache { + inner: Mutex, +} + +struct ArrayCacheInner { + map: IndexMap>, + capacity: usize, +} + +impl Default for ArrayCache { + fn default() -> Self { + Self::new(DEFAULT_ARRAY_CACHE_SIZE) + } +} + +impl ArrayCache { + fn new(capacity: usize) -> Self { + Self { + inner: Mutex::new(ArrayCacheInner { + map: IndexMap::with_capacity(capacity), + capacity, + }), + } + } + + fn global() -> &'static ArrayCache { + GLOBAL_ARRAY_CACHE.get_or_init(ArrayCache::default) + } + + fn capacity(&self) -> usize { + self.inner.lock().unwrap().capacity + } + + fn set_capacity(&self, capacity: usize) { + let mut inner = self.inner.lock().unwrap(); + inner.capacity = capacity; + while inner.map.len() > inner.capacity { + inner.map.shift_remove_index(0); + } + } + + fn get(&self, key: &Q) -> Option>> + where + Q: ?Sized + Hash + Equivalent, + T: Any + Send + Sync, + { + let mut inner = self.inner.lock().unwrap(); + if let Some(idx) = inner.map.get_index_of(key) { + let (key, array) = inner.map.shift_remove_index(idx).unwrap(); + inner.map.insert(key, array.clone()); + array.downcast::>().ok() + } else { + None + } + } + + fn insert(&self, key: ArrayKey, array: ArrayD) { + let mut inner = self.inner.lock().unwrap(); + inner.map.insert(key, Arc::new(array)); + while inner.map.len() > inner.capacity { + inner.map.shift_remove_index(0); + } + } + + #[cfg(test)] + fn len(&self) -> usize { + self.inner.lock().unwrap().map.len() + } +} + +impl Debug for ArrayCache { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ArrayCache").finish_non_exhaustive() + } +} + /// sliceable view on an image file #[serde_as] #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] @@ -337,6 +512,18 @@ impl View { FrameCache::global().capacity() } + /// set the maximum number of materialized arrays the process-wide cache + /// may hold, evicting the least recently used arrays + pub fn with_array_cache_capacity(self, capacity: usize) -> Self { + ArrayCache::global().set_capacity(capacity); + self + } + + /// the maximum number of materialized arrays the process-wide cache may hold + pub fn array_cache_capacity(&self) -> usize { + ArrayCache::global().capacity() + } + fn with_operations(mut self, operations: IndexMap) -> Self { self.operations = operations; self @@ -799,6 +986,10 @@ impl View { Array1: MinMax>, Array2: MinMax>, { + let key = self.array_key::(); + if let Some(arr) = ArrayCache::global().get(&key) { + return Ok(arr.as_ref().clone()); + } let mut op_xy = IndexMap::new(); if let Some((&ax, op)) = self.operations.first() && ((ax == Axis::X) || (ax == Axis::Y)) @@ -954,7 +1145,7 @@ impl View { } } _ => { - panic!("xy cannot be 3d or more"); + unreachable!("xy cannot be 3d or more"); } }; if let Some((_, op)) = op_czt.first() { @@ -1020,6 +1211,7 @@ impl View { let m = T::from_usize(n).unwrap_or_else(|| T::zero()); out.take().unwrap().mapv(|x| x / m.clone()) }; + ArrayCache::global().insert(key.to_owned(), array.clone()); Ok(array) } @@ -1051,16 +1243,33 @@ impl View { fn frame_key(&self, c: usize, z: usize, t: usize) -> FrameKeyRef<'_> { FrameKeyRef { - name: self.reader.reader_name(), - path: self.reader.path(), - series: self.reader.series(), - position: self.reader.position(), + reader: ReaderKeyRef { + name: self.reader.reader_name(), + path: self.reader.path(), + series: self.reader.series(), + position: self.reader.position(), + }, c, z, t, } } + fn array_key(&self) -> ArrayKeyRef<'_> { + ArrayKeyRef { + reader: ReaderKeyRef { + name: self.reader.reader_name(), + path: self.reader.path(), + series: self.reader.series(), + position: self.reader.position(), + }, + dtype: type_name::(), + slice: &self.slice, + axes: &self.axes, + operations: &self.operations, + } + } + fn get_cached_frame(&self, c: usize, z: usize, t: usize) -> Result, Error> { let key = self.frame_key(c, z, t); if let Some(frame) = FrameCache::global().get(&key) { @@ -1297,14 +1506,17 @@ to_bytes_vec_impl!( #[cfg(test)] mod tests { - use crate::axes::Axis; + use crate::axes::{Axis, Operation}; use crate::error::Error; use crate::readers::{DynReader, Frame, Reader}; use crate::stats::MinMax; - use crate::view::{FrameCache, Item, ReaderKey}; + use crate::view::{ + ArrayCache, ArrayKey, ArrayKeyRef, FrameCache, Item, ReaderKey, ReaderKeyRef, + }; + use indexmap::IndexMap; use ndarray::{Array, Array4, Array5, NewAxis}; - use ndarray::{Array2, s}; - use std::path::PathBuf; + use ndarray::{Array2, ArrayD, IxDyn, SliceInfoElem, s}; + use std::path::{Path, PathBuf}; use std::sync::Arc; fn open(file: &str) -> Result { @@ -1610,6 +1822,51 @@ mod tests { Ok(()) } + #[test] + fn array_cache() -> Result<(), Error> { + let cache = ArrayCache::new(2); + let key = |dtype: &'static str, index: isize, axes: &[Axis]| ArrayKey { + reader: ReaderKey { + name: "test".to_string(), + path: PathBuf::from("test.tif"), + series: 0, + position: 0, + }, + dtype, + slice: vec![SliceInfoElem::Index(index)], + axes: axes.to_vec(), + operations: vec![], + }; + let k0 = key("u16", 0, &[Axis::T]); + let k1 = key("u16", 1, &[Axis::T]); + cache.insert(k0.clone(), ArrayD::::zeros(IxDyn(&[2, 2]))); + cache.insert(k1.clone(), ArrayD::::zeros(IxDyn(&[2, 2]))); + assert!(cache.get::(&k0).is_some()); + cache.insert( + key("u16", 2, &[Axis::T]), + ArrayD::::zeros(IxDyn(&[2, 2])), + ); + assert!(cache.get::(&k0).is_some()); + assert!(cache.get::(&k1).is_none()); + assert_eq!(cache.len(), 2); + let ops: IndexMap = IndexMap::new(); + let borrowed = ArrayKeyRef { + reader: ReaderKeyRef { + name: "test", + path: Path::new("test.tif"), + series: 0, + position: 0, + }, + dtype: "u16", + slice: &[SliceInfoElem::Index(0)], + axes: &[Axis::T], + operations: &ops, + }; + assert!(cache.get::(&borrowed).is_none()); + assert!(cache.get::(&borrowed).is_some()); + Ok(()) + } + #[test] fn as_array_caches_frames() -> Result<(), Error> { let file = "tiffseq/YTL1841B2-2-1_1hr_DMSO_galinduction_1";