Surface mode in vz.to_icechunk, keeping current defaults

docs/about/releases.md

@@ -4,11 +4,16 @@
 
 ### 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
 
 ### 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

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
 

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,

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",
@@ -996,9 +1109,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",

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,7 +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:
-        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,