Skip to main content

rectangular_array_write_read/
rectangular_array_write_read.rs

1#![allow(missing_docs)]
2
3use std::num::NonZeroU64;
4use std::sync::Arc;
5
6use ndarray::ArrayD;
7
8use zarrs::array::chunk_grid::RectangularChunkGridConfiguration;
9use zarrs::metadata::v3::MetadataV3;
10use zarrs::storage::ReadableWritableListableStorage;
11use zarrs::storage::storage_adapter::usage_log::UsageLogStorageAdapter;
12
13fn rectangular_array_write_read() -> Result<(), Box<dyn std::error::Error>> {
14    use rayon::prelude::{IntoParallelIterator, ParallelIterator};
15    use zarrs::array::{ArraySubset, ZARR_NAN_F32, codec, data_type};
16    use zarrs::node::Node;
17    use zarrs::storage::store;
18
19    // Create a store
20    // let path = tempfile::TempDir::new()?;
21    // let mut store: ReadableWritableListableStorage =
22    //     Arc::new(zarrs::filesystem::FilesystemStore::new(path.path())?);
23    let mut store: ReadableWritableListableStorage = Arc::new(store::MemoryStore::new());
24    if let Some(arg1) = std::env::args().collect::<Vec<_>>().get(1)
25        && arg1 == "--usage-log"
26    {
27        let log_writer = Arc::new(std::sync::Mutex::new(
28            // std::io::BufWriter::new(
29            std::io::stdout(),
30            //    )
31        ));
32        store = Arc::new(UsageLogStorageAdapter::new(store, log_writer, || {
33            chrono::Utc::now().format("[%T%.3f] ").to_string()
34        }));
35    }
36
37    // Create the root group
38    zarrs::group::GroupBuilder::new()
39        .build(store.clone(), "/")?
40        .store_metadata()?;
41
42    // Create a group with attributes
43    let group_path = "/group";
44    let mut group = zarrs::group::GroupBuilder::new().build(store.clone(), group_path)?;
45    group
46        .attributes_mut()
47        .insert("foo".into(), serde_json::Value::String("bar".into()));
48    group.store_metadata()?;
49
50    println!(
51        "The group metadata is:\n{}\n",
52        group.metadata().to_string_pretty()
53    );
54
55    // Create an array
56    let array_path = "/group/array";
57    let array = zarrs::array::ArrayBuilder::new(
58        vec![8, 8], // array shape
59        MetadataV3::new_with_configuration(
60            "rectangular",
61            RectangularChunkGridConfiguration {
62                chunk_shape: vec![
63                    vec![
64                        NonZeroU64::new(1).unwrap(),
65                        NonZeroU64::new(2).unwrap(),
66                        NonZeroU64::new(3).unwrap(),
67                        NonZeroU64::new(2).unwrap(),
68                    ]
69                    .into(),
70                    NonZeroU64::new(4).unwrap().into(),
71                ], // chunk sizes
72            },
73        ),
74        data_type::float32(),
75        ZARR_NAN_F32,
76    )
77    .bytes_to_bytes_codecs(vec![
78        #[cfg(feature = "gzip")]
79        Arc::new(codec::GzipCodec::new(5)?),
80    ])
81    .dimension_names(["y", "x"].into())
82    // .storage_transformers(vec![].into())
83    .build(store.clone(), array_path)?;
84
85    // Write array metadata to store
86    array.store_metadata()?;
87
88    // Write some chunks (in parallel)
89    (0..4).into_par_iter().try_for_each(|i| {
90        let chunk_grid = array.chunk_grid();
91        let chunk_indices = vec![i, 0];
92        if let Some(chunk_shape) = chunk_grid.chunk_shape(&chunk_indices)? {
93            let chunk_array = ndarray::ArrayD::<f32>::from_elem(
94                chunk_shape
95                    .iter()
96                    .map(|u| u.get() as usize)
97                    .collect::<Vec<_>>(),
98                i as f32,
99            );
100            array.store_chunk(&chunk_indices, chunk_array)
101        } else {
102            Err(zarrs::array::ArrayError::InvalidChunkGridIndicesError(
103                chunk_indices.to_vec(),
104            ))
105        }
106    })?;
107
108    println!(
109        "The array metadata is:\n{}\n",
110        array.metadata().to_string_pretty()
111    );
112
113    // Write a subset spanning multiple chunks, including updating chunks already written
114    array.store_array_subset(
115        &[3..6, 3..6], // start
116        ndarray::ArrayD::<f32>::from_shape_vec(
117            vec![3, 3],
118            vec![0.1f32, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
119        )?,
120    )?;
121
122    // Store elements directly, in this case set the 7th column to 123.0
123    array.store_array_subset(&[0..8, 6..7], &[123.0f32; 8])?;
124
125    // Store elements directly in a chunk, in this case set the last row of the bottom right chunk
126    array.store_chunk_subset(
127        // chunk indices
128        &[3, 1],
129        // subset within chunk
130        &[1..2, 0..4],
131        &[-4.0f32; 4],
132    )?;
133
134    // Read the whole array
135    let data_all: ArrayD<f32> = array.retrieve_array_subset(&array.subset_all())?;
136    println!("The whole array is:\n{data_all}\n");
137
138    // Read a chunk back from the store
139    let chunk_indices = vec![1, 0];
140    let data_chunk: ArrayD<f32> = array.retrieve_chunk(&chunk_indices)?;
141    println!("Chunk [1,0] is:\n{data_chunk}\n");
142
143    // Read the central 4x2 subset of the array
144    let subset_4x2 = ArraySubset::new_with_ranges(&[2..6, 3..5]); // the center 4x2 region
145    let data_4x2: ArrayD<f32> = array.retrieve_array_subset(&subset_4x2)?;
146    println!("The middle 4x2 subset is:\n{data_4x2}\n");
147
148    // Show the hierarchy
149    let node = Node::open(&store, "/").unwrap();
150    let tree = node.hierarchy_tree();
151    println!("The Zarr hierarchy tree is:\n{tree}");
152
153    Ok(())
154}
155
156fn main() {
157    if let Err(err) = rectangular_array_write_read() {
158        println!("{:?}", err);
159    }
160}