diff --git a/docs/about/releases.md b/docs/about/releases.md index 31f671e7..3e8e8459 100644 --- a/docs/about/releases.md +++ b/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 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 01e7465a..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", @@ -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", diff --git a/virtualizarr/writers/icechunk.py b/virtualizarr/writers/icechunk.py index 7c6741ff..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,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,