array_write_read_string/
array_write_read_string.rs

1#![allow(missing_docs)]
2
3use itertools::Itertools;
4use ndarray::{array, Array2, ArrayD};
5use zarrs::{
6    storage::storage_adapter::usage_log::UsageLogStorageAdapter,
7    storage::ReadableWritableListableStorage,
8};
9
10fn array_write_read() -> Result<(), Box<dyn std::error::Error>> {
11    use std::sync::Arc;
12    use zarrs::{
13        array::{DataType, FillValue},
14        array_subset::ArraySubset,
15        storage::store,
16    };
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/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        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![4, 4], // array shape
61        DataType::String,
62        vec![2, 2].try_into()?, // regular chunk shape
63        FillValue::from("_"),
64    )
65    // .bytes_to_bytes_codecs(vec![]) // uncompressed
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    println!(
74        "The array metadata is:\n{}\n",
75        array.metadata().to_string_pretty()
76    );
77
78    // Write some chunks
79    array.store_chunk_ndarray(
80        &[0, 0],
81        ArrayD::<&str>::from_shape_vec(vec![2, 2], vec!["a", "bb", "ccc", "dddd"]).unwrap(),
82    )?;
83    array.store_chunk_ndarray(
84        &[0, 1],
85        ArrayD::<&str>::from_shape_vec(vec![2, 2], vec!["4444", "333", "22", "1"]).unwrap(),
86    )?;
87    let subset_all = array.subset_all();
88    let data_all = array.retrieve_array_subset_ndarray::<String>(&subset_all)?;
89    println!("store_chunk [0, 0] and [0, 1]:\n{data_all}\n");
90
91    // Write a subset spanning multiple chunks, including updating chunks already written
92    let ndarray_subset: Array2<&str> = array![["!", "@@"], ["###", "$$$$"]];
93    array.store_array_subset_ndarray(
94        ArraySubset::new_with_ranges(&[1..3, 1..3]).start(),
95        ndarray_subset,
96    )?;
97    let data_all = array.retrieve_array_subset_ndarray::<String>(&subset_all)?;
98    println!("store_array_subset [1..3, 1..3]:\nndarray::ArrayD<String>\n{data_all}");
99
100    // Retrieve bytes directly, convert into a single string allocation, create a &str ndarray
101    // TODO: Add a convenience function for this?
102    let data_all = array.retrieve_array_subset(&subset_all)?;
103    let (bytes, offsets) = data_all.into_variable()?;
104    let string = String::from_utf8(bytes.into_owned())?;
105    let elements = offsets
106        .iter()
107        .tuple_windows()
108        .map(|(&curr, &next)| &string[curr..next])
109        .collect::<Vec<&str>>();
110    let ndarray = ArrayD::<&str>::from_shape_vec(subset_all.shape_usize(), elements)?;
111    println!("ndarray::ArrayD<&str>:\n{ndarray}");
112
113    Ok(())
114}
115
116fn main() {
117    if let Err(err) = array_write_read() {
118        println!("{:?}", err);
119    }
120}