rectilinear_array_write_read/
rectilinear_array_write_read.rs

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

rectilinear_array_write_read/
rectilinear_array_write_read.rs
