Skip to content

Commit cb77678

Browse files
committed
Remove custom open_store_variable method from GDSZarrStore class
Can directly rely on upstream xarray's ZarrStore.open_store_variable method since Zarr v3 compatibility was added in pydata/xarray#9552.
1 parent 7dd78e9 commit cb77678

File tree

1 file changed

+1
-69
lines changed

1 file changed

+1
-69
lines changed

cupy_xarray/kvikio.py

Lines changed: 1 addition & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,9 @@
77
import warnings
88

99
import cupy as cp
10-
import numpy as np
11-
from xarray import Variable
12-
from xarray.backends import zarr as zarr_backend
1310
from xarray.backends.common import _normalize_path # TODO: can this be public
1411
from xarray.backends.store import StoreBackendEntrypoint
15-
from xarray.backends.zarr import ZarrArrayWrapper, ZarrBackendEntrypoint, ZarrStore
16-
from xarray.core import indexing
12+
from xarray.backends.zarr import ZarrBackendEntrypoint, ZarrStore
1713
from xarray.core.utils import close_on_error # TODO: can this be public.
1814

1915
try:
@@ -29,41 +25,6 @@
2925
# TODO: minimum xarray version for ZarrArrayWrapper._array 2023.10.0?
3026

3127

32-
class DummyZarrArrayWrapper(ZarrArrayWrapper):
33-
def __init__(self, array: np.ndarray):
34-
assert isinstance(array, np.ndarray)
35-
self._array = array
36-
self.filters = None
37-
self.dtype = array.dtype
38-
self.shape = array.shape
39-
40-
def __array__(self):
41-
return self._array
42-
43-
def get_array(self):
44-
return self._array
45-
46-
def __getitem__(self, key):
47-
return self._array[key]
48-
49-
50-
class CupyZarrArrayWrapper(ZarrArrayWrapper):
51-
def __array__(self):
52-
return self.get_array()
53-
54-
55-
class EagerCupyZarrArrayWrapper(ZarrArrayWrapper):
56-
"""Used to wrap dimension coordinates."""
57-
58-
def __array__(self):
59-
return self._array[:].get()
60-
61-
def get_duck_array(self):
62-
# total hack: make a numpy array look like a Zarr array
63-
# this gets us through Xarray's backend layers
64-
return DummyZarrArrayWrapper(self._array[:].get())
65-
66-
6728
class GDSZarrStore(ZarrStore):
6829
@classmethod
6930
def open_group(
@@ -136,35 +97,6 @@ def open_group(
13697
safe_chunks,
13798
)
13899

139-
def open_store_variable(self, name, zarr_array):
140-
try_nczarr = self._mode == "r"
141-
dimensions, attributes = zarr_backend._get_zarr_dims_and_attrs(
142-
zarr_array, zarr_backend.DIMENSION_KEY, try_nczarr
143-
)
144-
145-
#### Changed from zarr array wrapper
146-
# we want indexed dimensions to be loaded eagerly
147-
# Right now we load in to device and then transfer to host
148-
# But these should be small-ish arrays
149-
# TODO: can we tell GDSStore to load as numpy array directly
150-
# not cupy array?
151-
array_wrapper = EagerCupyZarrArrayWrapper if name in dimensions else CupyZarrArrayWrapper
152-
data = indexing.LazilyIndexedArray(array_wrapper(zarr_array))
153-
154-
attributes = dict(attributes)
155-
encoding = {
156-
"chunks": zarr_array.chunks,
157-
"preferred_chunks": dict(zip(dimensions, zarr_array.chunks)),
158-
"compressor": zarr_array.compressor,
159-
"filters": zarr_array.filters,
160-
}
161-
# _FillValue needs to be in attributes, not encoding, so it will get
162-
# picked up by decode_cf
163-
if zarr_array.fill_value is not None:
164-
attributes["_FillValue"] = zarr_array.fill_value
165-
166-
return Variable(dimensions, data, attributes, encoding)
167-
168100

169101
class KvikioBackendEntrypoint(ZarrBackendEntrypoint):
170102
"""

0 commit comments

Comments
 (0)