Merge pull request #698 from Zentrik/fix-asarray-overhanging-chunks

Fix asarray when chunk dims are greater than shape dims

Author: FrancescAlted

src/blosc2/ndarray.py

@@ -315,28 +315,12 @@ def are_partitions_behaved(shape, chunks, blocks):
     bool
         True if the partitions are well-behaved, False otherwise.
     """
+    def check_contiguity(container, part):
+        if container and container[-1] != part[-1]:
+            return False
+        return builtins.all(size % unit == 0 for size, unit in zip(container[:-1], part[:-1], strict=True))
 
-    # Check C-contiguity among partitions
-    def check_contiguity(shape, part):
-        ndims = len(shape)
-        inner_dim = ndims - 1
-        for i, size, unit in zip(reversed(range(ndims)), reversed(shape), reversed(part), strict=True):
-            if size > unit:
-                if i < inner_dim:
-                    if size % unit != 0:
-                        return False
-                else:
-                    if size != unit:
-                        return False
-                inner_dim = i
-        return True
-
-    # Check C-contiguity for blocks inside chunks
-    if not check_contiguity(chunks, blocks):
-        return False
-
-    # Check C-contiguity for chunks inside shape
-    return check_contiguity(shape, chunks)
+    return check_contiguity(chunks, blocks) and check_contiguity(shape, chunks)
 
 
 def get_flat_slices_orig(shape: tuple[int], s: tuple[slice, ...]) -> list[slice]:

tests/ndarray/test_ndarray.py

@@ -103,6 +103,25 @@ def test_asarray(a):
         np.testing.assert_allclose(a, b[:])
 
 
+@pytest.mark.parametrize(
+    ("shape", "chunks", "blocks"),
+    [
+        ((146, 23802), (147, 23802), (1, 23802)),
+        ((2200, 1000), (1100, 1024), (10, 1024)),
+        ((20, 300, 500), (21, 300, 500), (1, 300, 500)),
+    ],
+)
+def test_asarray_chunks_larger_than_shape(shape, chunks, blocks):
+    # Above 16 MB, asarray fills chunk by chunk; padded chunks must skip update_data
+    a = np.arange(math.prod(shape), dtype=np.float64).reshape(shape)
+    assert a.nbytes > 2**24
+
+    b = blosc2.asarray(a, chunks=chunks, blocks=blocks)
+
+    np.testing.assert_array_equal(b[:], a)
+    assert not blosc2.are_partitions_behaved(shape, chunks, blocks)
+
+
 def test_asarray_persists_copy_with_urlpath(tmp_path):
     array = blosc2.asarray(np.arange(10, dtype=np.int64), chunks=(5,), blocks=(2,))
     path = tmp_path / "persisted_copy.b2nd"