zarrs/
hierarchy.rs

1//! Zarr hierarchies.
2//!
3//! A Zarr hierarchy is a tree structure, where each node in the tree is either a [`Group`] or an [`Array`].
4//!
5//! A [`Hierarchy`] holds a mapping of [`NodePath`]s to [`NodeMetadata`].
6//!
7//! The [`Hierarchy::tree`] function can be used to create a string representation of the hierarchy.
8//!
9//! See <https://zarr-specs.readthedocs.io/en/latest/v3/core/index.html#hierarchy>.
10
11use std::collections::BTreeMap;
12use std::sync::Arc;
13
14use crate::array::{Array, ArrayMetadata};
15use crate::config::MetadataRetrieveVersion;
16use crate::group::Group;
17use crate::node::get_all_nodes_of;
18pub use crate::node::{Node, NodeCreateError, NodePath, NodePathError};
19#[cfg(feature = "async")]
20use crate::{
21    node::async_get_all_nodes_of,
22    storage::{AsyncListableStorageTraits, AsyncReadableStorageTraits},
23};
24pub use zarrs_metadata::NodeMetadata;
25use zarrs_storage::{ListableStorageTraits, ReadableStorageTraits};
26
27/// A Zarr hierarchy.
28#[derive(Debug, Clone)]
29pub struct Hierarchy(BTreeMap<NodePath, NodeMetadata>);
30
31impl std::fmt::Display for Hierarchy {
32    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33        write!(f, "{}", self.tree())
34    }
35}
36
37/// A hierarchy creation error.
38pub type HierarchyCreateError = NodeCreateError;
39
40impl Hierarchy {
41    /// Create a new, empty hierarchy.
42    fn new() -> Self {
43        Hierarchy(BTreeMap::new())
44    }
45
46    /// Open a hierarchy at `path` and read metadata and children from `storage` with default [`MetadataRetrieveVersion`].
47    ///
48    /// # Errors
49    /// Returns [`HierarchyCreateError`] if metadata is invalid or there is a failure to list child nodes.
50    pub fn open<TStorage: ?Sized + ReadableStorageTraits + ListableStorageTraits>(
51        storage: &Arc<TStorage>,
52        path: &str,
53    ) -> Result<Self, HierarchyCreateError> {
54        Self::open_opt(storage, path, &MetadataRetrieveVersion::Default)
55    }
56
57    /// Open a hierarchy at a `path` and read metadata and children from `storage` with non-default [`MetadataRetrieveVersion`].
58    ///
59    /// # Errors
60    /// Returns [`HierarchyCreateError`] if metadata is invalid or there is a failure to list child nodes.
61    pub fn open_opt<TStorage: ?Sized + ReadableStorageTraits + ListableStorageTraits>(
62        storage: &Arc<TStorage>,
63        path: &str,
64        version: &MetadataRetrieveVersion,
65    ) -> Result<Self, HierarchyCreateError> {
66        let node_path = NodePath::try_from(path)?;
67        let node_metadata = Node::get_metadata(storage, &node_path, version)?;
68        let mut hierarchy = Hierarchy::new();
69
70        let nodes = match node_metadata {
71            NodeMetadata::Array(_) => Vec::default(),
72            // TODO: Add consolidated metadata support
73            NodeMetadata::Group(_) => get_all_nodes_of(storage, &node_path, version)?,
74        };
75
76        hierarchy.0.insert(node_path, node_metadata);
77        hierarchy.0.extend(nodes);
78
79        Ok(hierarchy)
80    }
81
82    #[cfg(feature = "async")]
83    /// Asynchronously open a hierarchy at `path` and read metadata and children from `storage` with default [`MetadataRetrieveVersion`].
84    ///
85    /// # Errors
86    /// Returns [`HierarchyCreateError`] if metadata is invalid or there is a failure to list child nodes.
87    pub async fn async_open<
88        TStorage: ?Sized + AsyncReadableStorageTraits + AsyncListableStorageTraits,
89    >(
90        storage: &Arc<TStorage>,
91        path: &str,
92    ) -> Result<Self, HierarchyCreateError> {
93        Self::async_open_opt(storage, path, &MetadataRetrieveVersion::Default).await
94    }
95
96    #[cfg(feature = "async")]
97    /// Asynchronously open a hierarchy at a `path` and read metadata and children from `storage` with non-default [`MetadataRetrieveVersion`].
98    ///
99    /// # Errors
100    /// Returns [`HierarchyCreateError`] if metadata is invalid or there is a failure to list child nodes.
101    pub async fn async_open_opt<
102        TStorage: ?Sized + AsyncReadableStorageTraits + AsyncListableStorageTraits,
103    >(
104        storage: &Arc<TStorage>,
105        path: &str,
106        version: &MetadataRetrieveVersion,
107    ) -> Result<Self, HierarchyCreateError> {
108        let node_path = NodePath::try_from(path)?;
109        let node_metadata = Node::async_get_metadata(storage, &node_path, version).await?;
110        let mut hierarchy = Hierarchy::new();
111
112        let nodes = match node_metadata {
113            NodeMetadata::Array(_) => Vec::default(),
114            // TODO: Add consolidated metadata support
115            NodeMetadata::Group(_) => async_get_all_nodes_of(storage, &node_path, version).await?,
116        };
117
118        hierarchy.0.insert(node_path, node_metadata);
119        hierarchy.0.extend(nodes);
120
121        Ok(hierarchy)
122    }
123
124    /// Convenience method to create a `Hierarchy` from a `Group` with synchronous storage.
125    ///
126    /// # Errors
127    /// Returns [`HierarchyCreateError`] if group metadata is invalid or there is a failure to list child nodes.
128    pub fn try_from_group<TStorage: ?Sized + ReadableStorageTraits + ListableStorageTraits>(
129        group: &Group<TStorage>,
130    ) -> Result<Self, HierarchyCreateError> {
131        let mut hierarchy = Hierarchy::new();
132        hierarchy.0.insert(
133            group.path().clone(),
134            NodeMetadata::Group(group.metadata().clone()),
135        );
136        hierarchy.0.extend(get_all_nodes_of(
137            &group.storage(),
138            group.path(),
139            &MetadataRetrieveVersion::Default,
140        )?);
141        Ok(hierarchy)
142    }
143
144    #[cfg(feature = "async")]
145    /// Convenience method to create a `Hierarchy` from a Group with asynchronous storage.
146    ///
147    /// # Errors
148    /// Returns [`HierarchyCreateError`] if group metadata is invalid or there is a failure to list child nodes.
149    pub async fn try_from_async_group<
150        TStorage: ?Sized + AsyncReadableStorageTraits + AsyncListableStorageTraits,
151    >(
152        group: &Group<TStorage>,
153    ) -> Result<Hierarchy, HierarchyCreateError> {
154        let mut hierarchy = Hierarchy::new();
155        hierarchy.0.insert(
156            group.path().clone(),
157            NodeMetadata::Group(group.metadata().clone()),
158        );
159        hierarchy.0.extend(
160            async_get_all_nodes_of(
161                &group.storage(),
162                group.path(),
163                &MetadataRetrieveVersion::Default,
164            )
165            .await?,
166        );
167        Ok(hierarchy)
168    }
169
170    // /// Insert a node into the hierarchy.
171    // pub fn insert(&mut self, path: NodePath, metadata: NodeMetadata) -> Option<NodeMetadata> {
172    //     self.0.insert(path, metadata)
173    // }
174
175    /// Create a string representation of the hierarchy starting from the root.
176    #[must_use]
177    pub fn tree(&self) -> String {
178        self.tree_of(&NodePath::root())
179    }
180
181    /// Create a string representation of the hierarchy starting from `path`.
182    #[must_use]
183    pub fn tree_of(&self, path: &NodePath) -> String {
184        fn print_metadata(name: &str, string: &mut String, metadata: &NodeMetadata) {
185            match metadata {
186                NodeMetadata::Array(array_metadata) => {
187                    let s = match array_metadata {
188                        ArrayMetadata::V3(array_metadata) => {
189                            format!(
190                                "{} {:?} {}",
191                                name, array_metadata.shape, array_metadata.data_type
192                            )
193                        }
194                        ArrayMetadata::V2(array_metadata) => {
195                            format!(
196                                "{} {:?} {:?}",
197                                name, array_metadata.shape, array_metadata.dtype
198                            )
199                        }
200                    };
201                    string.push_str(&s);
202                }
203                NodeMetadata::Group(_) => {
204                    string.push_str(name);
205                }
206            }
207            string.push('\n');
208        }
209
210        let mut s = String::from(path.as_str());
211        s.push('\n');
212
213        let prefix = path.as_str();
214        let depth = path.as_path().components().count();
215
216        for node in self
217            .0
218            .iter()
219            .filter(|(path, _)| path.as_str().starts_with(prefix) && !path.as_str().eq(prefix))
220            .map(|(p, md)| Node::new_with_metadata(p.clone(), md.clone(), vec![]))
221        {
222            let depth = node
223                .path()
224                .as_path()
225                .components()
226                .count()
227                .saturating_sub(depth);
228
229            s.push_str(&" ".repeat(depth * 2));
230            print_metadata(node.name().as_str(), &mut s, node.metadata());
231        }
232        s
233    }
234}
235
236// impl Extend<(NodePath, NodeMetadata)> for Hierarchy {
237//     fn extend<T: IntoIterator<Item = (NodePath, NodeMetadata)>>(&mut self, iter: T) {
238//         for (path, metadata) in iter {
239//             self.insert(path, metadata);
240//         }
241//     }
242// }
243
244impl<TStorage: ?Sized> TryFrom<&Array<TStorage>> for Hierarchy {
245    type Error = HierarchyCreateError;
246    fn try_from(array: &Array<TStorage>) -> Result<Self, Self::Error> {
247        let mut hierarchy = Hierarchy::new();
248        hierarchy.0.insert(
249            array.path().clone(),
250            NodeMetadata::Array(array.metadata().clone()),
251        );
252        Ok(hierarchy)
253    }
254}
255
256impl<TStorage: ?Sized> TryFrom<Array<TStorage>> for Hierarchy {
257    type Error = HierarchyCreateError;
258    fn try_from(array: Array<TStorage>) -> Result<Self, Self::Error> {
259        (&array).try_into()
260    }
261}
262
263#[cfg(test)]
264mod tests {
265    use std::num::NonZeroU64;
266
267    use super::*;
268    use crate::array::ArrayBuilder;
269    use crate::group::GroupBuilder;
270    use zarrs_metadata::GroupMetadata;
271    use zarrs_metadata::v2::{ArrayMetadataV2, GroupMetadataV2};
272    use zarrs_metadata::v3::GroupMetadataV3;
273    #[cfg(feature = "async")]
274    use zarrs_storage::AsyncReadableWritableListableStorageTraits;
275    use zarrs_storage::store::MemoryStore;
276    use zarrs_storage::{StoreKey, WritableStorageTraits};
277
278    const EXPECTED_TREE: &str = "/\n  array [10, 10] float32\n  group\n    array [10, 10] float32\n    subgroup\n      mysubarray [10, 10] float32\n";
279
280    fn helper_create_dataset(store: &Arc<MemoryStore>) -> Group<MemoryStore> {
281        let group_builder = GroupBuilder::default();
282
283        let root = group_builder
284            .build(store.clone(), NodePath::root().as_str())
285            .unwrap();
286        let group = group_builder.build(store.clone(), "/group").unwrap();
287        let array_builder = ArrayBuilder::new(
288            vec![10, 10],
289            vec![5, 5],
290            crate::array::data_type::float32(),
291            0.0f32,
292        );
293
294        let array = array_builder.build(store.clone(), "/array").unwrap();
295        let group_array = array_builder.build(store.clone(), "/group/array").unwrap();
296        let subgroup = group_builder
297            .build(store.clone(), "/group/subgroup")
298            .unwrap();
299        let subgroup_array = array_builder
300            .build(store.clone(), "/group/subgroup/mysubarray")
301            .unwrap();
302
303        root.store_metadata().unwrap();
304        array.store_metadata().unwrap();
305        group.store_metadata().unwrap();
306        group_array.store_metadata().unwrap();
307        subgroup.store_metadata().unwrap();
308        subgroup_array.store_metadata().unwrap();
309
310        root
311    }
312
313    #[test]
314    fn hierarchy_try_from_array() {
315        let store = Arc::new(MemoryStore::new());
316        let array_builder =
317            ArrayBuilder::new(vec![1], vec![1], crate::array::data_type::float32(), 0.0f32);
318
319        let array = array_builder
320            .build(store, "/store/of/data.zarr/path/to/an/array")
321            .expect("Faulty test array");
322
323        let hierarchy = Hierarchy::try_from(&array).unwrap();
324        assert_eq!(hierarchy.0.len(), 1);
325        let hierarchy = Hierarchy::try_from(array).unwrap();
326        assert_eq!(hierarchy.0.len(), 1);
327    }
328
329    #[cfg(feature = "async")]
330    #[test]
331    fn hierarchy_try_from_async_array() {
332        let store = std::sync::Arc::new(zarrs_object_store::AsyncObjectStore::new(
333            object_store::memory::InMemory::new(),
334        ));
335        let array_builder =
336            ArrayBuilder::new(vec![1], vec![1], crate::array::data_type::float32(), 0.0f32);
337
338        let array = array_builder
339            .build(store, "/store/of/data.zarr/path/to/an/array")
340            .expect("Faulty test array");
341
342        let hierarchy = Hierarchy::try_from(&array).unwrap();
343        assert_eq!(hierarchy.0.len(), 1);
344        let hierarchy = Hierarchy::try_from(array).unwrap();
345        assert_eq!(hierarchy.0.len(), 1);
346    }
347
348    #[test]
349    fn hierarchy_try_from_group() {
350        let store = Arc::new(MemoryStore::new());
351        let group = helper_create_dataset(&store);
352        let hierarchy = Hierarchy::try_from_group(&group).unwrap();
353
354        assert_eq!(hierarchy.0.len(), 6);
355    }
356
357    #[cfg(feature = "async")]
358    #[tokio::test]
359    async fn hierarchy_async_try_from_group() {
360        let store = std::sync::Arc::new(zarrs_object_store::AsyncObjectStore::new(
361            object_store::memory::InMemory::new(),
362        ));
363        let group_path = "/group";
364        let group_builder = GroupBuilder::new();
365        let group = group_builder.build(store.clone(), group_path).unwrap();
366
367        let subgroup = group_builder
368            .build(store.clone(), "/group/subgroup")
369            .unwrap();
370
371        let array = ArrayBuilder::new(
372            vec![10, 10],
373            vec![5, 5],
374            crate::array::data_type::float32(),
375            0.0f32,
376        )
377        .build(store.clone(), "/group/subgroup/array")
378        .unwrap();
379
380        group.async_store_metadata().await.unwrap();
381        subgroup.async_store_metadata().await.unwrap();
382        array.async_store_metadata().await.unwrap();
383
384        let hierarchy = Hierarchy::try_from_async_group(&group).await;
385        assert!(hierarchy.is_ok());
386        let hierarchy = hierarchy.unwrap();
387        assert!(
388            "/\n  group\n    subgroup\n      array [10, 10] float32\n" == hierarchy.to_string()
389        );
390    }
391
392    #[cfg(feature = "async")]
393    #[tokio::test]
394    async fn hierarchy_async_try_from_invalid_async_group() {
395        let store = std::sync::Arc::new(zarrs_object_store::AsyncObjectStore::new(
396            object_store::memory::InMemory::new(),
397        ));
398
399        let root = Group::new_with_metadata(
400            store.clone(),
401            "/",
402            GroupMetadata::V3(GroupMetadataV3::default()),
403        )
404        .unwrap();
405
406        assert!(Hierarchy::try_from_async_group(&root).await.is_ok());
407
408        use zarrs_storage::AsyncWritableStorageTraits;
409        // Inject fauly subgroup
410        store
411            .set(
412                &StoreKey::new("subgroup/zarr.json").unwrap(),
413                vec![0].into(),
414            )
415            .await
416            .unwrap();
417
418        assert!(Hierarchy::try_from_async_group(&root).await.is_err());
419    }
420
421    #[test]
422    fn hierarchy_try_from_invalid_group() {
423        let store = Arc::new(MemoryStore::new());
424
425        let root = Group::new_with_metadata(
426            store.clone(),
427            "/",
428            GroupMetadata::V3(GroupMetadataV3::default()),
429        )
430        .unwrap();
431
432        assert!(Hierarchy::try_from_group(&root).is_ok());
433
434        // Inject fauly subgroup
435        store
436            .set(
437                &StoreKey::new("subgroup/zarr.json").unwrap(),
438                vec![0].into(),
439            )
440            .unwrap();
441
442        assert!(Hierarchy::try_from_group(&root).is_err());
443    }
444
445    #[test]
446    fn hierarchy_tree_of() {
447        let store = Arc::new(MemoryStore::new());
448
449        let group = helper_create_dataset(&store);
450
451        let hierarchy = Hierarchy::try_from_group(&group).unwrap();
452
453        assert_eq!(
454            "/group/subgroup\n  mysubarray [10, 10] float32\n",
455            hierarchy.tree_of(&NodePath::try_from("/group/subgroup").unwrap())
456        );
457    }
458
459    #[test]
460    fn hierarchy_tree() {
461        let store = Arc::new(MemoryStore::new());
462
463        let group = helper_create_dataset(&store);
464
465        let hierarchy = Hierarchy::try_from_group(&group).unwrap();
466
467        assert_eq!(
468            "/\n  array [10, 10] float32\n  group\n    array [10, 10] float32\n    subgroup\n      mysubarray [10, 10] float32\n",
469            hierarchy.tree()
470        );
471
472        let store = Arc::new(MemoryStore::new());
473        let groupv2 = Group::new_with_metadata(
474            store.clone(),
475            "/groupv2",
476            GroupMetadata::V2(GroupMetadataV2::new()),
477        )
478        .expect("Unexpected issue when greating a Group for testing.");
479
480        let arrayv2 = Array::new_with_metadata(
481            store.clone(),
482            "/groupv2/arrayv2",
483            ArrayMetadata::V2(ArrayMetadataV2::new(
484                vec![1],
485                crate::array::ChunkShape::from(vec![NonZeroU64::new(1).unwrap()]),
486                zarrs_metadata::v2::DataTypeMetadataV2::Simple("<f8".into()),
487                zarrs_metadata::FillValueMetadata::from(f64::NAN),
488                None,
489                None,
490            )),
491        )
492        .expect("Unexpected issue when creating a v2 Array for testing.");
493
494        let _ = groupv2.store_metadata();
495        let _ = arrayv2.store_metadata();
496
497        let h = Hierarchy::try_from_group(&groupv2);
498        assert!(h.is_ok());
499        assert!("/\n  groupv2\n    arrayv2 [1] Simple(\"<f8\")\n" == h.unwrap().tree());
500    }
501
502    #[test]
503    fn hierarchy_tree_empty() {
504        let hierarchy = Hierarchy::new();
505        let tree_str = hierarchy.tree();
506        assert_eq!(tree_str, "/\n");
507    }
508
509    #[test]
510    fn hierarchy_open() {
511        let store: std::sync::Arc<MemoryStore> = std::sync::Arc::new(MemoryStore::new());
512
513        let _group = helper_create_dataset(&store);
514
515        // Open a group node
516        let h = Hierarchy::open(&store, "/").unwrap();
517        assert_eq!(EXPECTED_TREE, h.tree());
518
519        // Open an array node
520        let h = Hierarchy::open(&store, "/array").unwrap();
521        assert_eq!("/\n  array [10, 10] float32\n", h.tree());
522    }
523
524    #[cfg(feature = "async")]
525    async fn async_helper_create_dataset<
526        AStore: ?Sized + AsyncReadableWritableListableStorageTraits + 'static,
527    >(
528        store: &Arc<AStore>,
529    ) -> Group<AStore> {
530        let group_builder = GroupBuilder::default();
531
532        let root = group_builder
533            .build(store.clone(), NodePath::root().as_str())
534            .unwrap();
535        let group = group_builder.build(store.clone(), "/group").unwrap();
536        let array_builder = ArrayBuilder::new(
537            vec![10, 10],
538            vec![5, 5],
539            crate::array::data_type::float32(),
540            0.0f32,
541        );
542
543        let array = array_builder.build(store.clone(), "/array").unwrap();
544        let group_array = array_builder.build(store.clone(), "/group/array").unwrap();
545        let subgroup = group_builder
546            .build(store.clone(), "/group/subgroup")
547            .unwrap();
548        let subgroup_array = array_builder
549            .build(store.clone(), "/group/subgroup/mysubarray")
550            .unwrap();
551
552        root.async_store_metadata().await.unwrap();
553        array.async_store_metadata().await.unwrap();
554        group.async_store_metadata().await.unwrap();
555        group_array.async_store_metadata().await.unwrap();
556        subgroup.async_store_metadata().await.unwrap();
557        subgroup_array.async_store_metadata().await.unwrap();
558
559        root
560    }
561
562    #[cfg(feature = "async")]
563    #[tokio::test]
564    async fn hierarchy_async_open() {
565        use zarrs_storage::AsyncReadableWritableListableStorage;
566
567        let store: AsyncReadableWritableListableStorage = Arc::new(
568            zarrs_object_store::AsyncObjectStore::new(object_store::memory::InMemory::new()),
569        );
570
571        let _group = async_helper_create_dataset(&store).await;
572
573        // Open a Group node
574        let h = Hierarchy::async_open(&store, "/").await;
575        assert!(h.is_ok());
576        assert_eq!(EXPECTED_TREE, h.unwrap().tree());
577
578        // Open an Array node
579        let h = Hierarchy::async_open(&store, "/array").await;
580        assert!(h.is_ok());
581        assert_eq!("/\n  array [10, 10] float32\n", h.unwrap().tree());
582    }
583}
zarrs/hierarchy.rs

zarrs/
hierarchy.rs
