Add Kvikio backend entrypoint

ci/doc.yml

@@ -15,6 +15,7 @@ dependencies:
   - furo>=2024.8.6
   - myst-nb
   - xarray
+  - zarr
   - pip:
       # 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,230 @@
+"""
+:doc:`kvikIO <kvikio:index>` backend for xarray to read Zarr stores directly into CuPy
+arrays in GPU memory.
+"""
+
+import os
+import warnings
+
+import cupy as cp
+import numpy as np
+from xarray import Variable
+from xarray.backends import zarr as zarr_backend
+from xarray.backends.common import _normalize_path  # TODO: can this be public
+from xarray.backends.store import StoreBackendEntrypoint
+from xarray.backends.zarr import ZarrArrayWrapper, ZarrBackendEntrypoint, ZarrStore
+from xarray.core import indexing
+from xarray.core.utils import close_on_error  # TODO: can this be public.
+
+try:
+    import kvikio.zarr
+    import zarr
+
+    has_kvikio = True
+except ImportError:
+    has_kvikio = False
+
+
+#  TODO: minimum kvikio version for supporting consolidated
+#  TODO: minimum xarray version for ZarrArrayWrapper._array 2023.10.0?
+
+
+class DummyZarrArrayWrapper(ZarrArrayWrapper):
+    def __init__(self, array: np.ndarray):
+        assert isinstance(array, np.ndarray)
+        self._array = array
+        self.filters = None
+        self.dtype = array.dtype
+        self.shape = array.shape
+
+    def __array__(self):
+        return self._array
+
+    def get_array(self):
+        return self._array
+
+    def __getitem__(self, key):
+        return self._array[key]
+
+
+class CupyZarrArrayWrapper(ZarrArrayWrapper):
+    def __array__(self):
+        return self.get_array()
+
+
+class EagerCupyZarrArrayWrapper(ZarrArrayWrapper):
+    """Used to wrap dimension coordinates."""
+
+    def __array__(self):
+        return self._array[:].get()
+
+    def get_duck_array(self):
+        # total hack: make a numpy array look like a Zarr array
+        # this gets us through Xarray's backend layers
+        return DummyZarrArrayWrapper(self._array[:].get())
+
+
+class GDSZarrStore(ZarrStore):
+    @classmethod
+    def open_group(
+        cls,
+        store,
+        mode="r",
+        synchronizer=None,
+        group=None,
+        consolidated=False,
+        consolidate_on_close=False,
+        chunk_store=None,
+        storage_options=None,
+        append_dim=None,
+        write_region=None,
+        safe_chunks=True,
+        stacklevel=2,
+    ):
+        # zarr doesn't support pathlib.Path objects yet. zarr-python#601
+        if isinstance(store, os.PathLike):
+            store = os.fspath(store)
+
+        open_kwargs = {
+            "mode": mode,
+            "synchronizer": synchronizer,
+            "path": group,
+            ########## NEW STUFF
+            "meta_array": cp.empty(()),
+        }
+        open_kwargs["storage_options"] = storage_options
+
+        if chunk_store:
+            open_kwargs["chunk_store"] = chunk_store
+            if consolidated is None:
+                consolidated = False
+
+        store = kvikio.zarr.GDSStore(store)
+
+        if consolidated is None:
+            try:
+                zarr_group = zarr.open_consolidated(store, **open_kwargs)
+            except KeyError:
+                warnings.warn(
+                    "Failed to open Zarr store with consolidated metadata, "
+                    "falling back to try reading non-consolidated metadata. "
+                    "This is typically much slower for opening a dataset. "
+                    "To silence this warning, consider:\n"
+                    "1. Consolidating metadata in this existing store with "
+                    "zarr.consolidate_metadata().\n"
+                    "2. Explicitly setting consolidated=False, to avoid trying "
+                    "to read consolidate metadata, or\n"
+                    "3. Explicitly setting consolidated=True, to raise an "
+                    "error in this case instead of falling back to try "
+                    "reading non-consolidated metadata.",
+                    RuntimeWarning,
+                    stacklevel=stacklevel,
+                )
+                zarr_group = zarr.open_group(store, **open_kwargs)
+        elif consolidated:
+            # TODO: an option to pass the metadata_key keyword
+            zarr_group = zarr.open_consolidated(store, **open_kwargs)
+        else:
+            zarr_group = zarr.open_group(store, **open_kwargs)
+
+        return cls(
+            zarr_group,
+            mode,
+            consolidate_on_close,
+            append_dim,
+            write_region,
+            safe_chunks,
+        )
+
+    def open_store_variable(self, name, zarr_array):
+        try_nczarr = self._mode == "r"
+        dimensions, attributes = zarr_backend._get_zarr_dims_and_attrs(
+            zarr_array, zarr_backend.DIMENSION_KEY, try_nczarr
+        )
+
+        #### Changed from zarr array wrapper
+        # we want indexed dimensions to be loaded eagerly
+        # Right now we load in to device and then transfer to host
+        # But these should be small-ish arrays
+        # TODO: can we tell GDSStore to load as numpy array directly
+        # not cupy array?
+        array_wrapper = EagerCupyZarrArrayWrapper if name in dimensions else CupyZarrArrayWrapper
+        data = indexing.LazilyIndexedArray(array_wrapper(zarr_array))
+
+        attributes = dict(attributes)
+        encoding = {
+            "chunks": zarr_array.chunks,
+            "preferred_chunks": dict(zip(dimensions, zarr_array.chunks)),
+            "compressor": zarr_array.compressor,
+            "filters": zarr_array.filters,
+        }
+        # _FillValue needs to be in attributes, not encoding, so it will get
+        # picked up by decode_cf
+        if zarr_array.fill_value is not None:
+            attributes["_FillValue"] = zarr_array.fill_value
+
+        return Variable(dimensions, data, attributes, encoding)
+
+
+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,
+        stacklevel=3,
+    ):
+        filename_or_obj = _normalize_path(filename_or_obj)
+        store = GDSZarrStore.open_group(
+            filename_or_obj,
+            group=group,
+            mode=mode,
+            synchronizer=synchronizer,
+            consolidated=consolidated,
+            consolidate_on_close=False,
+            chunk_store=chunk_store,
+            storage_options=storage_options,
+            stacklevel=stacklevel + 1,
+        )
+
+        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["compressor"] = 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 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**: