Skip to content

Commit

Permalink
update import path to match Zarr v2.12 and v2.13 experimental API
Browse files Browse the repository at this point in the history
remove path='xarray' default for zarr v3

path=None should work as of Zarr v2.13
  • Loading branch information
grlee77 committed Sep 21, 2022
1 parent d590b76 commit dbdf63f
Show file tree
Hide file tree
Showing 3 changed files with 72 additions and 31 deletions.
14 changes: 5 additions & 9 deletions xarray/backends/zarr.py
Original file line number Diff line number Diff line change
Expand Up @@ -371,11 +371,7 @@ def open_group(

if zarr_version is None:
# default to 2 if store doesn't specify it's version (e.g. a path)
zarr_version = getattr(store, '_store_version', 2)

if zarr_version > 2 and group is None:
# v3 stores require a group name: use 'xarray' as a default one.
group = 'xarray'
zarr_version = getattr(store, "_store_version", 2)

open_kwargs = dict(
mode=mode,
Expand Down Expand Up @@ -463,7 +459,7 @@ def open_store_variable(self, name, zarr_array):
# TODO: how to properly handle 'filters' for v3 stores
# currently these use a hack to store 'filters' within attributes
# need to drop this here for V3 store tests to succeed
attributes.pop('filters', None)
attributes.pop("filters", None)

encoding = {
"chunks": zarr_array.chunks,
Expand Down Expand Up @@ -592,11 +588,11 @@ def store(
self.set_variables(
variables_encoded, check_encoding_set, writer, unlimited_dims=unlimited_dims
)
zarr_version = getattr(self.zarr_group.store, '_store_version', 3)
zarr_version = getattr(self.zarr_group.store, "_store_version", 3)
consolidate_kwargs = {}
if zarr_version > 2:
# zarr v3 spec requires providing a path
consolidate_kwargs['path'] = self.zarr_group.path
# If the group has a path, it needs to also be passed to consolidate_metadata
consolidate_kwargs["path"] = self.zarr_group.path
if self._consolidate_on_close:
zarr.consolidate_metadata(self.zarr_group.store, **consolidate_kwargs)

Expand Down
1 change: 1 addition & 0 deletions xarray/core/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -1954,6 +1954,7 @@ def to_zarr(
region: Mapping[str, slice] | None = None,
safe_chunks: bool = True,
storage_options: dict[str, str] | None = None,
zarr_version : int | None = None,
) -> ZarrStore | Delayed:
"""Write dataset contents to a zarr group.
Expand Down
88 changes: 66 additions & 22 deletions xarray/tests/test_backends.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,13 +104,17 @@
have_zarr_kvstore = False
try:
from zarr.storage import KVStore

have_zarr_kvstore = True
except ImportError:
KVStore = None

have_zarr_v3 = False
try:
from zarr.storage_v3 import DirectoryStoreV3, KVStoreV3
# as of Zarr v2.13 these imports require environment variable
# ZARR_V3_EXPERIMENTAL_API=1
from zarr import DirectoryStoreV3, KVStoreV3

have_zarr_v3 = True
except ImportError:
KVStoreV3 = None
Expand Down Expand Up @@ -1705,14 +1709,18 @@ def create_zarr_target(self):
@contextlib.contextmanager
def create_store(self):
with self.create_zarr_target() as store_target:
yield backends.ZarrStore.open_group(store_target, mode="w", **self.version_kwargs)
yield backends.ZarrStore.open_group(
store_target, mode="w", **self.version_kwargs
)

def save(self, dataset, store_target, **kwargs):
return dataset.to_zarr(store=store_target, **kwargs, **self.version_kwargs)

@contextlib.contextmanager
def open(self, store_target, **kwargs):
with xr.open_dataset(store_target, engine="zarr", **kwargs, **self.version_kwargs) as ds:
with xr.open_dataset(
store_target, engine="zarr", **kwargs, **self.version_kwargs
) as ds:
yield ds

@contextlib.contextmanager
Expand Down Expand Up @@ -2040,9 +2048,13 @@ def test_write_persistence_modes(self, group):
ds, ds_to_append, _ = create_append_test_data()
with self.create_zarr_target() as store_target:
ds.to_zarr(store_target, mode="w", group=group, **self.version_kwargs)
ds_to_append.to_zarr(store_target, append_dim="time", group=group, **self.version_kwargs)
ds_to_append.to_zarr(
store_target, append_dim="time", group=group, **self.version_kwargs
)
original = xr.concat([ds, ds_to_append], dim="time")
actual = xr.open_dataset(store_target, group=group, engine="zarr", **self.version_kwargs)
actual = xr.open_dataset(
store_target, group=group, engine="zarr", **self.version_kwargs
)
assert_identical(original, actual)

def test_compressor_encoding(self):
Expand Down Expand Up @@ -2104,13 +2116,19 @@ def test_append_with_invalid_dim_raises(self):
with pytest.raises(
ValueError, match="does not match any existing dataset dimensions"
):
ds_to_append.to_zarr(store_target, append_dim="notvalid", **self.version_kwargs)
ds_to_append.to_zarr(
store_target, append_dim="notvalid", **self.version_kwargs
)

def test_append_with_no_dims_raises(self):
with self.create_zarr_target() as store_target:
Dataset({"foo": ("x", [1])}).to_zarr(store_target, mode="w", **self.version_kwargs)
Dataset({"foo": ("x", [1])}).to_zarr(
store_target, mode="w", **self.version_kwargs
)
with pytest.raises(ValueError, match="different dimension names"):
Dataset({"foo": ("y", [2])}).to_zarr(store_target, mode="a", **self.version_kwargs)
Dataset({"foo": ("y", [2])}).to_zarr(
store_target, mode="a", **self.version_kwargs
)

def test_append_with_append_dim_not_set_raises(self):
ds, ds_to_append, _ = create_append_test_data()
Expand All @@ -2124,7 +2142,9 @@ def test_append_with_mode_not_a_raises(self):
with self.create_zarr_target() as store_target:
ds.to_zarr(store_target, mode="w", **self.version_kwargs)
with pytest.raises(ValueError, match="cannot set append_dim unless"):
ds_to_append.to_zarr(store_target, mode="w", append_dim="time", **self.version_kwargs)
ds_to_append.to_zarr(
store_target, mode="w", append_dim="time", **self.version_kwargs
)

def test_append_with_existing_encoding_raises(self):
ds, ds_to_append, _ = create_append_test_data()
Expand Down Expand Up @@ -2161,11 +2181,15 @@ def test_check_encoding_is_consistent_after_append(self):
encoding = {"da": {"compressor": compressor}}
ds.to_zarr(store_target, mode="w", encoding=encoding, **self.version_kwargs)
ds_to_append.to_zarr(store_target, append_dim="time", **self.version_kwargs)
actual_ds = xr.open_dataset(store_target, engine="zarr", **self.version_kwargs)
actual_ds = xr.open_dataset(
store_target, engine="zarr", **self.version_kwargs
)
actual_encoding = actual_ds["da"].encoding["compressor"]
assert actual_encoding.get_config() == compressor.get_config()
assert_identical(
xr.open_dataset(store_target, engine="zarr", **self.version_kwargs).compute(),
xr.open_dataset(
store_target, engine="zarr", **self.version_kwargs
).compute(),
xr.concat([ds, ds_to_append], dim="time"),
)

Expand All @@ -2182,7 +2206,8 @@ def test_append_with_new_variable(self):
combined = xr.concat([ds, ds_to_append], dim="time")
combined["new_var"] = ds_with_new_var["new_var"]
assert_identical(
combined, xr.open_dataset(store_target, engine="zarr", **self.version_kwargs)
combined,
xr.open_dataset(store_target, engine="zarr", **self.version_kwargs),
)

@requires_dask
Expand Down Expand Up @@ -2289,9 +2314,14 @@ def test_write_region(self, consolidated, compute, use_dask):
for i in range(0, 10, 2):
region = {"x": slice(i, i + 2)}
nonzeros.isel(region).to_zarr(
store, region=region, consolidated=consolidated, **self.version_kwargs,
store,
region=region,
consolidated=consolidated,
**self.version_kwargs,
)
with xr.open_zarr(store, consolidated=consolidated, **self.version_kwargs) as actual:
with xr.open_zarr(
store, consolidated=consolidated, **self.version_kwargs
) as actual:
assert_identical(actual, nonzeros)

@pytest.mark.parametrize("mode", [None, "r+", "a"])
Expand All @@ -2301,7 +2331,9 @@ def test_write_region_mode(self, mode):
with self.create_zarr_target() as store:
zeros.to_zarr(store, **self.version_kwargs)
for region in [{"x": slice(5)}, {"x": slice(5, 10)}]:
nonzeros.isel(region).to_zarr(store, region=region, mode=mode, **self.version_kwargs)
nonzeros.isel(region).to_zarr(
store, region=region, mode=mode, **self.version_kwargs
)
with xr.open_zarr(store, **self.version_kwargs) as actual:
assert_identical(actual, nonzeros)

Expand Down Expand Up @@ -2343,7 +2375,9 @@ def test_write_preexisting_override_metadata(self):
with self.create_zarr_target() as store:
original.to_zarr(store, compute=False, **self.version_kwargs)
# with region, the default mode becomes r+
both_modified.to_zarr(store, region={"x": slice(None)}, **self.version_kwargs)
both_modified.to_zarr(
store, region={"x": slice(None)}, **self.version_kwargs
)
with self.open(store) as actual:
assert_identical(actual, only_new_data)

Expand Down Expand Up @@ -2371,7 +2405,9 @@ def setup_and_verify_store(expected=data):
"cannot set region unless mode='a', mode='r+' or mode=None"
),
):
data.to_zarr(store, region={"x": slice(None)}, mode="w", **self.version_kwargs)
data.to_zarr(
store, region={"x": slice(None)}, mode="w", **self.version_kwargs
)

with setup_and_verify_store() as store:
with pytest.raises(TypeError, match=r"must be a dict"):
Expand All @@ -2383,7 +2419,9 @@ def setup_and_verify_store(expected=data):

with setup_and_verify_store() as store:
with pytest.raises(ValueError, match=r"step on all slices"):
data2.to_zarr(store, region={"x": slice(None, None, 2)}, **self.version_kwargs)
data2.to_zarr(
store, region={"x": slice(None, None, 2)}, **self.version_kwargs
)

with setup_and_verify_store() as store:
with pytest.raises(
Expand All @@ -2397,13 +2435,20 @@ def setup_and_verify_store(expected=data):
ValueError,
match=r"all variables in the dataset to write must have at least one dimension in common",
):
data2.assign(v=2).to_zarr(store, region={"x": slice(2)}, **self.version_kwargs)
data2.assign(v=2).to_zarr(
store, region={"x": slice(2)}, **self.version_kwargs
)

with setup_and_verify_store() as store:
with pytest.raises(
ValueError, match=r"cannot list the same dimension in both"
):
data.to_zarr(store, region={"x": slice(None)}, append_dim="x", **self.version_kwargs)
data.to_zarr(
store,
region={"x": slice(None)},
append_dim="x",
**self.version_kwargs,
)

with setup_and_verify_store() as store:
with pytest.raises(
Expand Down Expand Up @@ -2517,7 +2562,6 @@ def create_store(self):
yield group



class ZarrBaseV3(ZarrBase):
def test_roundtrip_coordinates_with_space(self):
original = Dataset(coords={"x": 0, "y z": 1})
Expand Down Expand Up @@ -2548,7 +2592,7 @@ def create_zarr_target(self):
class TestZarrDirectoryStoreV3FromPath(TestZarrDirectoryStoreV3):
# Must specify zarr_version=3 to get a v3 store because create_zarr_target
# is a string path.
version_kwargs = {'zarr_version': 3}
version_kwargs = {"zarr_version": 3}

@contextlib.contextmanager
def create_zarr_target(self):
Expand Down

0 comments on commit dbdf63f

Please sign in to comment.