zarrs/array/
array_sharded_ext.rs

1use super::codec::ShardingCodecConfiguration;
2use super::{Array, ArrayShape, ChunkGrid, ChunkShape};
3use crate::array::codec::array_to_bytes::sharding::ShardingCodec;
4use zarrs_metadata::ConfigurationSerialize;
5
6/// An [`Array`] extension trait to simplify working with arrays using the `sharding_indexed` codec.
7pub trait ArrayShardedExt: private::Sealed {
8    /// Returns true if the array to bytes codec of the array is `sharding_indexed`.
9    fn is_sharded(&self) -> bool;
10
11    /// Returns true if the array-to-bytes codec of the array is `sharding_indexed` and the array has no array-to-array or bytes-to-bytes codecs.
12    fn is_exclusively_sharded(&self) -> bool;
13
14    /// Return the subchunk shape as defined in the `sharding_indexed` codec metadata.
15    ///
16    /// Returns [`None`] for an unsharded array.
17    fn subchunk_shape(&self) -> Option<ChunkShape>;
18
19    /// The effective subchunk shape.
20    ///
21    /// The effective subchunk shape is the "read granularity" of the sharded array that accounts for array-to-array codecs preceding the sharding codec.
22    /// For example, the transpose codec changes the shape of an array subset that corresponds to a single subchunk.
23    /// The effective subchunk shape is used when determining the subchunk grid of a sharded array.
24    ///
25    /// Returns [`None`] for an unsharded array of if the effective subchunk shape is indeterminate.
26    fn effective_subchunk_shape(&self) -> Option<ChunkShape>;
27
28    /// Retrieve the subchunk grid.
29    ///
30    /// This uses the effective subchunk shape so that reading a subchunk reads only one contiguous byte range.
31    ///
32    /// Returns the normal chunk grid for an unsharded array.
33    fn subchunk_grid(&self) -> ChunkGrid;
34
35    /// Return the shape of the subchunk grid (i.e., the number of subchunks).
36    ///
37    /// Returns the normal chunk grid shape for an unsharded array.
38    fn subchunk_grid_shape(&self) -> ArrayShape;
39}
40
41impl<TStorage: ?Sized> ArrayShardedExt for Array<TStorage> {
42    fn is_sharded(&self) -> bool {
43        self.codecs
44            .array_to_bytes_codec()
45            .as_any()
46            .is::<ShardingCodec>()
47    }
48
49    fn is_exclusively_sharded(&self) -> bool {
50        self.is_sharded()
51            && self.codecs.array_to_array_codecs().is_empty()
52            && self.codecs.bytes_to_bytes_codecs().is_empty()
53    }
54
55    fn subchunk_shape(&self) -> Option<ChunkShape> {
56        let configuration = self
57            .codecs
58            .array_to_bytes_codec()
59            .configuration_v3(self.metadata_options.codec_metadata_options())
60            .expect("the array to bytes codec should have metadata");
61        if let Ok(ShardingCodecConfiguration::V1(sharding_configuration)) =
62            ShardingCodecConfiguration::try_from_configuration(configuration)
63        {
64            Some(sharding_configuration.chunk_shape)
65        } else {
66            None
67        }
68    }
69
70    fn effective_subchunk_shape(&self) -> Option<ChunkShape> {
71        let mut subchunk_shape = self.subchunk_shape()?;
72        for codec in self.codecs().array_to_array_codecs().iter().rev() {
73            if let Ok(Some(subchunk_shape_)) = codec.decoded_shape(&subchunk_shape) {
74                subchunk_shape = subchunk_shape_;
75            } else {
76                return None;
77            }
78        }
79        Some(subchunk_shape)
80    }
81
82    fn subchunk_grid(&self) -> ChunkGrid {
83        // FIXME: Create the subchunk grid in `Array` and return a ref
84        if let Some(subchunk_shape) = self.effective_subchunk_shape() {
85            ChunkGrid::new(
86                crate::array::chunk_grid::RegularChunkGrid::new(
87                    self.shape().to_vec(),
88                    subchunk_shape,
89                ).expect("the subchunk grid dimensionality is already confirmed to match the array dimensionality"),
90            )
91        } else {
92            self.chunk_grid().clone()
93        }
94    }
95
96    fn subchunk_grid_shape(&self) -> ArrayShape {
97        self.subchunk_grid().grid_shape().to_vec()
98    }
99}
100
101mod private {
102    use super::Array;
103
104    pub trait Sealed {}
105
106    impl<TStorage: ?Sized> Sealed for Array<TStorage> {}
107}
zarrs/array/array_sharded_ext.rs

zarrs/array/
array_sharded_ext.rs
