|
7 | 7 | import warnings
|
8 | 8 |
|
9 | 9 | import cupy as cp
|
10 |
| -import numpy as np |
11 |
| -from xarray import Variable |
12 |
| -from xarray.backends import zarr as zarr_backend |
13 | 10 | from xarray.backends.common import _normalize_path # TODO: can this be public
|
14 | 11 | 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 |
17 | 13 | from xarray.core.utils import close_on_error # TODO: can this be public.
|
18 | 14 |
|
19 | 15 | try:
|
|
29 | 25 | # TODO: minimum xarray version for ZarrArrayWrapper._array 2023.10.0?
|
30 | 26 |
|
31 | 27 |
|
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 |
| - |
67 | 28 | class GDSZarrStore(ZarrStore):
|
68 | 29 | @classmethod
|
69 | 30 | def open_group(
|
@@ -136,35 +97,6 @@ def open_group(
|
136 | 97 | safe_chunks,
|
137 | 98 | )
|
138 | 99 |
|
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 |
| - |
168 | 100 |
|
169 | 101 | class KvikioBackendEntrypoint(ZarrBackendEntrypoint):
|
170 | 102 | """
|
|
0 commit comments