ChunkIterator: account for smaller last slice (TGSAI#702)

Author: dmitriyrepin

src/mdio/core/indexing.py

@@ -74,11 +74,13 @@ def __next__(self) -> dict[str, slice]:
             # We build slices here. It is dimension agnostic
             current_start = next(self._ranges)
 
-            # TODO (Dmitriy Repin): Enhance ChunkIterator to make the last slice, if needed, smaller
-            # https://github.com/TGSAI/mdio-python/issues/586
             start_indices = tuple(dim * chunk for dim, chunk in zip(current_start, self.len_chunks, strict=True))
 
-            stop_indices = tuple((dim + 1) * chunk for dim, chunk in zip(current_start, self.len_chunks, strict=True))
+            # Calculate stop indices, making the last slice fit the data exactly
+            stop_indices = tuple(
+                min((dim + 1) * chunk, self.arr_shape[i])
+                for i, (dim, chunk) in enumerate(zip(current_start, self.len_chunks, strict=True))
+            )
 
             slices = tuple(slice(start, stop) for start, stop in zip(start_indices, stop_indices, strict=True))
 

tests/unit/test_indexing.py

@@ -25,8 +25,8 @@ def test_chunk_iterator_returning_dict() -> None:
     assert iter2.dim_chunks == (2, 3, 4)
     assert iter2.num_chunks == 24
 
-    # Its purpose is to confirm that all slices are created of the same size,
-    # even if the last slice should have been smaller.
+    # Its purpose is to confirm that the last slice is adjusted to fit the data exactly
+    # when the array size doesn't align perfectly with chunk boundaries.
     for _ in range(13):  # element index 12
         region = iter1.__next__()
     assert region == {
@@ -38,7 +38,7 @@ def test_chunk_iterator_returning_dict() -> None:
     for _ in range(13):  # element index 12
         region = iter2.__next__()
     assert region == {
-        "inline": slice(3, 6, None),
+        "inline": slice(3, 5, None),
         "crossline": slice(0, 4, None),
         "depth": slice(0, 5, None),
     }
@@ -61,15 +61,15 @@ def test_chunk_iterator_returning_tuple() -> None:
     assert iter2.dim_chunks == (2, 3, 4)
     assert iter2.num_chunks == 24
 
-    # Its purpose is to confirm that all slices are created of the same size,
-    # even if the last slice should have been smaller.
+    # Its purpose is to confirm that the last slice is adjusted to fit the data exactly
+    # when the array size doesn't align perfectly with chunk boundaries.
     for _ in range(13):  # element index 12
         region = iter1.__next__()
     assert region == (slice(3, 6, None), slice(0, 4, None), slice(0, 5, None))
 
     for _ in range(13):  # element index 12
         region = iter2.__next__()
-    assert region == (slice(3, 6, None), slice(0, 4, None), slice(0, 5, None))
+    assert region == (slice(3, 5, None), slice(0, 4, None), slice(0, 5, None))
 
 
 def val(shape: tuple[int, int, int], i: int, j: int, k: int) -> int: