From b313ffc2bd6ffc1c221f8950fff3cceb24cb775f Mon Sep 17 00:00:00 2001 From: Carl Andersson Date: Fri, 1 Dec 2023 03:20:14 +0100 Subject: [PATCH] Properly closes zarr groups in zarr store (#8425) Co-authored-by: Maximilian Roos <5635139+max-sixty@users.noreply.github.com> Co-authored-by: Anderson Banihirwe <13301940+andersy005@users.noreply.github.com> --- doc/whats-new.rst | 3 +++ xarray/backends/zarr.py | 8 +++++++- xarray/tests/test_backends.py | 10 ++++++++++ 3 files changed, 20 insertions(+), 1 deletion(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 9fc1b0ba80a..abc3760df6e 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -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ó `_. +- 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 _`. Documentation ~~~~~~~~~~~~~ diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 30b306ceb34..7ff59e0e7bf 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -381,6 +381,7 @@ class ZarrStore(AbstractWritableDataStore): "_write_region", "_safe_chunks", "_write_empty", + "_close_store_on_close", ) @classmethod @@ -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, @@ -472,6 +474,7 @@ def open_group( write_region, safe_chunks, write_empty, + close_store_on_close, ) def __init__( @@ -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 @@ -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): @@ -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( diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 8c5f2f8b98a..d6f76df067c 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -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):