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