check non-CUDA compatible compressor

rapidsai · Oct 5, 2023 · 9700c57 · 9700c57
1 parent cbc2cca
commit 9700c57
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/python/kvikio/zarr.py b/python/kvikio/zarr.py
@@ -381,6 +381,12 @@ def open_cupy_array(
                 meta_array=meta_array,
                 **kwargs,
             )
+        elif not isinstance(ret.compressor, CudaCodec):
+            raise ValueError(
+                "The Zarr file was written using a non-CUDA compatible "
+                f"compressor, {ret.compressor}, please use something "
+                "like kvikio.zarr.CompatCompressor"
+            )
         return ret
 
     if isinstance(compressor, CompatCompressor):

diff --git a/python/tests/test_zarr.py b/python/tests/test_zarr.py
@@ -245,3 +245,11 @@ def test_open_cupy_array(tmp_path):
     assert isinstance(z[:], numpy.ndarray)
     assert z.compressor == kvikio_zarr.CompatCompressor.lz4().cpu
     numpy.testing.assert_array_equal(a.get(), z[:])
+
+
+@pytest.mark.parametrize("mode", ["r", "r+"])
+def test_open_cupy_array_incompatible_compressor(tmp_path, mode):
+    zarr.create((10,), store=tmp_path, compressor=numcodecs.Blosc())
+
+    with pytest.raises(ValueError, match="non-CUDA compatible compressor"):
+        kvikio_zarr.open_cupy_array(tmp_path, mode=mode)