From a2d7a13a1c21601a434ceaa80618fed084e5123a Mon Sep 17 00:00:00 2001 From: Aaron Spring Date: Wed, 8 Jul 2026 11:34:11 +0200 Subject: [PATCH 1/2] Open existing groups when writing a DataTree with region/append_dim virtual_datatree_to_icechunk unconditionally created each node's group with Group.from_store, so forwarding region or append_dim (which require the group and its arrays to already exist) always failed with ContainsGroupError. Mirror the branch virtual_dataset_to_icechunk already had: open the existing group for region/append_dim, create it otherwise. Removes the now-passing xfail on test_write_datatree_region, whose stated reason (an xarray region limitation) was a misdiagnosis of this ContainsGroupError. Co-Authored-By: Claude Fable 5 --- docs/about/releases.md | 3 +++ virtualizarr/tests/test_writers/test_icechunk.py | 3 --- virtualizarr/writers/icechunk.py | 7 ++++++- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/docs/about/releases.md b/docs/about/releases.md index 31f671e7..5cf67b0a 100644 --- a/docs/about/releases.md +++ b/docs/about/releases.md @@ -9,6 +9,9 @@ ### Bug fixes +- Writing a virtual `DataTree` with `region` or `append_dim` (forwarded to each node) previously always failed with `ContainsGroupError`, because every group was unconditionally created; the existing groups are now opened instead. + By [Aaron Spring](https://github.com/aaronspring). + ### Documentation ### Internal changes diff --git a/virtualizarr/tests/test_writers/test_icechunk.py b/virtualizarr/tests/test_writers/test_icechunk.py index 01e7465a..fcbbb17f 100644 --- a/virtualizarr/tests/test_writers/test_icechunk.py +++ b/virtualizarr/tests/test_writers/test_icechunk.py @@ -996,9 +996,6 @@ def test_write_region_unaligned_chunks_raises( write_session.store, region={"y": "auto", "x": "auto"} ) - @pytest.mark.xfail( - reason="This doesn't work in xarray either, maybe not even intended?", - ) def test_write_datatree_region( self, icechunk_repo: "Repository", diff --git a/virtualizarr/writers/icechunk.py b/virtualizarr/writers/icechunk.py index 7c6741ff..522d30ad 100644 --- a/virtualizarr/writers/icechunk.py +++ b/virtualizarr/writers/icechunk.py @@ -238,7 +238,12 @@ def get_store_path(subtree, vdt) -> StorePath: # TODO this serial loop could be slow writing lots of groups to high-latency store, see https://github.com/pydata/xarray/issues/9455 for store_path, vds in paths_and_virtual_datasets: - group = Group.from_store(store=store_path, zarr_format=3) + if kwargs.get("append_dim") or kwargs.get("region"): + # both require the group (and arrays) to already exist + group = Group.open(store=store_path, zarr_format=3) + else: + # create the group if it doesn't already exist + group = Group.from_store(store=store_path, zarr_format=3) write_virtual_dataset_to_icechunk_group( vds=vds, From d787844dd7d082bd7a6e26af53808b43b18012fa Mon Sep 17 00:00:00 2001 From: Aaron Spring Date: Wed, 8 Jul 2026 11:35:48 +0200 Subject: [PATCH 2/2] Surface mode in to_icechunk, keeping current defaults Add a mode parameter ("w", "w-", "a") to vds.vz.to_icechunk and dt.vz.to_icechunk. The default (None) keeps the existing behaviour of raising ContainsGroupError when the target group already exists, while mode="a" opens the existing group so datasets can e.g. be split across commits by variable, and mode="w" overwrites existing contents. A _resolve_mode helper validates mode and resolves it to the effective zarr group-open mode ("r+" when append_dim/region require an existing group), so both writers open groups through a single zarr.open_group call. Closes #1001. Co-Authored-By: Claude Fable 5 --- docs/about/releases.md | 2 + docs/how_to/scaling.md | 2 +- virtualizarr/accessor.py | 21 ++++ .../tests/test_writers/test_icechunk.py | 115 +++++++++++++++++- virtualizarr/writers/icechunk.py | 72 ++++++++--- 5 files changed, 196 insertions(+), 16 deletions(-) diff --git a/docs/about/releases.md b/docs/about/releases.md index 5cf67b0a..3e8e8459 100644 --- a/docs/about/releases.md +++ b/docs/about/releases.md @@ -4,6 +4,8 @@ ### New Features - Added [VirtualiZarrDatasetAccessor.nrefs][virtualizarr.accessor.VirtualiZarrDatasetAccessor.nrefs] — a method that returns the total number of virtual chunk references in the dataset, ignoring non-virtual variables. Closes #573. +- `vds.vz.to_icechunk` and `vdt.vz.to_icechunk` now accept a `mode` parameter controlling how a pre-existing group at the target path is handled: `"w-"` (create, error if the group exists — the previous and still-default behaviour), `"w"` (overwrite existing contents), or `"a"` (open the existing group and add/update variables in it). `mode="a"` enables e.g. splitting a large virtual dataset across commits by variable. Closes #1001. + By [Aaron Spring](https://github.com/aaronspring). ### Breaking changes diff --git a/docs/how_to/scaling.md b/docs/how_to/scaling.md index 2153ea47..7d9cd2f8 100644 --- a/docs/how_to/scaling.md +++ b/docs/how_to/scaling.md @@ -334,7 +334,7 @@ for i, start in enumerate(range(0, vds.sizes["time"], step)): If the slice boundaries don't align with chunk edges along that axis, the indexing call raises `SubChunkIndexingError`. -(Remember you can also subset the Dataset to specific variables and commit those separately too if necessary.) +(Remember you can also subset the Dataset to specific variables and commit those separately too if necessary — pass `mode="a"` to `to_icechunk` from the second write onwards, so that writing into the already-existing group doesn't raise a `ContainsGroupError`.) ### Retries diff --git a/virtualizarr/accessor.py b/virtualizarr/accessor.py index 5f51833f..b4ff12e0 100644 --- a/virtualizarr/accessor.py +++ b/virtualizarr/accessor.py @@ -75,6 +75,7 @@ def to_icechunk( store: "IcechunkStore", *, group: str | None = None, + mode: Literal["w", "w-", "a"] | None = None, append_dim: str | None = None, region: Literal["auto"] | Mapping[str, Literal["auto"] | slice] | None = None, validate_containers: bool = True, @@ -104,6 +105,16 @@ def to_icechunk( Store to write dataset into. group Path of the group to write the dataset into (default: the root group). + mode + How to handle a pre-existing group at the target path: + + - ``"w-"``: create the group, raising a ``ContainsGroupError`` if it already exists. + - ``"w"``: create the group, overwriting any existing contents at that path. + - ``"a"``: open the group if it exists (keeping existing arrays), otherwise create it. + - ``None`` (default): equivalent to ``"w-"``, unless ``append_dim`` or ``region`` + is given, in which case the existing group is opened. + + ``mode="w"`` and ``mode="w-"`` are incompatible with ``append_dim`` and ``region``. append_dim Dimension along which to append the virtual dataset. region @@ -133,6 +144,7 @@ def to_icechunk( self.ds, store, group=group, + mode=mode, append_dim=append_dim, region=region, validate_containers=validate_containers, @@ -337,6 +349,7 @@ def to_icechunk( self, store: "IcechunkStore", *, + mode: Literal["w", "w-", "a"] | None = None, write_inherited_coords: bool = False, validate_containers: bool = True, last_updated_at: datetime | None = None, @@ -361,6 +374,13 @@ def to_icechunk( ---------- store Store to write dataset into. + mode + How to handle pre-existing groups at the target paths: + + - ``"w-"`` or ``None`` (default): create each group, raising a + ``ContainsGroupError`` if it already exists. + - ``"w"``: create each group, overwriting any existing contents at that path. + - ``"a"``: open each group if it exists (keeping existing arrays), otherwise create it. write_inherited_coords If ``True``, replicate inherited coordinates on all descendant nodes. Otherwise, only write coordinates at the level at which they are @@ -400,6 +420,7 @@ def to_icechunk( virtual_datatree_to_icechunk( self.dt, store, + mode=mode, write_inherited_coords=write_inherited_coords, validate_containers=validate_containers, last_updated_at=last_updated_at, diff --git a/virtualizarr/tests/test_writers/test_icechunk.py b/virtualizarr/tests/test_writers/test_icechunk.py index fcbbb17f..a7fcdf82 100644 --- a/virtualizarr/tests/test_writers/test_icechunk.py +++ b/virtualizarr/tests/test_writers/test_icechunk.py @@ -15,6 +15,7 @@ from zarr.codecs import BytesCodec from zarr.core.metadata import ArrayV3Metadata from zarr.dtype import parse_data_type +from zarr.errors import ContainsGroupError from virtualizarr import open_virtual_dataset from virtualizarr.manifests import ChunkManifest, ManifestArray @@ -62,7 +63,7 @@ def icechunk_filestore(icechunk_repo: "Repository") -> "IcechunkStore": return session.store -@pytest.mark.parametrize("kwarg", [("group", {}), ("append_dim", {})]) +@pytest.mark.parametrize("kwarg", [("group", {}), ("mode", {}), ("append_dim", {})]) def test_invalid_kwarg_type( icechunk_filestore: "IcechunkStore", vds_with_manifest_arrays: xr.Dataset, @@ -113,6 +114,118 @@ def test_write_new_virtual_variable( assert arr.metadata.dimension_names == ("x", "y") +@pytest.mark.parametrize("mode", [None, "w-"]) +def test_write_to_existing_group_fails_by_default( + icechunk_filestore: "IcechunkStore", + vds_with_manifest_arrays: xr.Dataset, + mode: Optional[str], +): + vds = vds_with_manifest_arrays + vds.vz.to_icechunk(icechunk_filestore, validate_containers=False) + + with pytest.raises(ContainsGroupError): + vds.vz.to_icechunk(icechunk_filestore, mode=mode, validate_containers=False) + + +def test_write_variables_across_commits_with_mode_a( + icechunk_repo: "Repository", + synthetic_vds_multiple_vars, +): + # regression test for https://github.com/zarr-developers/VirtualiZarr/issues/1001 + vds, arr = synthetic_vds_multiple_vars + + session1 = icechunk_repo.writable_session("main") + vds[["foo"]].vz.to_icechunk(session1.store) + session1.commit("wrote foo") + + session2 = icechunk_repo.writable_session("main") + vds[["bar"]].vz.to_icechunk(session2.store, mode="a") + session2.commit("wrote bar") + + with xr.open_zarr( + icechunk_repo.readonly_session("main").store, zarr_format=3, consolidated=False + ) as ds: + # both variables have encoding={"scale_factor": 2} + np.testing.assert_equal(ds["foo"].data, arr * 2) + np.testing.assert_equal(ds["bar"].data, arr * 2) + + +def test_write_parent_group_after_child_group_with_mode_a( + icechunk_filestore: "IcechunkStore", + vds_with_manifest_arrays: xr.Dataset, +): + # regression test for https://github.com/zarr-developers/VirtualiZarr/issues/1001 + vds = vds_with_manifest_arrays + vds.vz.to_icechunk( + icechunk_filestore, group="supgroup/subgroup", validate_containers=False + ) + vds.vz.to_icechunk( + icechunk_filestore, group="supgroup", mode="a", validate_containers=False + ) + + assert "a" in zarr.group(store=icechunk_filestore, path="supgroup") + assert "a" in zarr.group(store=icechunk_filestore, path="supgroup/subgroup") + + +def test_mode_w_overwrites_existing_group( + icechunk_filestore: "IcechunkStore", + synthetic_vds_multiple_vars, +): + vds, arr = synthetic_vds_multiple_vars + vds.vz.to_icechunk(icechunk_filestore) + vds[["foo"]].vz.to_icechunk(icechunk_filestore, mode="w") + + group = zarr.group(store=icechunk_filestore) + assert "foo" in group + assert "bar" not in group + + +def test_invalid_mode( + icechunk_filestore: "IcechunkStore", + vds_with_manifest_arrays: xr.Dataset, +): + with pytest.raises(ValueError, match="mode"): + vds_with_manifest_arrays.vz.to_icechunk(icechunk_filestore, mode="r+") + + +@pytest.mark.parametrize("mode", ["w", "w-"]) +def test_mode_incompatible_with_append_dim( + icechunk_filestore: "IcechunkStore", + vds_with_manifest_arrays: xr.Dataset, + mode: str, +): + with pytest.raises(ValueError, match="append_dim or region"): + vds_with_manifest_arrays.vz.to_icechunk( + icechunk_filestore, mode=mode, append_dim="x" + ) + + +def test_write_datatree_to_existing_groups_with_mode_a( + icechunk_repo: "Repository", + synthetic_vds_multiple_vars, +): + vds, arr = synthetic_vds_multiple_vars + + session1 = icechunk_repo.writable_session("main") + vdt1 = xr.DataTree.from_dict({"nested/group": vds[["foo"]]}) + vdt1.vz.to_icechunk(session1.store) + session1.commit("wrote foo") + + session2 = icechunk_repo.writable_session("main") + vdt2 = xr.DataTree.from_dict({"nested/group": vds[["bar"]]}) + vdt2.vz.to_icechunk(session2.store, mode="a") + session2.commit("wrote bar") + + with xr.open_zarr( + icechunk_repo.readonly_session("main").store, + zarr_format=3, + consolidated=False, + group="nested/group", + ) as ds: + np.testing.assert_equal(ds["foo"].data, arr * 2) + np.testing.assert_equal(ds["bar"].data, arr * 2) + + def test_set_single_virtual_ref_without_encoding( icechunk_filestore: "IcechunkStore", icechunk_repo: "Repository", diff --git a/virtualizarr/writers/icechunk.py b/virtualizarr/writers/icechunk.py index 522d30ad..2e3ee79a 100644 --- a/virtualizarr/writers/icechunk.py +++ b/virtualizarr/writers/icechunk.py @@ -7,7 +7,7 @@ import xarray as xr from xarray.backends.zarr import ZarrStore as XarrayZarrStore from xarray.backends.zarr import encode_zarr_attr_value -from zarr import Array, Group +from zarr import Array, Group, open_group from zarr.core.buffer import default_buffer_prototype from zarr.core.chunk_key_encodings import DefaultChunkKeyEncoding from zarr.core.sync import sync @@ -34,12 +34,39 @@ ENCODING_KEYS = {"_FillValue", "missing_value", "scale_factor", "add_offset"} +VALID_MODES = ("w", "w-", "a") + + +def _resolve_mode( + mode: Optional[Literal["w", "w-", "a"]], + append_dim: Optional[str] = None, + region: object = None, +) -> Literal["w", "w-", "a", "r+"]: + """Validate ``mode`` and resolve it to the effective zarr group-open mode.""" + if not isinstance(mode, (type(None), str)): + raise TypeError(f"mode: expected type Optional[str], but got type {type(mode)}") + + if mode is not None and mode not in VALID_MODES: + raise ValueError(f"mode: expected one of {VALID_MODES}, but got {mode!r}") + + if append_dim or region: + if mode in ("w", "w-"): + raise ValueError( + f"mode {mode!r} cannot be used together with append_dim or region, " + "which require opening an existing group" + ) + # appending or writing to a region requires the group (and arrays) to already exist + return "r+" + + return mode or "w-" + def virtual_dataset_to_icechunk( vds: xr.Dataset, store: "IcechunkStore", *, group: Optional[str] = None, + mode: Optional[Literal["w", "w-", "a"]] = None, append_dim: Optional[str] = None, region: Optional[Literal["auto"] | Mapping[str, Literal["auto"] | slice]] = None, validate_containers: bool = True, @@ -58,6 +85,16 @@ def virtual_dataset_to_icechunk( Store to write the dataset to, which must not be read-only. group Path to the group in which to store the dataset, defaulting to the root group. + mode + How to handle a pre-existing group at the target path: + + - ``"w-"``: create the group, raising a ``ContainsGroupError`` if it already exists. + - ``"w"``: create the group, overwriting any existing contents at that path. + - ``"a"``: open the group if it exists (keeping existing arrays), otherwise create it. + - ``None`` (default): equivalent to ``"w-"``, unless ``append_dim`` or ``region`` + is given, in which case the existing group is opened. + + ``mode="w"`` and ``mode="w-"`` are incompatible with ``append_dim`` and ``region``. append_dim Name of the dimension along which to append data. If provided, the dataset must have a dimension with this name. @@ -87,7 +124,6 @@ def virtual_dataset_to_icechunk( """ try: from icechunk import IcechunkStore # type: ignore[import-not-found] - from zarr import Group # type: ignore[import-untyped] from zarr.storage import StorePath # type: ignore[import-untyped] except ImportError: raise ImportError( @@ -104,6 +140,8 @@ def virtual_dataset_to_icechunk( f"group: expected type Optional[str], but got type {type(group)}" ) + open_mode = _resolve_mode(mode, append_dim=append_dim, region=region) + if not isinstance(append_dim, (type(None), str)): raise TypeError( f"append_dim: expected type Optional[str], but got type {type(append_dim)}" @@ -134,11 +172,9 @@ def virtual_dataset_to_icechunk( if validate_containers: validate_virtual_chunk_containers(store.session.config, [vds]) - if append_dim or region: - group_object = Group.open(store=store_path, zarr_format=3) - else: - # create the group if it doesn't already exist - group_object = Group.from_store(store=store_path, zarr_format=3) + group_object = open_group( + store_path, mode=open_mode, zarr_format=3, use_consolidated=False + ) write_virtual_dataset_to_icechunk_group( vds=vds, @@ -154,6 +190,7 @@ def virtual_datatree_to_icechunk( vdt: xr.DataTree, store: "IcechunkStore", *, + mode: Optional[Literal["w", "w-", "a"]] = None, write_inherited_coords: bool = False, validate_containers: bool = True, last_updated_at: datetime | None = None, @@ -170,6 +207,13 @@ def virtual_datatree_to_icechunk( DataTree to write to an Icechunk store. Can contain both "virtual" variables (backed by ManifestArray objects) and "loadable" variables (backed by numpy arrays). store Store to write the dataset to, which must not be read-only. + mode + How to handle pre-existing groups at the target paths: + + - ``"w-"`` or ``None`` (default): create each group, raising a + ``ContainsGroupError`` if it already exists. + - ``"w"``: create each group, overwriting any existing contents at that path. + - ``"a"``: open each group if it exists (keeping existing arrays), otherwise create it. write_inherited_coords If ``True``, replicate inherited coordinates on all descendant nodes of the tree. Otherwise, only write coordinates at the level at which they are @@ -197,7 +241,6 @@ def virtual_datatree_to_icechunk( """ try: from icechunk import IcechunkStore # type: ignore[import-not-found] - from zarr import Group # type: ignore[import-untyped] from zarr.storage import StorePath # type: ignore[import-untyped] except ImportError: raise ImportError( @@ -209,6 +252,10 @@ def virtual_datatree_to_icechunk( f"store: expected type IcechunkStore, but got type {type(store)}" ) + open_mode = _resolve_mode( + mode, append_dim=kwargs.get("append_dim"), region=kwargs.get("region") + ) + if not isinstance(last_updated_at, (type(None), datetime)): raise TypeError( "last_updated_at: expected type datetime," @@ -238,12 +285,9 @@ def get_store_path(subtree, vdt) -> StorePath: # TODO this serial loop could be slow writing lots of groups to high-latency store, see https://github.com/pydata/xarray/issues/9455 for store_path, vds in paths_and_virtual_datasets: - if kwargs.get("append_dim") or kwargs.get("region"): - # both require the group (and arrays) to already exist - group = Group.open(store=store_path, zarr_format=3) - else: - # create the group if it doesn't already exist - group = Group.from_store(store=store_path, zarr_format=3) + group = open_group( + store_path, mode=open_mode, zarr_format=3, use_consolidated=False + ) write_virtual_dataset_to_icechunk_group( vds=vds,