Fixed missing fast path in Array::[async_]retrieve_chunk_subset_opt

LDeakin · Sep 19, 2024 · d5fa6ab · d5fa6ab
1 parent d31e129
commit d5fa6ab
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 17 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 - Fixed an unnecessary copy in `Array::[async_]retrieve_chunk_if_exists_opt`
 - Fixed `CodecOptions` not being forwarded in `Array::retrieve_chunk_subset_opt` on the fast path
+- Fixed missing fast path in `Array::[async_]retrieve_chunk_subset_opt`
 
 ## [0.17.0-beta.1] - 2024-09-16
 

diff --git a/zarrs/src/array/array_async_readable.rs b/zarrs/src/array/array_async_readable.rs
@@ -683,23 +683,29 @@ impl<TStorage: ?Sized + AsyncReadableStorageTraits + 'static> Array<TStorage> {
             ));
         }
 
-        let storage_handle = Arc::new(StorageHandle::new(self.storage.clone()));
-        let storage_transformer = self
-            .storage_transformers()
-            .create_async_readable_transformer(storage_handle);
-        let input_handle = Arc::new(AsyncStoragePartialDecoder::new(
-            storage_transformer,
-            self.chunk_key(chunk_indices),
-        ));
-
-        let bytes = self
-            .codecs()
-            .async_partial_decoder(input_handle, &chunk_representation, options)
-            .await?
-            .partial_decode_opt(&[chunk_subset.clone()], options)
-            .await?
-            .remove(0)
-            .into_owned();
+        let bytes = if chunk_subset.start().iter().all(|&o| o == 0)
+            && chunk_subset.shape() == chunk_representation.shape_u64()
+        {
+            // Fast path if `chunk_subset` encompasses the whole chunk
+            self.async_retrieve_chunk_opt(chunk_indices, options)
+                .await?
+        } else {
+            let storage_handle = Arc::new(StorageHandle::new(self.storage.clone()));
+            let storage_transformer = self
+                .storage_transformers()
+                .create_async_readable_transformer(storage_handle);
+            let input_handle = Arc::new(AsyncStoragePartialDecoder::new(
+                storage_transformer,
+                self.chunk_key(chunk_indices),
+            ));
+            self.codecs()
+                .async_partial_decoder(input_handle, &chunk_representation, options)
+                .await?
+                .partial_decode_opt(&[chunk_subset.clone()], options)
+                .await?
+                .remove(0)
+                .into_owned()
+        };
         bytes.validate(chunk_subset.num_elements(), self.data_type().size())?;
         Ok(bytes)
     }