Fast path generalized to slices aligned to chunks

Author: FrancescAlted

bench/ndarray/aligned_chunks.py

@@ -0,0 +1,109 @@
+#######################################################################
+# Copyright (c) 2019-present, Blosc Development Team <blosc@blosc.org>
+# All rights reserved.
+#
+# This source code is licensed under a BSD-style license (found in the
+# LICENSE file in the root directory of this source tree)
+#######################################################################
+
+# Benchmark for comparing speeds of NDArray.slice() when using
+# different slices containing consecutive and non-consecutive chunks,
+# as well as aligned and unaligned.
+
+import math
+from time import time
+import numpy as np
+import blosc2
+
+# Dimensions and type properties for the arrays
+shape = (50, 100, 300)
+chunks = (5, 25, 50)
+blocks = (1, 5, 10)
+dtype = np.dtype(np.int32)
+
+# Non-consecutive slices
+nc_slices = [
+    (slice(0, 50), slice(0, 100), slice(0, 300-1)),
+    (slice(0, 10), slice(0, 100-1), slice(0, 300)),
+    (slice(0, 5-1), slice(0, 25), slice(0, 300)),
+    (slice(0, 5), slice(0, 25), slice(0, 50-1)),
+    ]
+# Consecutive slices
+c_slices = [
+    (slice(0, 50), slice(0, 100), slice(0, 300)),
+    (slice(0, 10), slice(0, 100), slice(0, 300)),
+    (slice(0, 5), slice(0, 25), slice(0, 300)),
+    (slice(0, 5), slice(0, 25), slice(0, 50)),
+    ]
+# Non-aligned slices
+na_slices = [
+    (slice(10, 50-1), slice(25, 100), slice(50, 300)),
+    (slice(10, 40), slice(25, 75-1), slice(100, 200)),
+    (slice(20, 35), slice(50, 75), slice(100, 300-1)),
+    (slice(20+1, 25), slice(25, 50), slice(50, 100)),
+    ]
+# Aligned slices
+a_slices = [
+    (slice(10, 50), slice(25, 100), slice(50, 300)),
+    (slice(10, 40), slice(25, 75), slice(100, 200)),
+    (slice(20, 35), slice(50, 75), slice(100, 300)),
+    (slice(20, 25), slice(25, 50), slice(50, 100)),
+    ]
+
+print("Creating array with shape:", shape)
+t0 = time()
+arr = blosc2.arange(math.prod(shape), dtype=dtype, shape=shape, chunks=chunks, blocks=blocks)
+print(f"Time to create array: {time() - t0 : .5f}")
+
+print("Timing non-consecutive slices...")
+nc_times = []
+t0 = time()
+for s in nc_slices:
+    t1 = time()
+    arr2 = arr.slice(s)
+    nc_times.append(time() - t1)
+    # print(arr2.schunk.nbytes, arr[s].nbytes)
+    # np.testing.assert_array_equal(arr2[:], arr[s])
+print(f"Time to get non-consecutive slices: {time() - t0 : .5f}")
+
+print("Timing consecutive slices...")
+c_times = []
+c_speedup = []
+t0 = time()
+for i, s in enumerate(c_slices):
+    t1 = time()
+    arr2 = arr.slice(s)
+    c_times.append(time() - t1)
+    c_speedup.append(nc_times[i] / c_times[i])
+    # print(arr2.shape, arr[s].shape)
+    # print(arr2.schunk.nbytes, arr[s].nbytes)
+    # np.testing.assert_array_equal(arr2[:], arr[s])
+print(f"Time to get consecutive slices: {time() - t0 : .5f}")
+print(f"Speedups for consecutive slices: ", [f"{s:.2f}x" for s in c_speedup])
+
+print("Timing non-aligned slices...")
+na_times = []
+t0 = time()
+for i, s in enumerate(na_slices):
+    t1 = time()
+    arr2 = arr.slice(s)
+    na_times.append(time() - t1)
+    # print(arr2.shape, arr[s].shape)
+    # print(arr2.schunk.nbytes, arr[s].nbytes)
+    # np.testing.assert_array_equal(arr2[:], arr[s])
+print(f"Time to get non-aligned slices: {time() - t0 : .5f}")
+
+print("Timing aligned slices...")
+a_times = []
+a_speedup = []
+t0 = time()
+for i, s in enumerate(a_slices):
+    t1 = time()
+    arr2 = arr.slice(s)
+    a_times.append(time() - t1)
+    a_speedup.append(na_times[i] / a_times[i])
+    # print(arr2.shape, arr[s].shape)
+    # print(arr2.schunk.nbytes, arr[s].nbytes)
+    # np.testing.assert_array_equal(arr2[:], arr[s])
+print(f"Time to get aligned slices: {time() - t0 : .5f}")
+print(f"Speedups for aligned slices: ", [f"{s:.2f}x" for s in a_speedup])

