Skip to main content

array_write_read/
array_write_read.rs

1#![allow(missing_docs)]
2
3use ndarray::ArrayD;
4use rayon::iter::{IntoParallelIterator, ParallelIterator};
5use zarrs::storage::ReadableWritableListableStorage;
6use zarrs::storage::storage_adapter::usage_log::UsageLogStorageAdapter;
7
8fn array_write_read() -> Result<(), Box<dyn std::error::Error>> {
9    use std::sync::Arc;
10
11    use zarrs::array::{ArraySubset, ZARR_NAN_F32, data_type};
12    use zarrs::node::Node;
13    use zarrs::storage::store;
14
15    // Create a store
16    // let path = tempfile::TempDir::new()?;
17    // let mut store: ReadableWritableListableStorage =
18    //     Arc::new(zarrs::filesystem::FilesystemStore::new(path.path())?);
19    // let mut store: ReadableWritableListableStorage = Arc::new(
20    //     zarrs::filesystem::FilesystemStore::new("zarrs/tests/data/array_write_read.zarr")?,
21    // );
22    let mut store: ReadableWritableListableStorage = Arc::new(store::MemoryStore::new());
23    if let Some(arg1) = std::env::args().collect::<Vec<_>>().get(1)
24        && arg1 == "--usage-log"
25    {
26        let log_writer = Arc::new(std::sync::Mutex::new(
27            // std::io::BufWriter::new(
28            std::io::stdout(),
29            //    )
30        ));
31        store = Arc::new(UsageLogStorageAdapter::new(store, log_writer, || {
32            chrono::Utc::now().format("[%T%.3f] ").to_string()
33        }));
34    }
35
36    // Create the root group
37    zarrs::group::GroupBuilder::new()
38        .build(store.clone(), "/")?
39        .store_metadata()?;
40
41    // Create a group with attributes
42    let group_path = "/group";
43    let mut group = zarrs::group::GroupBuilder::new().build(store.clone(), group_path)?;
44    group
45        .attributes_mut()
46        .insert("foo".into(), serde_json::Value::String("bar".into()));
47    group.store_metadata()?;
48
49    println!(
50        "The group metadata is:\n{}\n",
51        group.metadata().to_string_pretty()
52    );
53
54    // Create an array
55    let array_path = "/group/array";
56    let array = zarrs::array::ArrayBuilder::new(
57        vec![8, 8], // array shape
58        vec![4, 4], // regular chunk shape
59        data_type::float32(),
60        ZARR_NAN_F32,
61    )
62    // .bytes_to_bytes_codecs(vec![]) // uncompressed
63    .dimension_names(["y", "x"].into())
64    // .storage_transformers(vec![].into())
65    .build(store.clone(), array_path)?;
66
67    // Write array metadata to store
68    array.store_metadata()?;
69
70    println!(
71        "The array metadata is:\n{}\n",
72        array.metadata().to_string_pretty()
73    );
74
75    // Write some chunks
76    (0..2).into_par_iter().try_for_each(|i| {
77        let chunk_indices: Vec<u64> = vec![0, i];
78        let chunk_subset = array.chunk_grid().subset(&chunk_indices)?.ok_or_else(|| {
79            zarrs::array::ArrayError::InvalidChunkGridIndicesError(chunk_indices.to_vec())
80        })?;
81        array.store_chunk(
82            &chunk_indices,
83            vec![i as f32 * 0.1; chunk_subset.num_elements() as usize],
84        )
85    })?;
86
87    let subset_all = array.subset_all();
88    let data_all: ArrayD<f32> = array.retrieve_array_subset(&subset_all)?;
89    println!("store_chunk [0, 0] and [0, 1]:\n{data_all:+4.1}\n");
90
91    // Store multiple chunks
92    array.store_chunks(
93        &[1..2, 0..2],
94        &[
95            //
96            1.0f32, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1, 1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1,
97            //
98            1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1, 1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1,
99        ],
100    )?;
101    let data_all: ArrayD<f32> = array.retrieve_array_subset(&subset_all)?;
102    println!("store_chunks [1..2, 0..2]:\n{data_all:+4.1}\n");
103
104    // Write a subset spanning multiple chunks, including updating chunks already written
105    array.store_array_subset(
106        &[3..6, 3..6],
107        &[-3.3f32, -3.4, -3.5, -4.3, -4.4, -4.5, -5.3, -5.4, -5.5],
108    )?;
109    let data_all: ArrayD<f32> = array.retrieve_array_subset(&subset_all)?;
110    println!("store_array_subset [3..6, 3..6]:\n{data_all:+4.1}\n");
111
112    // Store array subset
113    array.store_array_subset(
114        &[0..8, 6..7],
115        &[-0.6f32, -1.6, -2.6, -3.6, -4.6, -5.6, -6.6, -7.6],
116    )?;
117    let data_all: ArrayD<f32> = array.retrieve_array_subset(&subset_all)?;
118    println!("store_array_subset [0..8, 6..7]:\n{data_all:+4.1}\n");
119
120    // Store chunk subset
121    array.store_chunk_subset(
122        // chunk indices
123        &[1, 1],
124        // subset within chunk
125        &[3..4, 0..4],
126        &[-7.4f32, -7.5, -7.6, -7.7],
127    )?;
128    let data_all: ArrayD<f32> = array.retrieve_array_subset(&subset_all)?;
129    println!("store_chunk_subset [3..4, 0..4] of chunk [1, 1]:\n{data_all:+4.1}\n");
130
131    // Erase a chunk
132    array.erase_chunk(&[0, 0])?;
133    let data_all: ArrayD<f32> = array.retrieve_array_subset(&subset_all)?;
134    println!("erase_chunk [0, 0]:\n{data_all:+4.1}\n");
135
136    // Read a chunk
137    let chunk_indices = vec![0, 1];
138    let data_chunk: ArrayD<f32> = array.retrieve_chunk(&chunk_indices)?;
139    println!("retrieve_chunk [0, 1]:\n{data_chunk:+4.1}\n");
140
141    // Read chunks
142    let chunks = ArraySubset::new_with_ranges(&[0..2, 1..2]);
143    let data_chunks: ArrayD<f32> = array.retrieve_chunks(&chunks)?;
144    println!("retrieve_chunks [0..2, 1..2]:\n{data_chunks:+4.1}\n");
145
146    // Retrieve an array subset
147    let subset = ArraySubset::new_with_ranges(&[2..6, 3..5]); // the center 4x2 region
148    let data_subset: ArrayD<f32> = array.retrieve_array_subset(&subset)?;
149    println!("retrieve_array_subset [2..6, 3..5]:\n{data_subset:+4.1}\n");
150
151    // Show the hierarchy
152    let node = Node::open(&store, "/").unwrap();
153    let tree = node.hierarchy_tree();
154    println!("hierarchy_tree:\n{}", tree);
155
156    Ok(())
157}
158
159fn main() {
160    if let Err(err) = array_write_read() {
161        println!("{:?}", err);
162    }
163}