modewise_store_zarr/
lib.rs

1//! Stores artifacts on disk as Zarr v3, readable by zarr-python 3 and xarray.
2//!
3//! Each artifact is a group at `<first two key characters>/<key>/` holding its
4//! arrays (`data`, `modes` or `power`) and its metadata as group attributes.
5//! The group's metadata is written last, so an interrupted write never counts as stored.
6
7mod codec;
8
9use std::path::{Path, PathBuf};
10use std::sync::Arc;
11
12use modewise_pipeline::{Artifact, ArtifactKey, ArtifactStore, StoreError};
13use zarrs::filesystem::FilesystemStore;
14use zarrs::group::GroupBuilder;
15
16pub struct ZarrStore {
17    root: PathBuf,
18    storage: Arc<FilesystemStore>,
19}
20
21fn store_error(e: impl std::fmt::Display) -> StoreError {
22    StoreError(e.to_string())
23}
24
25impl ZarrStore {
26    /// Opens the store at `root`, creating it if needed.
27    pub fn open(root: &Path) -> Result<Self, StoreError> {
28        std::fs::create_dir_all(root).map_err(store_error)?;
29        let storage = Arc::new(FilesystemStore::new(root).map_err(store_error)?);
30        if !root.join("zarr.json").exists() {
31            let group = GroupBuilder::new()
32                .build(storage.clone(), "/")
33                .map_err(store_error)?;
34            group.store_metadata().map_err(store_error)?;
35        }
36        Ok(Self {
37            root: root.to_path_buf(),
38            storage,
39        })
40    }
41
42    pub fn root(&self) -> &Path {
43        &self.root
44    }
45
46    fn node(key: &ArtifactKey) -> String {
47        let key = key.as_str();
48        format!("/{}/{key}", &key[..2.min(key.len())])
49    }
50}
51
52impl ArtifactStore for ZarrStore {
53    fn get(&self, key: &ArtifactKey) -> Result<Option<Arc<Artifact>>, StoreError> {
54        if !self.contains(key)? {
55            return Ok(None);
56        }
57        codec::read(&self.storage, &Self::node(key)).map(|a| Some(Arc::new(a)))
58    }
59
60    fn put(&self, key: &ArtifactKey, artifact: Arc<Artifact>) -> Result<(), StoreError> {
61        codec::write(&self.storage, &Self::node(key), &artifact)
62    }
63
64    fn contains(&self, key: &ArtifactKey) -> Result<bool, StoreError> {
65        let group = Self::node(key);
66        Ok(self.root.join(&group[1..]).join("zarr.json").exists())
67    }
68}
modewise_store_zarr/lib.rs

modewise_store_zarr/
lib.rs
