genegraph_storage/zzarr/
mod.rs

1//! Zarr v3 ("zzarr") array IO over the [`zarrs`] crate (#3).
2//!
3//! Contract: Zarr v3 metadata (`zarr.json`), regular chunk grid, default
4//! chunk key encoding (`/` separator), `bytes` (little endian) plus
5//! optional `zstd` codecs. Unsupported encodings surface
6//! [`StorageError::UnsupportedFormat`] — never guessed.
7//!
8//! All operations are synchronous; async wrappers belong at the serving
9//! boundary (`spawn_blocking`), keeping this module runtime-free.
10
11pub mod csr;
12
13pub use zarrs::array::ElementOwned;
14
15use std::ops::Range;
16use std::path::Path;
17use std::sync::Arc;
18
19use zarrs::array::codec::ZstdCodec;
20use zarrs::array::{Array, data_type};
21use zarrs::array::{ArraySubset, DataType, Element, FillValue};
22use zarrs::filesystem::FilesystemStore;
23use zarrs::storage::ReadableWritableListableStorageTraits;
24
25use crate::{StorageError, StorageResult};
26
27/// Element types this module can store. Maps a Rust type to its Zarr v3
28/// data type and its zero fill value (the zzarr contract uses `0`).
29pub trait ZzarrElement: Element + ElementOwned + Clone {
30    /// The Zarr v3 data type.
31    fn zzarr_data_type() -> DataType;
32    /// The zero fill value.
33    fn zzarr_fill() -> FillValue;
34}
35
36macro_rules! zzarr_element {
37    ($t:ty, $dt:expr) => {
38        impl ZzarrElement for $t {
39            fn zzarr_data_type() -> DataType {
40                $dt()
41            }
42            fn zzarr_fill() -> FillValue {
43                FillValue::new((0 as $t).to_le_bytes().to_vec())
44            }
45        }
46    };
47}
48
49zzarr_element!(f32, data_type::float32);
50zzarr_element!(f64, data_type::float64);
51zzarr_element!(i64, data_type::int64);
52zzarr_element!(u64, data_type::uint64);
53
54fn store_for(path: &Path) -> StorageResult<Arc<dyn ReadableWritableListableStorageTraits>> {
55    let store: Arc<dyn ReadableWritableListableStorageTraits> =
56        Arc::new(FilesystemStore::new(path).map_err(|e| StorageError::Io(e.to_string()))?);
57    Ok(store)
58}
59
60fn map_open_err(err: impl std::fmt::Display) -> StorageError {
61    StorageError::UnsupportedFormat(format!("not a readable Zarr v3 array: {}", err))
62}
63
64fn map_err(err: impl std::fmt::Display) -> StorageError {
65    StorageError::Invalid(format!("zzarr error: {}", err))
66}
67
68/// A Zarr v3 array opened on the filesystem.
69pub struct ZarrArray {
70    array: Array<dyn ReadableWritableListableStorageTraits>,
71}
72
73impl std::fmt::Debug for ZarrArray {
74    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75        f.debug_struct("ZarrArray")
76            .field("shape", &self.array.shape())
77            .finish()
78    }
79}
80
81impl ZarrArray {
82    /// Array shape, outermost dimension first.
83    pub fn shape(&self) -> Vec<u64> {
84        self.array.shape().to_vec()
85    }
86
87    /// Read every element, row-major.
88    pub fn read_all<T: ElementOwned>(&self) -> StorageResult<Vec<T>> {
89        self.read_subset::<T>(
90            &(0..self.array.shape().len())
91                .map(|i| 0..self.array.shape()[i])
92                .collect::<Vec<_>>(),
93        )
94    }
95
96    /// Read a hyperrectangular subset, row-major.
97    pub fn read_subset<T: ElementOwned>(&self, ranges: &[Range<u64>]) -> StorageResult<Vec<T>> {
98        let subset = ArraySubset::new_with_ranges(ranges);
99        self.array
100            .retrieve_array_subset::<Vec<T>>(&subset)
101            .map_err(map_err)
102    }
103
104    /// Write a hyperrectangular subset, row-major (the counterpart of
105    /// [`Self::read_subset`]). Only the touched chunks are stored.
106    pub fn write_subset<T: Element>(
107        &mut self,
108        ranges: &[Range<u64>],
109        values: &[T],
110    ) -> StorageResult<()> {
111        let subset = ArraySubset::new_with_ranges(ranges);
112        self.array
113            .store_array_subset(&subset, values)
114            .map_err(map_err)
115    }
116
117    /// Append elements along the leading axis: resize, persist the new
118    /// shape, then write the tail. Values must tile the trailing extent.
119    ///
120    /// Accepted crash window (Python parity, #5 review): the resized shape
121    /// is published before the tail data lands, so a crash in between
122    /// leaves zero-filled phantom rows visible to readers.
123    pub fn append<T: Element>(&mut self, values: &[T]) -> StorageResult<()> {
124        let shape = self.array.shape().to_vec();
125        let trailing: u64 = shape[1..].iter().product::<u64>().max(1);
126        if shape.is_empty() || !(values.len() as u64).is_multiple_of(trailing) {
127            return Err(StorageError::Invalid(format!(
128                "append length {} does not tile trailing extent {}",
129                values.len(),
130                trailing
131            )));
132        }
133        let old_len = shape[0];
134        let added = values.len() as u64 / trailing;
135        let new_shape = shape
136            .iter()
137            .enumerate()
138            .map(|(i, s)| if i == 0 { old_len + added } else { *s })
139            .collect::<Vec<_>>();
140        self.array.set_shape(new_shape).map_err(map_err)?;
141        self.array.store_metadata().map_err(map_err)?;
142        let mut ranges: Vec<Range<u64>> = Vec::new();
143        ranges.push(old_len..old_len + added);
144        ranges.extend(shape[1..].iter().map(|s| 0..*s));
145        let subset = ArraySubset::new_with_ranges(&ranges);
146        self.array
147            .store_array_subset(&subset, values)
148            .map_err(map_err)
149    }
150}
151
152/// Open an existing Zarr v3 array directory.
153///
154/// Fails with [`StorageError::UnsupportedFormat`] when the directory does
155/// not carry `zarr.json` or carries unreadable metadata.
156pub fn open(path: &Path) -> StorageResult<ZarrArray> {
157    if !path.join("zarr.json").is_file() {
158        return Err(StorageError::UnsupportedFormat(format!(
159            "not a Zarr v3 array (missing zarr.json): {}",
160            path.display()
161        )));
162    }
163    let store = store_for(path)?;
164    let array = Array::open(store, "/").map_err(map_open_err)?;
165    Ok(ZarrArray { array })
166}
167
168/// Create a Zarr v3 array and write every element.
169///
170/// `compress` selects the `bytes`-only or `bytes`+`zstd` codec chain.
171/// Fails with [`StorageError::InvalidState`] when the path already holds
172/// an array.
173pub fn write_array<T: ZzarrElement>(
174    path: &Path,
175    shape: &[u64],
176    chunks: &[u64],
177    values: &[T],
178    compress: bool,
179) -> StorageResult<()> {
180    if path.join("zarr.json").is_file() {
181        return Err(StorageError::InvalidState(format!(
182            "zzarr array already exists at {}",
183            path.display()
184        )));
185    }
186    let total: u64 = shape.iter().product();
187    if total != values.len() as u64 {
188        return Err(StorageError::Invalid(format!(
189            "shape {:?} holds {} elements but {} values were given",
190            shape,
191            total,
192            values.len()
193        )));
194    }
195    std::fs::create_dir_all(path).map_err(|e| StorageError::Io(e.to_string()))?;
196    let store = store_for(path)?;
197    let mut builder = zarrs::array::ArrayBuilder::new(
198        shape.to_vec(),
199        chunks.to_vec(),
200        T::zzarr_data_type(),
201        T::zzarr_fill(),
202    );
203    if compress {
204        builder.bytes_to_bytes_codecs(vec![Arc::new(ZstdCodec::new(3, false))]);
205    }
206    let array = builder.build(store, "/").map_err(map_open_err)?;
207    array.store_metadata().map_err(map_err)?;
208    let subset_all = array.subset_all();
209    array
210        .store_array_subset(&subset_all, values)
211        .map_err(map_err)
212}
genegraph_storage/zzarr/mod.rs

genegraph_storage/zzarr/
mod.rs
