data_beans/sparse_backend/
zarr.rs

1#![allow(dead_code)]
2
3use crate::sparse_io::*;
4use legume_numeric::matrix::common_io::*;
5use log::info;
6use std::ops::Range;
7use std::sync::{Arc, OnceLock};
8use zarrs::array::chunk_cache::ChunkCacheDecodedLruChunkLimit;
9use zarrs::array::{data_type, ArraySubset, DataType};
10use zarrs::filesystem::FilesystemStore;
11use zarrs::storage::ReadableListableStorageTraits as ZReadStorageTraits;
12
13/// Decoded-chunk LRU capacity, per zarr array. Keeps cache memory low
14/// per backend so it scales when many backends are loaded simultaneously
15/// (e.g. multi-sample workloads). With ~1 MiB/chunk, default 4 chunks per
16/// array × 4 arrays = ~16 MiB upper bound per backend. Override with the
17/// `LEGUME_ZARR_CACHE_CAP` env var when total backend count × default
18/// would exceed the memory budget.
19const DEFAULT_CACHE_CHUNK_CAP: u64 = 4;
20
21fn cache_chunk_cap() -> u64 {
22    static CAP: OnceLock<u64> = OnceLock::new();
23    *CAP.get_or_init(|| {
24        std::env::var("LEGUME_ZARR_CACHE_CAP")
25            .ok()
26            .and_then(|s| s.parse().ok())
27            .unwrap_or(DEFAULT_CACHE_CHUNK_CAP)
28    })
29}
30
31const KEY_BY_COLUMN_DATA: &str = "/by_column/data";
32const KEY_BY_COLUMN_INDICES: &str = "/by_column/indices";
33const KEY_BY_ROW_DATA: &str = "/by_row/data";
34const KEY_BY_ROW_INDICES: &str = "/by_row/indices";
35
36use anyhow::anyhow;
37
38use crate::sparse_backend::shared;
39use crate::utilities::io_helpers::{chunk_elems, parse_name_file};
40
41const COMPRESSION_LEVEL: i32 = 5;
42
43/// Block size (in elements) for streaming the CSC value/row-index arrays when
44/// exporting to Matrix Market. ~1M elements ≈ 4 MiB f32 + 8 MiB u64 transient
45/// per block — large enough to amortize per-retrieve overhead, small enough to
46/// bound memory regardless of matrix size.
47const MTX_STREAM_BLOCK: u64 = 1 << 20;
48
49/// 10x-like cell-feature matrix with `zarr` backend (feature x cell)
50///
51/// ```text
52/// (root)
53///     ├── nrow
54///     ├── ncell
55///     ├── by_column
56///     │   ├── data
57///     │   ├── indices (row indices)
58///     │   └── indptr (column pointers)
59///     └── by_row
60///         ├── data
61///         ├── indices (column indices)
62///         └── indptr (row pointers)
63/// ```
64///
65#[derive(Clone)]
66pub struct SparseMtxData {
67    read_store: Arc<dyn ZReadStorageTraits>,
68    write_store: Option<Arc<FilesystemStore>>,
69    file_name: String,
70    max_row_name_idx: usize,
71    max_column_name_idx: usize,
72    by_column_indptr: Vec<u64>,
73    /// Streaming-write cursor: entries appended so far (see `note_streamed_nnz`).
74    streamed_nnz: u64,
75    by_row_indptr: Vec<u64>,
76    by_column_indices: Option<Vec<u64>>,
77    by_column_data: Option<Vec<f32>>,
78    by_row_indices: Option<Vec<u64>>,
79    by_row_data: Option<Vec<f32>>,
80    /// Persistent decoded-chunk LRU caches. `Arc<OnceLock<_>>` so that
81    /// `Clone`s share state, and the underlying `moka::sync::Cache` inside
82    /// `ChunkCacheDecodedLruChunkLimit` is internally thread-safe — no
83    /// external locking needed.
84    by_column_data_cache: Arc<OnceLock<ChunkCacheDecodedLruChunkLimit>>,
85    by_column_indices_cache: Arc<OnceLock<ChunkCacheDecodedLruChunkLimit>>,
86    by_row_data_cache: Arc<OnceLock<ChunkCacheDecodedLruChunkLimit>>,
87    by_row_indices_cache: Arc<OnceLock<ChunkCacheDecodedLruChunkLimit>>,
88}
89
90impl SparseMtxData {
91    /// Get the writable store, or error if this is a read-only backend (e.g. zip archive).
92    fn write_store(&self) -> anyhow::Result<&Arc<FilesystemStore>> {
93        self.write_store
94            .as_ref()
95            .ok_or_else(|| anyhow!("store is read-only (zip archive)"))
96    }
97}
98
99impl SparseMtxData {
100    /// Create an empty new `SparseMtxData` instance with a zarr
101    /// backend file If no `backend_file` is provided, a temporary
102    /// file will be created.
103    ///
104    /// * `backend_file` - Optional zarr backend file
105    pub fn new(zarr_file: Option<&str>) -> anyhow::Result<Self> {
106        Self::create_backend(zarr_file)
107    }
108
109    /// Helper to create a backend file (from provided path or temp file)
110    fn create_backend(zarr_file: Option<&str>) -> anyhow::Result<Self> {
111        match zarr_file {
112            Some(backend_file) => Self::register_backend_file(backend_file),
113            None => {
114                let backend_file = create_temp_dir_file(".zarr")?;
115                let backend_file = backend_file
116                    .to_str()
117                    .ok_or_else(|| anyhow::anyhow!("Failed to convert path to string"))?;
118                Self::register_backend_file(backend_file)
119            }
120        }
121    }
122
123    /// Create `SparseMtxData` instance from an existing zarr backend file
124    /// * `zarr_file` - zarr backend file (directory or `.zarr.zip`)
125    pub fn open(backend_file: &str) -> anyhow::Result<Self> {
126        let (read_store, write_store) = crate::zarr_io::open_zarr_store_rw(backend_file)?;
127
128        if (
129            Self::_num_rows(read_store.clone()),
130            Self::_num_columns(read_store.clone()),
131            Self::_num_nnz(read_store.clone()),
132        ) == (None, None, None)
133        {
134            anyhow::bail!("Couldn't figure out the size of this sparse matrix data");
135        }
136
137        let mut ret = Self {
138            read_store,
139            write_store,
140            file_name: backend_file.to_string(),
141            max_row_name_idx: MAX_ROW_NAME_IDX,
142            max_column_name_idx: MAX_COLUMN_NAME_IDX,
143            by_column_indptr: vec![],
144            streamed_nnz: 0,
145            by_row_indptr: vec![],
146            by_column_indices: None,
147            by_column_data: None,
148            by_row_indices: None,
149            by_row_data: None,
150            by_column_data_cache: Arc::new(OnceLock::new()),
151            by_column_indices_cache: Arc::new(OnceLock::new()),
152            by_row_data_cache: Arc::new(OnceLock::new()),
153            by_row_indices_cache: Arc::new(OnceLock::new()),
154        };
155
156        ret.read_column_indptr()?;
157        ret.read_row_indptr()?;
158
159        Ok(ret)
160    }
161
162    /// Create `SparseMtxData` from mtx file with `backend_file` as
163    /// the backend file.  If no `backend_file` is provided, it will
164    /// be the same as `mtx_file` with `.zarr` extension.
165    /// * `mtx_file`: mtx file to be read into zarr backend
166    /// * `backend_file`: zarr file to be associated with
167    /// * `index_by_row`: if true, the matrix will be indexed by row
168    pub fn from_mtx_file(
169        mtx_file: &str,
170        backend_file: Option<&str>,
171        index_by_row: Option<bool>,
172    ) -> anyhow::Result<Self> {
173        let zarr_file = backend_file
174            .map(|s| s.to_string())
175            .unwrap_or_else(|| format!("{}.zarr", mtx_file));
176
177        info!("backend file: {}", zarr_file);
178        let mut ret = Self::register_backend_file(&zarr_file)?;
179
180        ret.import_mtx_file(mtx_file, index_by_row == Some(true))?;
181
182        info!("created sparse backend from {}", mtx_file);
183        Ok(ret)
184    }
185
186    /// Create a new `SparseMtxData` instance from an `ndarray` array
187    /// * `array` - 2D array to be added to the backend
188    /// * `backend_file` - Optional zarr backend file
189    /// * `index_by_row` - Optional flag to index by row (CSR format)
190    pub fn from_ndarray(
191        array: &Array2<f32>,
192        zarr_file: Option<&str>,
193        index_by_row: Option<bool>,
194    ) -> anyhow::Result<Self> {
195        let mut ret = Self::create_backend(zarr_file)?;
196
197        ret.import_ndarray_by_col(array)?;
198        ret.read_column_indptr()?;
199
200        if index_by_row == Some(true) {
201            ret.import_ndarray_by_row(array)?;
202            ret.read_row_indptr()?;
203        }
204        Ok(ret)
205    }
206
207    /// Create a new `SparseMtxData` instance from an `DMatrix` array
208    /// * `array` - 2D array to be added to the backend
209    /// * `backend_file` - Optional zarr backend file
210    /// * `index_by_row` - Optional flag to index by row (CSR format)
211    pub fn from_dmatrix(
212        matrix: &DMatrix<f32>,
213        zarr_file: Option<&str>,
214        index_by_row: Option<bool>,
215    ) -> anyhow::Result<Self> {
216        let mut ret = Self::create_backend(zarr_file)?;
217
218        ret.import_dmatrix_by_col(matrix)?;
219        ret.read_column_indptr()?;
220
221        if index_by_row == Some(true) {
222            ret.import_dmatrix_by_row(matrix)?;
223            ret.read_row_indptr()?;
224        }
225        Ok(ret)
226    }
227
228    /// Show the hierarchy of the zarr store
229    pub fn print_hierarchy(&self) -> anyhow::Result<()> {
230        use zarrs::config::MetadataRetrieveVersion;
231        let node =
232            zarrs::node::Node::open_opt(&self.read_store, "/", &MetadataRetrieveVersion::Default)?;
233        let tree = node.hierarchy_tree();
234        info!("hierarchy_tree:\n{}", tree);
235        Ok(())
236    }
237
238    /// Helper function to create a new zarr backend file
239    fn register_backend_file(zarr_file: &str) -> anyhow::Result<Self> {
240        use zarrs::group::GroupBuilder;
241        let store = Arc::new(FilesystemStore::new(zarr_file)?);
242        let root = GroupBuilder::new().build(store.clone(), "/")?;
243        root.store_metadata()?;
244
245        Ok(Self {
246            read_store: store.clone(),
247            write_store: Some(store),
248            file_name: zarr_file.to_string(),
249            max_row_name_idx: MAX_ROW_NAME_IDX,
250            max_column_name_idx: MAX_COLUMN_NAME_IDX,
251            by_column_indptr: vec![],
252            streamed_nnz: 0,
253            by_row_indptr: vec![],
254            by_column_indices: None,
255            by_column_data: None,
256            by_row_indices: None,
257            by_row_data: None,
258            by_column_data_cache: Arc::new(OnceLock::new()),
259            by_column_indices_cache: Arc::new(OnceLock::new()),
260            by_row_data_cache: Arc::new(OnceLock::new()),
261            by_row_indices_cache: Arc::new(OnceLock::new()),
262        })
263    }
264
265    /////////////////////
266    // backend related //
267    /////////////////////
268
269    /// Helper function to create a filled 1D array with the given
270    /// data type and fill value. This is the most useful function to
271    /// create a vector like data.
272    ///
273    /// * `key` - the key name
274    /// * `dt` - the data type among `DataType`
275    /// * `vec` - the vector to be stored
276    ///
277    fn new_filled_vector<V>(&mut self, key: &str, dt: DataType, vec: &[V]) -> anyhow::Result<()>
278    where
279        V: zarrs::array::Element,
280    {
281        use zarrs::array::codec::ZstdCodec;
282        use zarrs::array::ArrayBuilder;
283        use zarrs::array::FillValue;
284
285        let ws = self.write_store()?;
286
287        let nelem = vec.len();
288        let chunk_size = chunk_elems(nelem, std::mem::size_of::<V>());
289
290        let fill = if dt == data_type::float32() {
291            FillValue::from(zarrs::array::ZARR_NAN_F32)
292        } else if dt == data_type::uint64() {
293            FillValue::from(0u64)
294        } else if dt == data_type::string() {
295            FillValue::from("")
296        } else {
297            FillValue::from(0)
298        };
299
300        let array = ArrayBuilder::new(
301            vec![vec.len() as u64],  // array shape
302            vec![chunk_size as u64], // chunk shape
303            dt,                      // data type
304            fill,                    //
305        )
306        .bytes_to_bytes_codecs(vec![Arc::new(ZstdCodec::new(COMPRESSION_LEVEL, false))])
307        .build(ws.clone(), key)?;
308
309        array.store_metadata()?;
310
311        let subset = Self::create_subset(0..vec.len() as u64);
312        array.store_array_subset(&subset, vec)?;
313
314        Ok(())
315    }
316
317    fn _open_vector(
318        &self,
319        key: &str,
320    ) -> anyhow::Result<zarrs::array::Array<dyn ZReadStorageTraits>> {
321        use zarrs::array::Array as ZArray;
322        let ret = ZArray::open(self.read_store.clone(), key)?;
323        Ok(ret)
324    }
325
326    /// Create an empty fixed-shape 1-D array with the given data type,
327    /// chunk layout, and fill value. No data is written.
328    ///
329    /// Used by the streaming write path to pre-create `/by_column/*` and
330    /// `/by_row/*` arrays at their final size before any triplets land.
331    fn create_shaped_vector(
332        &mut self,
333        key: &str,
334        dt: DataType,
335        elem_bytes: usize,
336        nelem: usize,
337    ) -> anyhow::Result<()> {
338        use zarrs::array::codec::ZstdCodec;
339        use zarrs::array::ArrayBuilder;
340        use zarrs::array::FillValue;
341
342        let ws = self.write_store()?;
343
344        let chunk_size = chunk_elems(nelem, elem_bytes);
345
346        let fill = if dt == data_type::float32() {
347            FillValue::from(zarrs::array::ZARR_NAN_F32)
348        } else if dt == data_type::uint64() {
349            FillValue::from(0u64)
350        } else {
351            FillValue::from(0)
352        };
353
354        let array = ArrayBuilder::new(
355            vec![nelem.max(1) as u64],
356            vec![chunk_size.max(1) as u64],
357            dt,
358            fill,
359        )
360        .bytes_to_bytes_codecs(vec![Arc::new(ZstdCodec::new(COMPRESSION_LEVEL, false))])
361        .build(ws.clone(), key)?;
362
363        array.store_metadata()?;
364        Ok(())
365    }
366
367    /// Open an existing array through the writable filesystem store.
368    /// `_open_vector` uses the read-only handle, so it can't be used
369    /// to stage streaming writes.
370    fn _open_writable_vector(
371        &self,
372        key: &str,
373    ) -> anyhow::Result<zarrs::array::Array<FilesystemStore>> {
374        use zarrs::array::Array as ZArray;
375        let ws = self.write_store()?.clone();
376        let ret = ZArray::open(ws, key)?;
377        Ok(ret)
378    }
379
380    /// Write a `u64` slab at the given offset. The target array must
381    /// already exist (see [`create_shaped_vector`]).
382    fn write_slab_u64(&mut self, key: &str, offset: u64, data: &[u64]) -> anyhow::Result<()> {
383        if data.is_empty() {
384            return Ok(());
385        }
386        let array = self._open_writable_vector(key)?;
387        let subset = Self::create_subset(offset..offset + data.len() as u64);
388        array.store_array_subset(&subset, data)?;
389        Ok(())
390    }
391
392    /// Write an `f32` slab at the given offset.
393    fn write_slab_f32(&mut self, key: &str, offset: u64, data: &[f32]) -> anyhow::Result<()> {
394        if data.is_empty() {
395            return Ok(());
396        }
397        let array = self._open_writable_vector(key)?;
398        let subset = Self::create_subset(offset..offset + data.len() as u64);
399        array.store_array_subset(&subset, data)?;
400        Ok(())
401    }
402
403    #[allow(clippy::type_complexity)]
404    fn open_csc_triplets(
405        &self,
406    ) -> anyhow::Result<(
407        zarrs::array::Array<dyn ZReadStorageTraits>,
408        zarrs::array::Array<dyn ZReadStorageTraits>,
409        zarrs::array::Array<dyn ZReadStorageTraits>,
410    )> {
411        Ok((
412            self._open_vector("/by_column/indptr")?,
413            self._open_vector("/by_column/data")?,
414            self._open_vector("/by_column/indices")?,
415        ))
416    }
417
418    /// Helper to create an ArraySubset from a range
419    #[inline]
420    fn create_subset(range: Range<u64>) -> ArraySubset {
421        ArraySubset::new_with_ranges(&[range])
422    }
423
424    /// Build a decoded-chunk LRU cache for one zarr array. The cache
425    /// short-circuits redundant zstd decompression when the same chunk
426    /// is touched by multiple non-mergeable subset reads, both within
427    /// a single batch and across repeated calls (e.g. minibatch loops
428    /// over the same backend).
429    fn open_chunk_cache(
430        read_store: &Arc<dyn ZReadStorageTraits>,
431        key: &str,
432    ) -> anyhow::Result<ChunkCacheDecodedLruChunkLimit> {
433        use zarrs::array::Array as ZArray;
434        use zarrs::storage::ReadableStorageTraits;
435
436        let storage_readable: Arc<dyn ReadableStorageTraits> = read_store.clone().readable();
437        let arr = ZArray::open(read_store.clone(), key)?;
438        let arr_arc = Arc::new(arr.with_storage(storage_readable));
439        Ok(ChunkCacheDecodedLruChunkLimit::new(
440            arr_arc,
441            cache_chunk_cap(),
442        ))
443    }
444
445    /// Lazily build the persistent decoded-chunk cache. Concurrent
446    /// first-callers may both open the array, but only one `set` wins;
447    /// the loser's cache is dropped and `get` returns the winner.
448    fn cache_for<'a>(
449        &'a self,
450        cell: &'a OnceLock<ChunkCacheDecodedLruChunkLimit>,
451        key: &str,
452    ) -> anyhow::Result<&'a ChunkCacheDecodedLruChunkLimit> {
453        if let Some(cache) = cell.get() {
454            return Ok(cache);
455        }
456        let _ = cell.set(Self::open_chunk_cache(&self.read_store, key)?);
457        Ok(cell.get().expect("OnceLock populated above"))
458    }
459
460    fn _retrieve_vector<V>(&self, key: &str) -> anyhow::Result<Vec<V>>
461    where
462        V: zarrs::array::ElementOwned,
463    {
464        let data = self._open_vector(key)?;
465        let ntot = data.shape()[0];
466        let subset = Self::create_subset(0..ntot);
467        Ok(data.retrieve_array_subset::<Vec<V>>(&subset)?)
468    }
469
470    /////////////////////////////
471    // purely helper functions //
472    /////////////////////////////
473
474    /// Helper function to set an attribute from a group named `group_name`
475    fn _set_group_attr<V>(
476        store: Arc<FilesystemStore>,
477        group_name: &str,
478        attr_name: &str,
479        value: &V,
480    ) -> anyhow::Result<()>
481    where
482        V: serde::Serialize,
483    {
484        use zarrs::group::Group;
485        let mut group = Group::open(store, group_name)?;
486
487        let new_value = serde_json::to_value(value)?;
488        group
489            .attributes_mut()
490            .insert((*attr_name).to_string(), new_value);
491        group.store_metadata()?;
492        Ok(())
493    }
494
495    /// Helper function to get an attribute from a group named `group_name`
496    fn _get_group_attr<V>(
497        store: Arc<dyn ZReadStorageTraits>,
498        group_name: &str,
499        attr_name: &str,
500    ) -> Option<V>
501    where
502        V: serde::de::DeserializeOwned,
503    {
504        zarrs::group::Group::open(store, group_name)
505            .ok()
506            .and_then(|grp| grp.attributes().get(attr_name).cloned())
507            .and_then(|attr| serde_json::from_value(attr).ok())
508    }
509
510    fn _num_nnz(store: Arc<dyn ZReadStorageTraits>) -> Option<usize> {
511        Self::_get_group_attr::<usize>(store, "/", "nnz")
512    }
513
514    fn _num_rows(store: Arc<dyn ZReadStorageTraits>) -> Option<usize> {
515        Self::_get_group_attr::<usize>(store, "/", "nrow")
516    }
517
518    fn _num_columns(store: Arc<dyn ZReadStorageTraits>) -> Option<usize> {
519        Self::_get_group_attr::<usize>(store, "/", "ncol")
520    }
521
522    /// Helper function to add a group in the writable store
523    fn _add_group(&mut self, group_name: &str) -> anyhow::Result<()> {
524        use zarrs::group::Group;
525        let ws = self.write_store()?;
526
527        if Group::open(ws.clone(), group_name).is_err() {
528            let new_group = zarrs::group::GroupBuilder::new().build(ws.clone(), group_name)?;
529            new_group.store_metadata()?;
530        }
531
532        Ok(())
533    }
534}
535
536impl SparseIo for SparseMtxData {
537    type IndexIter = Vec<usize>;
538
539    /// Read row index pointers
540    fn read_row_indptr(&mut self) -> anyhow::Result<()> {
541        use zarrs::array::Array as Zarray;
542        let key = "/by_row/indptr";
543        if let Ok(indptr) = Zarray::open(self.read_store.clone(), key) {
544            let indptr_vec = indptr.retrieve_array_subset::<Vec<u64>>(&indptr.subset_all())?;
545            self.by_row_indptr.clear();
546            self.by_row_indptr.extend(indptr_vec);
547        }
548        Ok(())
549    }
550
551    /// Read column index pointers
552    fn column_indptr(&self) -> &[u64] {
553        &self.by_column_indptr
554    }
555
556    fn reopen_backend(&mut self) -> anyhow::Result<()> {
557        // Path-addressed store: rebuild it, refresh the resident indptrs — and
558        // DROP the decoded-chunk LRU caches, which pin arrays of the store they
559        // were built on. Keeping them served pre-swap chunk contents against
560        // post-swap indptrs: a matrix that read back with row indices past its
561        // own nrow, from a file that was byte-for-byte correct on disk.
562        let store = Arc::new(FilesystemStore::new(&self.file_name)?);
563        self.read_store = store.clone();
564        self.write_store = Some(store);
565        self.by_column_data_cache = Arc::new(OnceLock::new());
566        self.by_column_indices_cache = Arc::new(OnceLock::new());
567        self.by_row_data_cache = Arc::new(OnceLock::new());
568        self.by_row_indices_cache = Arc::new(OnceLock::new());
569        self.streamed_nnz = 0;
570        self.read_column_indptr()?;
571        self.read_row_indptr()?;
572        Ok(())
573    }
574
575    fn note_streamed_nnz(&mut self, n: u64) {
576        self.streamed_nnz += n;
577    }
578
579    fn streamed_nnz(&self) -> u64 {
580        self.streamed_nnz
581    }
582
583    fn reset_streamed_nnz(&mut self) {
584        self.streamed_nnz = 0;
585    }
586
587    fn read_column_indptr(&mut self) -> anyhow::Result<()> {
588        use zarrs::array::Array as ZArray;
589        let key = "/by_column/indptr";
590        if let Ok(indptr) = ZArray::open(self.read_store.clone(), key) {
591            let indptr_vec = indptr.retrieve_array_subset::<Vec<u64>>(&indptr.subset_all())?;
592            self.by_column_indptr.clear();
593            self.by_column_indptr.extend(indptr_vec);
594        }
595        Ok(())
596    }
597
598    fn clean_preloaded_columns(&mut self) {
599        self.by_column_data = None;
600        self.by_column_indices = None;
601    }
602
603    /// preload columns' values and indices
604    fn preload_columns(&mut self) -> anyhow::Result<()> {
605        if let Some(nnz) = self.num_non_zeros() {
606            if !crate::sparse_io::preload_within_budget(nnz, "column") {
607                return Ok(());
608            }
609        }
610        use zarrs::array::Array as ZArray;
611
612        let key = "/by_column/data";
613        let data = ZArray::open(self.read_store.clone(), key)?;
614        let key = "/by_column/indices";
615        let indices = ZArray::open(self.read_store.clone(), key)?;
616
617        let data = data.retrieve_array_subset::<Vec<f32>>(&data.subset_all())?;
618        let indices = indices.retrieve_array_subset::<Vec<u64>>(&indices.subset_all())?;
619
620        self.by_column_indices = Some(indices);
621        self.by_column_data = Some(data);
622        Ok(())
623    }
624
625    fn clean_preloaded_rows(&mut self) {
626        self.by_row_data = None;
627        self.by_row_indices = None;
628    }
629
630    /// preload rows' values and indices
631    fn preload_rows(&mut self) -> anyhow::Result<()> {
632        if let Some(nnz) = self.num_non_zeros() {
633            if !crate::sparse_io::preload_within_budget(nnz, "row") {
634                return Ok(());
635            }
636        }
637        use zarrs::array::Array as ZArray;
638
639        let data = ZArray::open(self.read_store.clone(), KEY_BY_ROW_DATA)?;
640        let indices = ZArray::open(self.read_store.clone(), KEY_BY_ROW_INDICES)?;
641
642        let data = data.retrieve_array_subset::<Vec<f32>>(&data.subset_all())?;
643        let indices = indices.retrieve_array_subset::<Vec<u64>>(&indices.subset_all())?;
644
645        self.by_row_indices = Some(indices);
646        self.by_row_data = Some(data);
647        Ok(())
648    }
649
650    /// Helper function to keep the matrix shape
651    fn record_mtx_shape(&mut self, mtx_shape: Option<(usize, usize, usize)>) -> anyhow::Result<()> {
652        if let Some((nrow, ncol, nnz)) = mtx_shape {
653            let ws = self.write_store()?;
654            let read_store = self.read_store.clone();
655
656            let check_set_attr = |attr_name: &str, value: usize| -> anyhow::Result<()> {
657                let old_value = Self::_get_group_attr::<usize>(read_store.clone(), "/", attr_name);
658                let new_value = serde_json::to_value(value)?;
659
660                match old_value {
661                    Some(old_value) => {
662                        if old_value != new_value {
663                            return Err(anyhow!("{} mismatch", attr_name));
664                        }
665                    }
666                    _ => {
667                        Self::_set_group_attr(ws.clone(), "/", attr_name, &new_value)?;
668                    }
669                }
670                Ok(())
671            };
672
673            check_set_attr("nrow", nrow)?;
674            check_set_attr("ncol", ncol)?;
675            check_set_attr("nnz", nnz)?;
676        }
677        Ok(())
678    }
679
680    /// Helper function to create a new zarr backend file
681    fn initialize_backend(&mut self) -> anyhow::Result<()> {
682        use zarrs::group::GroupBuilder;
683
684        self.remove_backend_file()?;
685        let zarr_file = &self.file_name;
686        let store = Arc::new(FilesystemStore::new(zarr_file)?);
687        let root = GroupBuilder::new().build(store.clone(), "/")?;
688        root.store_metadata()?;
689
690        self.read_store = store.clone();
691        self.write_store = Some(store);
692        self.file_name = zarr_file.to_string();
693        self.max_column_name_idx = MAX_COLUMN_NAME_IDX;
694        self.max_row_name_idx = MAX_ROW_NAME_IDX;
695        self.by_column_indptr = vec![];
696        self.by_row_indptr = vec![];
697
698        Ok(())
699    }
700
701    /// Clean up the backend file
702    fn remove_backend_file(&self) -> anyhow::Result<()> {
703        let backend = std::path::Path::new(&self.file_name);
704        if backend.exists() {
705            if backend.is_file() {
706                std::fs::remove_file(backend)?;
707            } else {
708                std::fs::remove_dir_all(backend)?;
709            }
710        }
711        Ok(())
712    }
713
714    /// Access file name of the zarr backend
715    fn get_backend_file_name(&self) -> &str {
716        &self.file_name
717    }
718
719    fn backend_type(&self) -> SparseIoBackend {
720        SparseIoBackend::Zarr
721    }
722
723    /// Export the data to a mtx file. This will take time.
724    /// * `mtx_file`: mtx file to be written
725    fn to_mtx_file(&self, mtx_file: &str) -> anyhow::Result<()> {
726        if let (Some(ncol), Some(nrow), Some(nnz)) =
727            (self.num_columns(), self.num_rows(), self.num_non_zeros())
728        {
729            let (nrow, ncol, nnz) = (nrow, ncol, nnz);
730
731            let mut buf = open_buf_writer(mtx_file)?;
732            shared::write_mtx_header(&mut buf, nrow, ncol, nnz)?;
733
734            let (indptr, data, indices) = self.open_csc_triplets()?;
735            let indptr = indptr.retrieve_array_subset::<Vec<u64>>(&indptr.subset_all())?;
736            debug_assert!(indptr.len() == ncol + 1);
737
738            // Stream the CSC value/row-index arrays in large sequential blocks
739            // (instead of two retrieves per column) and walk `indptr` to map
740            // each stored nonzero back to its column. CSC values are laid out
741            // column-major in column order, so emitting them in storage order
742            // reproduces the same column-by-column output as a per-column scan;
743            // empty columns are skipped by advancing the column pointer.
744            let total_nnz = indptr[ncol];
745            let mut jj = 0usize; // column owning the running nnz position
746            let mut pos = 0u64;
747            while pos < total_nnz {
748                let end = (pos + MTX_STREAM_BLOCK).min(total_nnz);
749                let subset = Self::create_subset(pos..end);
750                let data_block = data.retrieve_array_subset::<Vec<f32>>(&subset)?;
751                let indices_block = indices.retrieve_array_subset::<Vec<u64>>(&subset)?;
752
753                for (k, (&val, &ii)) in data_block.iter().zip(&indices_block).enumerate() {
754                    let global = pos + k as u64;
755                    // advance to the column owning this nonzero (skips empties)
756                    while jj + 1 < indptr.len() && indptr[jj + 1] <= global {
757                        jj += 1;
758                    }
759                    // 1-based indices
760                    writeln!(buf, "{}\t{}\t{}", ii as usize + 1, jj + 1, val)?;
761                }
762                pos = end;
763            }
764            buf.flush()?;
765            Ok(())
766        } else {
767            Err(anyhow!("Unable to figure out the size of the backend data"))
768        }
769    }
770
771    /// Set row names for the matrix
772    /// * `row_name_file`: a file each line contains row name words
773    fn register_row_names_file(&mut self, row_name_file: &str) {
774        let _ = self.register_names_file(
775            "/row_names",
776            row_name_file,
777            0..self.max_row_name_idx,
778            ROW_SEP,
779        );
780    }
781
782    /// Set row names for the matrix
783    /// * `rows`: a vector of row names
784    fn register_row_names_vec(&mut self, rows: &[Box<str>]) {
785        let _ = self.register_names_vec("/row_names", rows);
786    }
787
788    /// Set column names for the matrix
789    /// * `column_name_file`: a file each line contains column name words
790    fn register_column_names_file(&mut self, column_name_file: &str) {
791        let _ = self.register_names_file(
792            "/column_names",
793            column_name_file,
794            0..self.max_column_name_idx,
795            COLUMN_SEP,
796        );
797    }
798
799    /// Set column names for the matrix
800    /// * `columns`: a vector of column names
801    fn register_column_names_vec(&mut self, columns: &[Box<str>]) {
802        let _ = self.register_names_vec("/column_names", columns);
803    }
804
805    /// Number of rows in the matrix
806    fn num_rows(&self) -> Option<usize> {
807        Self::_num_rows(self.read_store.clone())
808    }
809
810    /// Number of columns in the matrix
811    fn num_columns(&self) -> Option<usize> {
812        Self::_num_columns(self.read_store.clone())
813    }
814
815    /// Number of non-zero elements in the matrix
816    fn num_non_zeros(&self) -> Option<usize> {
817        Self::_num_nnz(self.read_store.clone())
818    }
819
820    /// Add arbitrary names (a vector of strings)
821    /// * `group_name`: group name
822    /// * `name_file`: a file each line contains name words
823    /// * `name_columns`: range of columns to be used for name
824    /// * `name_sep`: separator for name columns
825    fn register_names_file(
826        &mut self,
827        key: &str,
828        name_file: &str,
829        name_columns: Range<usize>,
830        name_sep: &str,
831    ) -> anyhow::Result<()> {
832        let names = parse_name_file(name_file, name_columns, name_sep)?;
833        self.new_filled_vector(key, data_type::string(), &names)?;
834        Ok(())
835    }
836
837    /// Add arbitrary names (a vector of strings)
838    /// * `group_name`: group name
839    /// * `names`: a file each line contains name words
840    fn register_names_vec(&mut self, key: &str, names: &[Box<str>]) -> anyhow::Result<()> {
841        let names_vec: Vec<String> = names.iter().map(|x| x.to_string()).collect();
842        self.new_filled_vector(key, data_type::string(), &names_vec)?;
843        Ok(())
844    }
845
846    fn row_names(&self) -> anyhow::Result<Vec<Box<str>>> {
847        self.retrieve_registered_names("/row_names")
848    }
849
850    fn column_names(&self) -> anyhow::Result<Vec<Box<str>>> {
851        self.retrieve_registered_names("/column_names")
852    }
853
854    /// Get back the registered names
855    /// * `key`: key for the registered names
856    fn retrieve_registered_names(&self, key: &str) -> anyhow::Result<Vec<Box<str>>> {
857        Ok(self
858            ._retrieve_vector::<String>(key)?
859            .into_iter()
860            .map(|s| s.into_boxed_str())
861            .collect())
862    }
863
864    /// Read columns within the range and return a vector of triplets (row, col, value)
865    /// * `col` : usize
866    ///
867    fn read_triplets_by_single_column(
868        &self,
869        j_data: usize,
870    ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)> {
871        use zarrs::array::Array as ZArray;
872
873        debug_assert!(!self.by_column_indptr.is_empty()); // pre-loaded
874        debug_assert!(j_data < self.num_columns().unwrap_or(0)); //
875
876        let indptr = &self.by_column_indptr;
877
878        debug_assert!((j_data + 1) < indptr.len());
879        debug_assert!(indptr.len() > self.num_columns().unwrap_or(0));
880
881        let nrow = self
882            .num_rows()
883            .ok_or(anyhow!("can't figure out the number of rows"))?;
884
885        if let (Some(data), Some(indices)) = (&self.by_column_data, &self.by_column_indices) {
886            let ncol_out = 1;
887            let jj = 0;
888
889            // [start, end)
890            let start = indptr[j_data] as usize;
891            let end = indptr[j_data + 1] as usize;
892            let ret: Vec<(u64, u64, f32)> = indices[start..end]
893                .iter()
894                .zip(data[start..end].iter())
895                .map(|(&ii, &x_ij)| (ii, jj, x_ij))
896                .collect();
897
898            Ok((nrow, ncol_out, ret))
899        } else {
900            let key = "/by_column/data";
901            let data = ZArray::open(self.read_store.clone(), key)?;
902            let key = "/by_column/indices";
903            let indices = ZArray::open(self.read_store.clone(), key)?;
904
905            let ncol_out = 1;
906            let jj = 0;
907
908            // [start, end)
909            let start = indptr[j_data];
910            let end = indptr[j_data + 1];
911
912            let mut ret: Vec<(u64, u64, f32)> = Vec::with_capacity((end - start) as usize);
913
914            if start < end {
915                let subset = Self::create_subset(start..end);
916                let data_slice = data.retrieve_array_subset::<Vec<f32>>(&subset)?;
917                let indices_slice = indices.retrieve_array_subset::<Vec<u64>>(&subset)?;
918
919                for k in 0..(end - start) {
920                    let x_ij = data_slice[k as usize];
921                    let ii = indices_slice[k as usize];
922                    debug_assert!((ii as usize) < nrow);
923                    ret.push((ii, jj, x_ij));
924                }
925            }
926
927            Ok((nrow, ncol_out, ret))
928        }
929    }
930
931    /// Read columns within the range and return dense `ndarray::Array2`
932    /// * `columns` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
933    ///
934    fn read_triplets_by_columns(
935        &self,
936        columns: Self::IndexIter,
937    ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)> {
938        debug_assert!(!self.by_column_indptr.is_empty());
939        let indptr = &self.by_column_indptr;
940        let columns_vec = columns.into_iter().collect::<Vec<usize>>();
941
942        debug_assert!(indptr.len() > self.num_columns().unwrap_or(0));
943
944        let nrow = self
945            .num_rows()
946            .ok_or(anyhow!("can't figure out the number of rows"))?;
947
948        let ncol = self
949            .num_columns()
950            .ok_or(anyhow!("can't figure out the number of columns"))?;
951
952        let ncol_out = columns_vec.len();
953
954        if let (Some(data), Some(indices)) = (&self.by_column_data, &self.by_column_indices) {
955            let min_start = columns_vec
956                .iter()
957                .map(|&j_data| indptr[j_data])
958                .min()
959                .unwrap_or(0);
960
961            let max_end = columns_vec
962                .iter()
963                .map(|&j_data| indptr[j_data + 1])
964                .max()
965                .unwrap_or(0);
966
967            let mut ret: Vec<(u64, u64, f32)> = Vec::with_capacity((max_end - min_start) as usize);
968
969            for (jj, &j_data) in columns_vec.iter().enumerate() {
970                let jj = jj as u64;
971                let start = indptr[j_data] as usize;
972                let end = indptr[j_data + 1] as usize;
973                for (&ii, &x_ij) in indices[start..end].iter().zip(data[start..end].iter()) {
974                    ret.push((ii, jj, x_ij));
975                }
976            }
977
978            Ok((nrow, ncol_out, ret))
979        } else {
980            // CSC: tag = output column, inner = row. Sort by indptr.start so
981            // abutting/overlapping ranges fuse into one retrieve; across-chunk
982            // redundancy for non-mergeable ranges is caught by the chunk cache.
983            let mut tagged: Vec<(u64, u64, u64)> = columns_vec
984                .iter()
985                .enumerate()
986                .filter_map(|(jj, &j_data)| {
987                    if j_data >= ncol {
988                        return None;
989                    }
990                    let start = indptr[j_data];
991                    let end = indptr[j_data + 1];
992                    (start < end).then_some((jj as u64, start, end))
993                })
994                .collect();
995            tagged.sort_by_key(|&(_, start, _)| start);
996
997            let data_cache = self.cache_for(&self.by_column_data_cache, KEY_BY_COLUMN_DATA)?;
998            let indices_cache =
999                self.cache_for(&self.by_column_indices_cache, KEY_BY_COLUMN_INDICES)?;
1000
1001            let opts = zarrs::array::CodecOptions::default();
1002            let ret = shared::coalesce_and_emit(
1003                &tagged,
1004                nrow,
1005                |jj, ii, val| (ii, jj, val),
1006                |s, e| {
1007                    let subset = Self::create_subset(s..e);
1008                    let data_buf =
1009                        <_ as zarrs::array::chunk_cache::ChunkCache>::retrieve_array_subset::<
1010                            Vec<f32>,
1011                        >(data_cache, &subset, &opts)?;
1012                    let indices_buf =
1013                        <_ as zarrs::array::chunk_cache::ChunkCache>::retrieve_array_subset::<
1014                            Vec<u64>,
1015                        >(indices_cache, &subset, &opts)?;
1016                    Ok((data_buf, indices_buf))
1017                },
1018            )?;
1019            Ok((nrow, ncol_out, ret))
1020        }
1021    }
1022
1023    fn csc_column_arrays(&self) -> Option<(&[u64], &[u64], &[f32])> {
1024        match (
1025            self.by_column_data.as_ref(),
1026            self.by_column_indices.as_ref(),
1027        ) {
1028            (Some(data), Some(indices)) if !self.by_column_indptr.is_empty() => Some((
1029                self.by_column_indptr.as_slice(),
1030                indices.as_slice(),
1031                data.as_slice(),
1032            )),
1033            _ => None,
1034        }
1035    }
1036
1037    /// Read rows within the range and return a vector of triplets (row, col, value)
1038    /// * `rows` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
1039    ///
1040    fn read_triplets_by_rows(
1041        &self,
1042        rows: Self::IndexIter,
1043    ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)> {
1044        debug_assert!(!self.by_row_indptr.is_empty());
1045        let indptr = &self.by_row_indptr;
1046        debug_assert!(indptr.len() > self.num_rows().unwrap_or(0));
1047
1048        let rows_vec = rows.into_iter().collect::<Vec<_>>();
1049
1050        let (nrow, ncol) = match (self.num_rows(), self.num_columns()) {
1051            (Some(nrow), Some(ncol)) => (nrow, ncol),
1052            _ => return Err(anyhow!("Unable to figure out the size of the backend data")),
1053        };
1054        let nrow_out = rows_vec.len();
1055
1056        if let (Some(data), Some(indices)) = (&self.by_row_data, &self.by_row_indices) {
1057            let mut nnz_total: usize = 0;
1058            let valid: Vec<(u64, usize)> = rows_vec
1059                .iter()
1060                .enumerate()
1061                .filter_map(|(ii, &i_data)| {
1062                    if i_data >= nrow {
1063                        return None;
1064                    }
1065                    nnz_total += (indptr[i_data + 1] - indptr[i_data]) as usize;
1066                    Some((ii as u64, i_data))
1067                })
1068                .collect();
1069
1070            let mut ret: Vec<(u64, u64, f32)> = Vec::with_capacity(nnz_total);
1071            for (ii, i_data) in valid {
1072                let start = indptr[i_data] as usize;
1073                let end = indptr[i_data + 1] as usize;
1074                for (&jj, &x_ij) in indices[start..end].iter().zip(data[start..end].iter()) {
1075                    ret.push((ii, jj, x_ij));
1076                }
1077            }
1078            return Ok((nrow_out, ncol, ret));
1079        }
1080
1081        // CSR: tag = output row, inner = column.
1082        let mut tagged: Vec<(u64, u64, u64)> = rows_vec
1083            .iter()
1084            .enumerate()
1085            .filter_map(|(ii, &i_data)| {
1086                if i_data >= nrow {
1087                    return None;
1088                }
1089                debug_assert!((i_data + 1) < indptr.len());
1090                let start = indptr[i_data];
1091                let end = indptr[i_data + 1];
1092                (start < end).then_some((ii as u64, start, end))
1093            })
1094            .collect();
1095        tagged.sort_by_key(|&(_, start, _)| start);
1096
1097        let data_cache = self.cache_for(&self.by_row_data_cache, KEY_BY_ROW_DATA)?;
1098        let indices_cache = self.cache_for(&self.by_row_indices_cache, KEY_BY_ROW_INDICES)?;
1099
1100        let opts = zarrs::array::CodecOptions::default();
1101        let ret = shared::coalesce_and_emit(
1102            &tagged,
1103            ncol,
1104            |ii, jj, val| (ii, jj, val),
1105            |s, e| {
1106                let subset = Self::create_subset(s..e);
1107                let data_buf = <_ as zarrs::array::chunk_cache::ChunkCache>::retrieve_array_subset::<
1108                    Vec<f32>,
1109                >(data_cache, &subset, &opts)?;
1110                let indices_buf =
1111                    <_ as zarrs::array::chunk_cache::ChunkCache>::retrieve_array_subset::<Vec<u64>>(
1112                        indices_cache,
1113                        &subset,
1114                        &opts,
1115                    )?;
1116                Ok((data_buf, indices_buf))
1117            },
1118        )?;
1119        Ok((nrow_out, ncol, ret))
1120    }
1121    /// CSR data structure in Zarr backend
1122    ///
1123    /// ```text
1124    ///     └── by_row
1125    ///         ├── data
1126    ///         ├── indices (column indices)
1127    ///         └── isndptr (row pointers)
1128    /// ```
1129    fn record_csr_dataset_backend(
1130        &mut self,
1131        csr_cols: &[u64],
1132        csr_vals: &[f32],
1133        csr_rowptr: &[u64],
1134    ) -> anyhow::Result<()> {
1135        // open or create the group "/by_row"
1136        let key = "/by_row";
1137        self._add_group(key)?;
1138
1139        let key = "/by_row/data";
1140        self.new_filled_vector(key, data_type::float32(), csr_vals)?;
1141        let key = "/by_row/indices";
1142        self.new_filled_vector(key, data_type::uint64(), csr_cols)?;
1143        let key = "/by_row/indptr";
1144        self.new_filled_vector(key, data_type::uint64(), csr_rowptr)?;
1145
1146        Ok(())
1147    }
1148
1149    /// CSC data structure in Zarr backend
1150    ///
1151    /// ```text
1152    /// Helper function to record the CSC dataset
1153    ///     ├── by_column
1154    ///     │   ├── data
1155    ///     │   ├── indices (row indices)
1156    ///     │   └── indptr (column pointers)
1157    /// ```
1158    fn record_csc_dataset_backend(
1159        &mut self,
1160        csc_rows: &[u64],
1161        csc_vals: &[f32],
1162        csc_colptr: &[u64],
1163    ) -> anyhow::Result<()> {
1164        // open or create the group "/by_column"
1165        let key = "/by_column";
1166        self._add_group(key)?;
1167
1168        let key = "/by_column/data";
1169        self.new_filled_vector(key, data_type::float32(), csc_vals)?;
1170        let key = "/by_column/indices";
1171        self.new_filled_vector(key, data_type::uint64(), csc_rows)?;
1172        let key = "/by_column/indptr";
1173        self.new_filled_vector(key, data_type::uint64(), csc_colptr)?;
1174
1175        Ok(())
1176    }
1177
1178    fn cs_create(&mut self, key: CsKey, len: usize) -> anyhow::Result<()> {
1179        let (group, path, dt, elem_bytes) = match key {
1180            CsKey::CscData => (
1181                "/by_column",
1182                "/by_column/data",
1183                data_type::float32(),
1184                std::mem::size_of::<f32>(),
1185            ),
1186            CsKey::CscIndices => (
1187                "/by_column",
1188                "/by_column/indices",
1189                data_type::uint64(),
1190                std::mem::size_of::<u64>(),
1191            ),
1192            CsKey::CscIndptr => (
1193                "/by_column",
1194                "/by_column/indptr",
1195                data_type::uint64(),
1196                std::mem::size_of::<u64>(),
1197            ),
1198            CsKey::CsrData => (
1199                "/by_row",
1200                "/by_row/data",
1201                data_type::float32(),
1202                std::mem::size_of::<f32>(),
1203            ),
1204            CsKey::CsrIndices => (
1205                "/by_row",
1206                "/by_row/indices",
1207                data_type::uint64(),
1208                std::mem::size_of::<u64>(),
1209            ),
1210            CsKey::CsrIndptr => (
1211                "/by_row",
1212                "/by_row/indptr",
1213                data_type::uint64(),
1214                std::mem::size_of::<u64>(),
1215            ),
1216        };
1217        self._add_group(group)?;
1218        self.create_shaped_vector(path, dt, elem_bytes, len)
1219    }
1220
1221    fn cs_write_u64(&mut self, key: CsKey, offset: u64, data: &[u64]) -> anyhow::Result<()> {
1222        let path = match key {
1223            CsKey::CscIndices => "/by_column/indices",
1224            CsKey::CscIndptr => "/by_column/indptr",
1225            CsKey::CsrIndices => "/by_row/indices",
1226            CsKey::CsrIndptr => "/by_row/indptr",
1227            CsKey::CscData | CsKey::CsrData => {
1228                return Err(anyhow!("cs_write_u64 called on f32 slot {:?}", key));
1229            }
1230        };
1231        self.write_slab_u64(path, offset, data)
1232    }
1233
1234    fn cs_write_f32(&mut self, key: CsKey, offset: u64, data: &[f32]) -> anyhow::Result<()> {
1235        let path = match key {
1236            CsKey::CscData => "/by_column/data",
1237            CsKey::CsrData => "/by_row/data",
1238            _ => {
1239                return Err(anyhow!("cs_write_f32 called on u64 slot {:?}", key));
1240            }
1241        };
1242        self.write_slab_f32(path, offset, data)
1243    }
1244}
data_beans/sparse_backend/zarr.rs

data_beans/sparse_backend/
zarr.rs
