Kvikio backend entrypoint with Zarr v3

ci/doc.yml

@@ -3,8 +3,9 @@ channels:
   - conda-forge
 dependencies:
   - cupy-core
+  - rapidsai-nightly::kvikio>=25.04.00a
   - pip
-  - python=3.10
+  - python=3.11
   - sphinx
   - sphinx-design
   - sphinx-copybutton
@@ -14,7 +15,8 @@ dependencies:
   - ipywidgets
   - furo>=2024.8.6
   - myst-nb
-  - xarray
+  - zarr>=3.0.3
   - pip:
+      - xarray @ git+https://github.com/pydata/xarray.git@0184702f16c3f744fc9096c7dac690626dcc6922 #https://github.com/pydata/xarray/pull/10078
       # relative to this file. Needs to be editable to be accepted.
       - --editable ..

cupy_xarray/__init__.py

@@ -1,4 +1,5 @@
 from . import _version
-from .accessors import CupyDataArrayAccessor, CupyDatasetAccessor  # noqa
+from .accessors import CupyDataArrayAccessor, CupyDatasetAccessor  # noqa: F401
+from .kvikio import KvikioBackendEntrypoint  # noqa: F401
 
 __version__ = _version.get_versions()["version"]

cupy_xarray/kvikio.py

@@ -0,0 +1,90 @@
+"""
+:doc:`kvikIO <kvikio:index>` backend for xarray to read Zarr stores directly into CuPy
+arrays in GPU memory.
+"""
+
+from xarray.backends.common import _normalize_path  # TODO: can this be public
+from xarray.backends.store import StoreBackendEntrypoint
+from xarray.backends.zarr import ZarrBackendEntrypoint, ZarrStore
+from xarray.core.dataset import Dataset
+from xarray.core.utils import close_on_error  # TODO: can this be public.
+
+try:
+    import kvikio.zarr
+
+    has_kvikio = True
+except ImportError:
+    has_kvikio = False
+
+
+class KvikioBackendEntrypoint(ZarrBackendEntrypoint):
+    """
+    Xarray backend to read Zarr stores using 'kvikio' engine.
+
+    For more information about the underlying library, visit
+    :doc:`kvikIO's Zarr page<kvikio:zarr>`.
+    """
+
+    available = has_kvikio
+    description = "Open zarr files (.zarr) using Kvikio"
+    url = "https://docs.rapids.ai/api/kvikio/stable/api/#zarr"
+
+    # disabled by default
+    # We need to provide this because of the subclassing from
+    # ZarrBackendEntrypoint
+    def guess_can_open(self, filename_or_obj):
+        return False
+
+    def open_dataset(
+        self,
+        filename_or_obj,
+        mask_and_scale=True,
+        decode_times=True,
+        concat_characters=True,
+        decode_coords=True,
+        drop_variables=None,
+        use_cftime=None,
+        decode_timedelta=None,
+        group=None,
+        mode="r",
+        synchronizer=None,
+        consolidated=None,
+        chunk_store=None,
+        storage_options=None,
+        zarr_version=None,
+        zarr_format=None,
+        store=None,
+        engine=None,
+        use_zarr_fill_value_as_mask=None,
+        cache_members: bool = True,
+    ) -> Dataset:
+        filename_or_obj = _normalize_path(filename_or_obj)
+        if not store:
+            store = ZarrStore.open_group(
+                store=kvikio.zarr.GDSStore(root=filename_or_obj),
+                group=group,
+                mode=mode,
+                synchronizer=synchronizer,
+                consolidated=consolidated,
+                consolidate_on_close=False,
+                chunk_store=chunk_store,
+                storage_options=storage_options,
+                zarr_version=zarr_version,
+                use_zarr_fill_value_as_mask=None,
+                zarr_format=zarr_format,
+                cache_members=cache_members,
+            )
+
+        store_entrypoint = StoreBackendEntrypoint()
+        with close_on_error(store):
+            ds = store_entrypoint.open_dataset(
+                store,
+                mask_and_scale=mask_and_scale,
+                decode_times=decode_times,
+                concat_characters=concat_characters,
+                decode_coords=decode_coords,
+                drop_variables=drop_variables,
+                use_cftime=use_cftime,
+                decode_timedelta=decode_timedelta,
+            )
+        return ds

cupy_xarray/tests/test_kvikio.py

@@ -0,0 +1,54 @@
+import cupy as cp
+import numpy as np
+import pytest
+import xarray as xr
+from xarray.core.indexing import ExplicitlyIndexedNDArrayMixin
+
+kvikio = pytest.importorskip("kvikio")
+zarr = pytest.importorskip("zarr")
+
+import kvikio.zarr  # noqa
+import xarray.core.indexing  # noqa
+
+
+@pytest.fixture
+def store(tmp_path):
+    ds = xr.Dataset(
+        {
+            "a": ("x", np.arange(10), {"foo": "bar"}),
+            "scalar": np.array(1),
+        },
+        coords={"x": ("x", np.arange(-5, 5))},
+    )
+
+    for var in ds.variables:
+        ds[var].encoding["compressors"] = None
+
+    store_path = tmp_path / "kvikio.zarr"
+    ds.to_zarr(store_path, consolidated=True)
+    return store_path
+
+
+def test_entrypoint():
+    assert "kvikio" in xr.backends.list_engines()
+
+
+@pytest.mark.parametrize("consolidated", [True, False])
+def test_lazy_load(consolidated, store):
+    with xr.open_dataset(store, engine="kvikio", consolidated=consolidated) as ds:
+        for _, da in ds.data_vars.items():
+            assert isinstance(da.variable._data, ExplicitlyIndexedNDArrayMixin)
+
+
+@pytest.mark.parametrize("indexer", [slice(None), slice(2, 4), 2, [2, 3, 5]])
+def test_lazy_indexing(indexer, store):
+    with zarr.config.enable_gpu(), xr.open_dataset(store, engine="kvikio") as ds:
+        ds = ds.isel(x=indexer)
+        for _, da in ds.data_vars.items():
+            assert isinstance(da.variable._data, ExplicitlyIndexedNDArrayMixin)
+
+        loaded = ds.compute()
+        for _, da in loaded.data_vars.items():
+            if da.ndim == 0:
+                continue
+            assert isinstance(da.data, cp.ndarray)

docs/api.rst

@@ -51,3 +51,16 @@ Methods
 
     Dataset.cupy.as_cupy
     Dataset.cupy.as_numpy
+
+
+KvikIO engine
+-------------
+
+.. currentmodule:: cupy_xarray
+
+.. automodule:: cupy_xarray.kvikio
+
+.. autosummary::
+   :toctree: generated/
+
+   KvikioBackendEntrypoint

docs/conf.py

@@ -57,6 +57,7 @@
     "python": ("https://docs.python.org/3/", None),
     "dask": ("https://docs.dask.org/en/latest", None),
     "cupy": ("https://docs.cupy.dev/en/latest", None),
+    "kvikio": ("https://docs.rapids.ai/api/kvikio/stable", None),
     "xarray": ("http://docs.xarray.dev/en/latest/", None),
 }
 

docs/index.md

@@ -59,6 +59,7 @@ Large parts of this documentations comes from [SciPy 2023 Xarray on GPUs tutoria
    source/high-level-api
    source/apply-ufunc
    source/real-example-1
+   source/kvikio
 
 
 **Tutorials & Presentations**: