sharded_array_write_read/
sharded_array_write_read.rs

1#![allow(missing_docs)]
2
3use itertools::Itertools;
4use ndarray::ArrayD;
5use zarrs::array::bytes_to_ndarray;
6use zarrs::storage::ReadableWritableListableStorage;
7use zarrs::storage::storage_adapter::usage_log::UsageLogStorageAdapter;
8use zarrs_codec::CodecOptions;
9
10fn sharded_array_write_read() -> Result<(), Box<dyn std::error::Error>> {
11    use std::sync::Arc;
12
13    use rayon::prelude::{IntoParallelIterator, ParallelIterator};
14    use zarrs::array::{ArraySubset, codec, data_type};
15    use zarrs::node::Node;
16    use zarrs::storage::store;
17
18    // Create a store
19    // let path = tempfile::TempDir::new()?;
20    // let mut store: ReadableWritableListableStorage =
21    //     Arc::new(zarrs::filesystem::FilesystemStore::new(path.path())?);
22    // let mut store: ReadableWritableListableStorage = Arc::new(
23    //     zarrs::filesystem::FilesystemStore::new("zarrs/tests/data/sharded_array_write_read.zarr")?,
24    // );
25    let mut store: ReadableWritableListableStorage = Arc::new(store::MemoryStore::new());
26    if let Some(arg1) = std::env::args().collect::<Vec<_>>().get(1)
27        && arg1 == "--usage-log"
28    {
29        let log_writer = Arc::new(std::sync::Mutex::new(
30            // std::io::BufWriter::new(
31            std::io::stdout(),
32            //    )
33        ));
34        store = Arc::new(UsageLogStorageAdapter::new(store, log_writer, || {
35            chrono::Utc::now().format("[%T%.3f] ").to_string()
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    // Create an array
53    let array_path = "/group/array";
54    let subchunk_shape = vec![4, 4];
55    let array = zarrs::array::ArrayBuilder::new(
56        vec![8, 8], // array shape
57        vec![4, 8], // chunk (shard) shape
58        data_type::uint16(),
59        0u16,
60    )
61    .subchunk_shape(subchunk_shape.clone())
62    .bytes_to_bytes_codecs(vec![
63        #[cfg(feature = "gzip")]
64        Arc::new(codec::GzipCodec::new(5)?),
65    ])
66    .dimension_names(["y", "x"].into())
67    // .storage_transformers(vec![].into())
68    .build(store.clone(), array_path)?;
69
70    // Write array metadata to store
71    array.store_metadata()?;
72
73    // The array metadata is
74    println!(
75        "The array metadata is:\n{}\n",
76        array.metadata().to_string_pretty()
77    );
78
79    // Use default codec options (concurrency etc)
80    let options = CodecOptions::default();
81
82    // Write some shards (in parallel)
83    (0..2).into_par_iter().try_for_each(|s| {
84        let chunk_grid = array.chunk_grid();
85        let chunk_indices = vec![s, 0];
86        if let Some(chunk_shape) = chunk_grid.chunk_shape(&chunk_indices)? {
87            let chunk_array = ndarray::ArrayD::<u16>::from_shape_fn(
88                chunk_shape
89                    .iter()
90                    .map(|u| u.get() as usize)
91                    .collect::<Vec<_>>(),
92                |ij| {
93                    (s * chunk_shape[0].get() * chunk_shape[1].get()
94                        + ij[0] as u64 * chunk_shape[1].get()
95                        + ij[1] as u64) as u16
96                },
97            );
98            array.store_chunk(&chunk_indices, chunk_array)
99        } else {
100            Err(zarrs::array::ArrayError::InvalidChunkGridIndicesError(
101                chunk_indices.to_vec(),
102            ))
103        }
104    })?;
105
106    // Read the whole array
107    let data_all: ArrayD<u16> = array.retrieve_array_subset(&array.subset_all())?;
108    println!("The whole array is:\n{data_all}\n");
109
110    // Read a shard back from the store
111    let shard_indices = vec![1, 0];
112    let data_shard: ArrayD<u16> = array.retrieve_chunk(&shard_indices)?;
113    println!("Shard [1,0] is:\n{data_shard}\n");
114
115    // Read a subchunk from the store
116    let subset_chunk_1_0 = ArraySubset::new_with_ranges(&[4..8, 0..4]);
117    let data_chunk: ArrayD<u16> = array.retrieve_array_subset(&subset_chunk_1_0)?;
118    println!("Chunk [1,0] is:\n{data_chunk}\n");
119
120    // Read the central 4x2 subset of the array
121    let subset_4x2 = ArraySubset::new_with_ranges(&[2..6, 3..5]); // the center 4x2 region
122    let data_4x2: ArrayD<u16> = array.retrieve_array_subset(&subset_4x2)?;
123    println!("The middle 4x2 subset is:\n{data_4x2}\n");
124
125    // Decode subchunks
126    // In some cases, it might be preferable to decode subchunks in a shard directly.
127    // If using the partial decoder, then the shard index will only be read once from the store.
128    let partial_decoder = array.partial_decoder(&[0, 0])?;
129    println!("Decoded subchunks:");
130    for subchunk_subset in [
131        ArraySubset::new_with_start_shape(vec![0, 0], subchunk_shape.clone())?,
132        ArraySubset::new_with_start_shape(vec![0, 4], subchunk_shape.clone())?,
133    ] {
134        println!("{subchunk_subset}");
135        let decoded_subchunk_bytes = partial_decoder.partial_decode(&subchunk_subset, &options)?;
136        let ndarray = bytes_to_ndarray::<u16>(
137            &subchunk_shape,
138            decoded_subchunk_bytes.into_fixed()?.into_owned(),
139        )?;
140        println!("{ndarray}\n");
141    }
142
143    // Show the hierarchy
144    let node = Node::open(&store, "/").unwrap();
145    let tree = node.hierarchy_tree();
146    println!("The Zarr hierarchy tree is:\n{}", tree);
147
148    println!(
149        "The keys in the store are:\n[{}]",
150        store.list().unwrap_or_default().iter().format(", ")
151    );
152
153    Ok(())
154}
155
156fn main() {
157    if let Err(err) = sharded_array_write_read() {
158        println!("{:?}", err);
159    }
160}
sharded_array_write_read/sharded_array_write_read.rs

sharded_array_write_read/
sharded_array_write_read.rs
