rectangular_array_write_read/
rectangular_array_write_read.rs

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

rectangular_array_write_read/
rectangular_array_write_read.rs
