Add fast path to retrieve_chunk_subset_opt if the entire chunk is r…

…equested
LDeakin · Feb 12, 2024 · 7005b2a · 7005b2a
1 parent 5631f30
commit 7005b2a
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 14 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -50,6 +50,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
      - `Indices`, `LinearisedIndices`, `ContiguousIndices`, `ContiguousLinearisedIndices`, `Chunks`
    - `Indices` and `Chunks` implement `IntoParallelIter`
    - **Breaking**: array subset iterators are moved into public `array_subset::iterators` and no longer in the `array_subset` namespace
+ - Add a fast path to `Array::retrieve_chunk_subset{_opt}` if the entire chunk is requested
 
 ### Removed
  - **Breaking**: remove `InvalidArraySubsetError` and `ArrayExtractElementsError`

diff --git a/src/array/array_sync_readable.rs b/src/array/array_sync_readable.rs
@@ -669,24 +669,31 @@ impl<TStorage: ?Sized + ReadableStorageTraits + 'static> Array<TStorage> {
             ));
         }
 
-        let storage_handle = Arc::new(StorageHandle::new(self.storage.clone()));
-        let storage_transformer = self
-            .storage_transformers()
-            .create_readable_transformer(storage_handle);
-        let input_handle = Box::new(StoragePartialDecoder::new(
-            storage_transformer,
-            data_key(self.path(), chunk_indices, self.chunk_key_encoding()),
-        ));
+        let decoded_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.retrieve_chunk(chunk_indices)?
+        } else {
+            let storage_handle = Arc::new(StorageHandle::new(self.storage.clone()));
+            let storage_transformer = self
+                .storage_transformers()
+                .create_readable_transformer(storage_handle);
+            let input_handle = Box::new(StoragePartialDecoder::new(
+                storage_transformer,
+                data_key(self.path(), chunk_indices, self.chunk_key_encoding()),
+            ));
 
-        let decoded_bytes = self
-            .codecs()
-            .partial_decoder_opt(input_handle, &chunk_representation, options)?
-            .partial_decode_opt(&[chunk_subset.clone()], options)?;
+            self.codecs()
+                .partial_decoder_opt(input_handle, &chunk_representation, options)?
+                .partial_decode_opt(&[chunk_subset.clone()], options)?
+                .concat()
+        };
 
-        let total_size = decoded_bytes.iter().map(Vec::len).sum::<usize>();
+        let total_size = decoded_bytes.len();
         let expected_size = chunk_subset.num_elements_usize() * self.data_type().size();
         if total_size == chunk_subset.num_elements_usize() * self.data_type().size() {
-            Ok(decoded_bytes.concat())
+            Ok(decoded_bytes)
         } else {
             Err(ArrayError::UnexpectedChunkDecodedSize(
                 total_size,