bench/ndarray/consecutive_chunks.py

@@ -16,6 +16,7 @@
 from itertools import product
 from typing import TYPE_CHECKING, Any, NamedTuple
 
+from dask.array.routines import aligned_coarsen_chunks
 from numpy.exceptions import ComplexWarning
 
 if TYPE_CHECKING:
@@ -1030,11 +1031,11 @@ def extract_values(arr, indices: np.ndarray[np.int_], max_cache_size: int = 10)
     return extracted_values
 
 
-def detect_consecutive_chunks(  # noqa: C901
-    key: Sequence[slice], shape: Sequence[int], chunks: Sequence[int]
+def detect_aligned_chunks(
+    key: Sequence[slice], shape: Sequence[int], chunks: Sequence[int], consecutive: bool = False
 ) -> list[int]:
     """
-    Detect whether a multidimensional slice matches a sequence of consecutive chunk boundaries.
+    Detect whether a multidimensional slice is aligned with chunk boundaries.
 
     Parameters
     ----------
@@ -1044,27 +1045,28 @@ def detect_consecutive_chunks(  # noqa: C901
         Shape of the NDArray.
     chunks : Sequence of int
         Chunk shape of the NDArray.
+    consecutive : bool, default=False
+        If True, check if the chunks are consecutive in storage order.
+        If False, only check for chunk boundary alignment.
 
     Returns
     -------
     list[int]
-        Index of the chunk (in C-order) if the slice matches exactly with a single chunk,
-        a list of chunk indices if the slice matches a consecutive sequence of chunks.
-        If it doesn't match any chunk(s) properly, return an empty list.
+        List of chunk indices (in C-order) that the slice overlaps with.
+        If the slice isn't aligned with chunk boundaries, returns an empty list.
+        If consecutive=True and chunks aren't consecutive, returns an empty list.
     """
     if len(key) != len(shape):
         return []
 
-    # Check that slice boundaries are exact multiple of chunk boundaries.
-    # We want to do that so we don't copy data, and hence, waste space,
-    # unnecessarily into destination.
+    # Check that slice boundaries are exact multiple of chunk boundaries
     for i, s in enumerate(key):
         if s.start is not None and s.start % chunks[i] != 0:
             return []
         if s.stop is not None and s.stop % chunks[i] != 0:
             return []
 
-    # Parse the slice boundaries and check for alignment
+    # Parse the slice boundaries
     start_indices = []
     end_indices = []
     n_chunks = []
@@ -1088,6 +1090,7 @@ def detect_consecutive_chunks(  # noqa: C901
         end_indices.append(end_idx)
         n_chunks.append(math.ceil(shape[i] / chunk_size))
 
+    # Get all chunk combinations in the slice
     indices = [range(start, end) for start, end in zip(start_indices, end_indices, strict=False)]
     result = []
 
@@ -1100,17 +1103,18 @@ def detect_consecutive_chunks(  # noqa: C901
 
         result.append(flat_index)
 
-    if not result:
-        return []
-
-    # The product() of ranges might not naturally produce indices in ascending order
-    result.sort()
-    is_consecutive = builtins.all(result[i] == result[i - 1] + 1 for i in range(1, len(result)))
+    # Check if chunks are consecutive if requested
+    if consecutive and result:
+        sorted_result = sorted(result)
+        if sorted_result[-1] - sorted_result[0] + 1 != len(sorted_result):
+            return []
 
-    if not is_consecutive:
-        return []
+        # The array of indices must be consecutive
+        for i in range(len(sorted_result) - 1):
+            if sorted_result[i + 1] - sorted_result[i] != 1:
+                return []
 
-    return result
+    return sorted(result)
 
 
 class NDOuterIterator:
@@ -1933,8 +1937,9 @@ def slice(self, key: int | slice | Sequence[slice], **kwargs: Any) -> NDArray:
 
         # Fast path for slices made with consecutive chunks
         if step == (1,) * self.ndim:
-            consecutive_chunks = detect_consecutive_chunks(key, self.shape, self.chunks)
-            if consecutive_chunks:
+            aligned_chunks = detect_aligned_chunks(key, self.shape, self.chunks, consecutive=False)
+            if aligned_chunks:
+                # print("Aligned chunks detected", aligned_chunks)
                 # Create a new ndarray for the key slice
                 new_shape = [
                     sp - st for sp, st in zip([k.stop for k in key], [k.start for k in key], strict=False)
@@ -1948,7 +1953,7 @@ def slice(self, key: int | slice | Sequence[slice], **kwargs: Any) -> NDArray:
                 )
                 # Get the chunks from the original array and update the new array
                 # No need for chunks to decompress and compress again
-                for order, nchunk in enumerate(consecutive_chunks):
+                for order, nchunk in enumerate(aligned_chunks):
                     chunk = self.schunk.get_chunk(nchunk)
                     newarr.schunk.update_chunk(order, chunk)
                 return newarr

src/blosc2/ndarray.py

@@ -21,11 +21,21 @@
     ((10, 100, 300), (5, 25, 50), (1, 5, 10), (slice(0, 5), slice(0, 100), slice(0, 300)), np.int32),
     ((10, 100, 300), (5, 25, 50), (1, 5, 10), (slice(0, 5), slice(0, 25), slice(0, 200)), np.int32),
     ((10, 100, 300), (5, 25, 50), (1, 5, 10), (slice(0, 5), slice(0, 25), slice(0, 50)), np.int32),
+    # Aligned slices
+    ((10, 100, 300), (5, 25, 50), (1, 5, 10), (slice(10, 50), slice(25, 100), slice(50, 300)), np.int32),
+    ((10, 100, 300), (5, 25, 50), (1, 5, 10), (slice(10, 40), slice(25, 75), slice(100, 200)), np.int32),
+    ((10, 100, 300), (5, 25, 50), (1, 5, 10), (slice(20, 35), slice(50, 75), slice(100, 300)), np.int32),
+    ((10, 100, 300), (5, 25, 50), (1, 5, 10), (slice(20, 25), slice(25, 50), slice(50, 100)), np.int32),
     # Non-consecutive slices
     ((10, 100, 300), (5, 25, 50), (1, 5, 10), (slice(0, 10), slice(0, 100), slice(0, 300 - 1)), np.int32),
     ((10, 100, 300), (5, 25, 50), (1, 5, 10), (slice(0, 5), slice(0, 100 - 1), slice(0, 300)), np.int32),
     ((10, 100, 300), (5, 25, 50), (1, 5, 10), (slice(0, 5 - 1), slice(0, 25), slice(0, 200)), np.int32),
     ((10, 100, 300), (5, 25, 50), (1, 5, 10), (slice(0, 5), slice(0, 25), slice(0, 50 - 1)), np.int32),
+    # Non-aligned slices
+    ((10, 100, 300), (5, 25, 50), (1, 5, 10), (slice(10, 50 - 1), slice(25, 100), slice(50, 300)), np.int32),
+    ((10, 100, 300), (5, 25, 50), (1, 5, 10), (slice(10, 40), slice(25, 75 - 1), slice(100, 200)), np.int32),
+    ((10, 100, 300), (5, 25, 50), (1, 5, 10), (slice(20, 35), slice(50, 75), slice(100, 300 - 1)), np.int32),
+    ((10, 100, 300), (5, 25, 50), (1, 5, 10), (slice(20 + 1, 25), slice(25, 50), slice(50, 100)), np.int32),
 ]