Skip to main content

array_write_read_string/
array_write_read_string.rs

1#![allow(missing_docs)]
2
3use itertools::Itertools;
4use ndarray::{Array2, ArrayD, array};
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::{ArrayBytes, data_type};
12    use zarrs::storage::store;
13
14    // Create a store
15    // let path = tempfile::TempDir::new()?;
16    // let mut store: ReadableWritableListableStorage =
17    //     Arc::new(zarrs::filesystem::FilesystemStore::new(path.path())?);
18    // let mut store: ReadableWritableListableStorage = Arc::new(
19    //     zarrs::filesystem::FilesystemStore::new("zarrs/tests/data/array_write_read.zarr")?,
20    // );
21    let mut store: ReadableWritableListableStorage = Arc::new(store::MemoryStore::new());
22    if let Some(arg1) = std::env::args().collect::<Vec<_>>().get(1)
23        && arg1 == "--usage-log"
24    {
25        let log_writer = Arc::new(std::sync::Mutex::new(
26            // std::io::BufWriter::new(
27            std::io::stdout(),
28            //    )
29        ));
30        store = Arc::new(UsageLogStorageAdapter::new(store, log_writer, || {
31            chrono::Utc::now().format("[%T%.3f] ").to_string()
32        }));
33    }
34
35    // Create the root group
36    zarrs::group::GroupBuilder::new()
37        .build(store.clone(), "/")?
38        .store_metadata()?;
39
40    // Create a group with attributes
41    let group_path = "/group";
42    let mut group = zarrs::group::GroupBuilder::new().build(store.clone(), group_path)?;
43    group
44        .attributes_mut()
45        .insert("foo".into(), serde_json::Value::String("bar".into()));
46    group.store_metadata()?;
47
48    println!(
49        "The group metadata is:\n{}\n",
50        group.metadata().to_string_pretty()
51    );
52
53    // Create an array
54    let array_path = "/group/array";
55    let array = zarrs::array::ArrayBuilder::new(
56        vec![4, 4], // array shape
57        vec![2, 2], // regular chunk shape
58        data_type::string(),
59        "_",
60    )
61    // .bytes_to_bytes_codecs(vec![]) // uncompressed
62    .dimension_names(["y", "x"].into())
63    // .storage_transformers(vec![].into())
64    .build(store.clone(), array_path)?;
65
66    // Write array metadata to store
67    array.store_metadata()?;
68
69    println!(
70        "The array metadata is:\n{}\n",
71        array.metadata().to_string_pretty()
72    );
73
74    // Write some chunks
75    array.store_chunk(
76        &[0, 0],
77        ArrayD::<&str>::from_shape_vec(vec![2, 2], vec!["a", "bb", "ccc", "dddd"]).unwrap(),
78    )?;
79    array.store_chunk(
80        &[0, 1],
81        ArrayD::<&str>::from_shape_vec(vec![2, 2], vec!["4444", "333", "22", "1"]).unwrap(),
82    )?;
83    let subset_all = array.subset_all();
84    let data_all: ArrayD<String> = array.retrieve_array_subset(&subset_all)?;
85    println!("store_chunk [0, 0] and [0, 1]:\n{data_all}\n");
86
87    // Write a subset spanning multiple chunks, including updating chunks already written
88    let ndarray_subset: Array2<&str> = array![["!", "@@"], ["###", "$$$$"]];
89    array.store_array_subset(&[1..3, 1..3], ndarray_subset)?;
90    let data_all: ArrayD<String> = array.retrieve_array_subset(&subset_all)?;
91    println!("store_array_subset [1..3, 1..3]:\nndarray::ArrayD<String>\n{data_all}");
92
93    // Retrieve bytes directly, convert into a single string allocation, create a &str ndarray
94    // TODO: Add a convenience function for this?
95    let data_all: ArrayBytes = array.retrieve_array_subset(&subset_all)?;
96    let (bytes, offsets) = data_all.into_variable()?.into_parts();
97    let string = String::from_utf8(bytes.into_owned())?;
98    let elements = offsets
99        .iter()
100        .tuple_windows()
101        .map(|(&curr, &next)| &string[curr..next])
102        .collect::<Vec<&str>>();
103    let ndarray = ArrayD::<&str>::from_shape_vec(subset_all.shape_usize(), elements)?;
104    println!("ndarray::ArrayD<&str>:\n{ndarray}");
105
106    Ok(())
107}
108
109fn main() {
110    if let Err(err) = array_write_read() {
111        println!("{err:?}");
112    }
113}