Merge pull request #18 from LDeakin/array_sharded_ext_traits

Add array extension traits to simplify working with sharded arrays

CHANGELOG.md

@@ -7,6 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+ - Added the `array_sharded_ext::{ArrayShardedExt,ArrayShardedReadableExt}` extension traits for `Array` to simplify working with sharded arrays
+   - Abstracts the chunk grid to an "inner chunk grid" to simplify inner chunk retrieval.
+   - Shard indexes are cached in a `ArrayShardedReadableExtCache`
+   - Retrieval and chunk grid methods have fallbacks for unsharded arrays. For example, an inner chunk in an unsharded array is just a chunk
+   - Sync API only, `AsyncArrayShardedReadableExt` and `AsyncArrayShardedReadableExtCache` are planned for a future release
+ - Added `ChunkGridTraits::chunks_subset()` with default implementation
+
 ### Changed
  - Allow float fill values to be created from int fill value metadata
  - Make `chunk_grid::{regular,rectangular}` public

src/array.rs

@@ -25,6 +25,11 @@ mod fill_value_metadata;
 mod nan_representations;
 mod unsafe_cell_slice;
 
+#[cfg(feature = "sharding")]
+mod array_sharded_ext;
+#[cfg(feature = "sharding")]
+mod array_sync_sharded_readable_ext;
+
 use std::sync::Arc;
 
 pub use self::{
@@ -49,6 +54,12 @@ pub use self::{
     unsafe_cell_slice::UnsafeCellSlice,
 };
 
+#[cfg(feature = "sharding")]
+pub use array_sharded_ext::ArrayShardedExt;
+#[cfg(feature = "sharding")]
+pub use array_sync_sharded_readable_ext::{ArrayShardedReadableExt, ArrayShardedReadableExtCache};
+// TODO: Add AsyncArrayShardedReadableExt and AsyncArrayShardedReadableExtCache
+
 use serde::Serialize;
 use thiserror::Error;
 

src/array/array_sharded_ext.rs

@@ -0,0 +1,53 @@
+use super::{codec::ShardingCodecConfiguration, Array, ChunkGrid, ChunkShape};
+
+/// An [`Array`] extension trait to simplify working with arrays using the `sharding_indexed` codec.
+pub trait ArrayShardedExt {
+    /// Returns true if the array to bytes codec of the array is `sharding_indexed`.
+    fn is_sharded(&self) -> bool;
+
+    /// Return the inner chunk shape.
+    ///
+    /// Returns [`None`] for an unsharded array.
+    fn inner_chunk_shape(&self) -> Option<ChunkShape>;
+
+    /// Retrieve the inner chunk grid.
+    ///
+    /// Returns the normal chunk grid for an unsharded array.
+    fn inner_chunk_grid(&self) -> ChunkGrid;
+}
+
+impl<TStorage: ?Sized> ArrayShardedExt for Array<TStorage> {
+    fn is_sharded(&self) -> bool {
+        self.codecs
+            .array_to_bytes_codec()
+            .create_metadata()
+            .expect("the array to bytes codec should have metadata")
+            .name() // TODO: Add codec::identifier()?
+            == super::codec::array_to_bytes::sharding::IDENTIFIER
+    }
+
+    fn inner_chunk_shape(&self) -> Option<ChunkShape> {
+        let codec_metadata = self
+            .codecs
+            .array_to_bytes_codec()
+            .create_metadata()
+            .expect("the array to bytes codec should have metadata");
+        if let Ok(ShardingCodecConfiguration::V1(sharding_configuration)) =
+            codec_metadata.to_configuration()
+        {
+            Some(sharding_configuration.chunk_shape)
+        } else {
+            None
+        }
+    }
+
+    fn inner_chunk_grid(&self) -> ChunkGrid {
+        if let Some(inner_chunk_shape) = self.inner_chunk_shape() {
+            ChunkGrid::new(crate::array::chunk_grid::RegularChunkGrid::new(
+                inner_chunk_shape,
+            ))
+        } else {
+            self.chunk_grid().clone()
+        }
+    }
+}