Skip to content

Commit

Permalink
Properly closes zarr groups in zarr store (#8425)
Browse files Browse the repository at this point in the history
Co-authored-by: Maximilian Roos <[email protected]>
Co-authored-by: Anderson Banihirwe <[email protected]>
  • Loading branch information
3 people authored Dec 1, 2023
1 parent b703102 commit b313ffc
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 1 deletion.
3 changes: 3 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -174,6 +174,9 @@ Bug fixes
- Fix a bug where :py:meth:`DataArray.to_dataset` silently drops a variable
if a coordinate with the same name already exists (:pull:`8433`, :issue:`7823`).
By `András Gunyhó <https://github.com/mgunyho>`_.
- Fix for :py:meth:`DataArray.to_zarr` & :py:meth:`Dataset.to_zarr` to close
the created zarr store when passing a path with `.zip` extension (:pull:`8425`).
By `Carl Andersson <https://github.com/CarlAndersson>_`.

Documentation
~~~~~~~~~~~~~
Expand Down
8 changes: 7 additions & 1 deletion xarray/backends/zarr.py
Original file line number Diff line number Diff line change
Expand Up @@ -381,6 +381,7 @@ class ZarrStore(AbstractWritableDataStore):
"_write_region",
"_safe_chunks",
"_write_empty",
"_close_store_on_close",
)

@classmethod
Expand Down Expand Up @@ -464,6 +465,7 @@ def open_group(
zarr_group = zarr.open_consolidated(store, **open_kwargs)
else:
zarr_group = zarr.open_group(store, **open_kwargs)
close_store_on_close = zarr_group.store is not store
return cls(
zarr_group,
mode,
Expand All @@ -472,6 +474,7 @@ def open_group(
write_region,
safe_chunks,
write_empty,
close_store_on_close,
)

def __init__(
Expand All @@ -483,6 +486,7 @@ def __init__(
write_region=None,
safe_chunks=True,
write_empty: bool | None = None,
close_store_on_close: bool = False,
):
self.zarr_group = zarr_group
self._read_only = self.zarr_group.read_only
Expand All @@ -494,6 +498,7 @@ def __init__(
self._write_region = write_region
self._safe_chunks = safe_chunks
self._write_empty = write_empty
self._close_store_on_close = close_store_on_close

@property
def ds(self):
Expand Down Expand Up @@ -762,7 +767,8 @@ def set_variables(self, variables, check_encoding_set, writer, unlimited_dims=No
writer.add(v.data, zarr_array, region)

def close(self):
pass
if self._close_store_on_close:
self.zarr_group.store.close()


def open_zarr(
Expand Down
10 changes: 10 additions & 0 deletions xarray/tests/test_backends.py
Original file line number Diff line number Diff line change
Expand Up @@ -5251,6 +5251,16 @@ def test_pickle_open_mfdataset_dataset():
assert_identical(ds, pickle.loads(pickle.dumps(ds)))


@requires_zarr
def test_zarr_closing_internal_zip_store():
store_name = "tmp.zarr.zip"
original_da = DataArray(np.arange(12).reshape((3, 4)))
original_da.to_zarr(store_name, mode="w")

with open_dataarray(store_name, engine="zarr") as loaded_da:
assert_identical(original_da, loaded_da)


@requires_zarr
class TestZarrRegionAuto:
def test_zarr_region_auto_all(self, tmp_path):
Expand Down

0 comments on commit b313ffc

Please sign in to comment.