Support writing xarray.DataTree to Icechunk store (#538)

Fixes #244

Author: chuckwondo

docs/api.rst

@@ -30,6 +30,7 @@ Serialization
 
     VirtualiZarrDatasetAccessor.to_kerchunk
     VirtualiZarrDatasetAccessor.to_icechunk
+    VirtualiZarrDataTreeAccessor.to_icechunk
 
 Information
 -----------

docs/releases.rst

@@ -14,6 +14,12 @@ New Features
   By `Tom Nicholas <https://github.com/TomNicholas>`_.
 - Added experimental :py:func:`open_virtual_mfdataset` function (:issue:`345`, :pull:`349`).
   By `Tom Nicholas <https://github.com/TomNicholas>`_.
+- Added :py:func:`datatree_to_icechunk` function for writing an ``xarray.DataTree`` to
+  an Icechunk store (:issue:`244`).  By `Chuck Daniels <https://github.com/chuckwondo>`_.
+- Added a ``.virtualize`` custom accessor to ``xarray.DataTree``, exposing the method
+  :py:meth:`xarray.DataTree.virtualize.to_icechunk()` for writing an ``xarray.DataTree``
+  to an Icechunk store (:issue:`244`).  By
+  `Chuck Daniels <https://github.com/chuckwondo>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~

virtualizarr/__init__.py

@@ -1,12 +1,24 @@
-from virtualizarr.manifests import ChunkManifest, ManifestArray  # type: ignore # noqa
-from virtualizarr.accessor import VirtualiZarrDatasetAccessor  # type: ignore # noqa
-from virtualizarr.backend import open_virtual_dataset, open_virtual_mfdataset  # noqa: F401
-
 from importlib.metadata import version as _version
 
+from virtualizarr.accessor import (
+    VirtualiZarrDatasetAccessor,
+    VirtualiZarrDataTreeAccessor,
+)
+from virtualizarr.backend import open_virtual_dataset, open_virtual_mfdataset
+from virtualizarr.manifests import ChunkManifest, ManifestArray
+
 try:
     __version__ = _version("virtualizarr")
 except Exception:
     # Local copy or not installed with setuptools.
     # Disable minimum version checks on downstream libraries.
     __version__ = "9999"
+
+__all__ = [
+    "ChunkManifest",
+    "ManifestArray",
+    "VirtualiZarrDatasetAccessor",
+    "VirtualiZarrDataTreeAccessor",
+    "open_virtual_dataset",
+    "open_virtual_mfdataset",
+]

virtualizarr/accessor.py

@@ -2,7 +2,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable, Literal, overload
 
-from xarray import Dataset, register_dataset_accessor
+import xarray as xr
 
 from virtualizarr.manifests import ManifestArray
 from virtualizarr.types.kerchunk import KerchunkStoreRefs
@@ -12,16 +12,16 @@
     from icechunk import IcechunkStore  # type: ignore[import-not-found]
 
 
-@register_dataset_accessor("virtualize")
+@xr.register_dataset_accessor("virtualize")
 class VirtualiZarrDatasetAccessor:
     """
     Xarray accessor for writing out virtual datasets to disk.
 
     Methods on this object are called via `ds.virtualize.{method}`.
     """
 
-    def __init__(self, ds: Dataset):
-        self.ds: Dataset = ds
+    def __init__(self, ds: xr.Dataset):
+        self.ds: xr.Dataset = ds
 
     def to_icechunk(
         self,
@@ -171,7 +171,7 @@ def to_kerchunk(
     def rename_paths(
         self,
         new: str | Callable[[str], str],
-    ) -> Dataset:
+    ) -> xr.Dataset:
         """
         Rename paths to chunks in every ManifestArray in this dataset.
 
@@ -232,3 +232,76 @@ def nbytes(self) -> int:
             else var.nbytes
             for var in self.ds.variables.values()
         )
+
+
+@xr.register_datatree_accessor("virtualize")
+class VirtualiZarrDataTreeAccessor:
+    """
+    Xarray accessor for writing out virtual datatrees to disk.
+
+    Methods on this object are called via `dt.virtualize.{method}`.
+    """
+
+    def __init__(self, dt: xr.DataTree):
+        self.dt = dt
+
+    def to_icechunk(
+        self,
+        store: "IcechunkStore",
+        *,
+        write_inherited_coords: bool = False,
+        last_updated_at: datetime | None = None,
+    ) -> None:
+        """
+        Write an xarray DataTree to an Icechunk store.
+
+        Any variables backed by ManifestArray objects will be be written as virtual
+        references. Any other variables will be loaded into memory before their binary
+        chunk data is written into the store.
+
+        If ``last_updated_at`` is provided, it will be used as a checksum for any
+        virtual chunks written to the store with this operation.  At read time, if any
+        of the virtual chunks have been updated since this provided datetime, an error
+        will be raised.  This protects against reading outdated virtual chunks that have
+        been updated since the last read.  When not provided, no check is performed.
+        This value is stored in Icechunk with seconds precision, so be sure to take that
+        into account when providing this value.
+
+        Parameters
+        ----------
+        store: IcechunkStore
+            Store to write dataset into.
+        write_inherited_coords : bool, default: False
+            If ``True``, replicate inherited coordinates on all descendant nodes.
+            Otherwise, only write coordinates at the level at which they are
+            originally defined. This saves disk space, but requires opening the
+            full tree to load inherited coordinates.
+        last_updated_at: datetime, optional
+            Datetime to use as a checksum for any virtual chunks written to the store
+            with this operation.  When not provided, no check is performed.
+
+        Raises
+        ------
+        ValueError
+            If the store is read-only.
+
+        Examples
+        --------
+        To ensure an error is raised if the files containing referenced virtual chunks
+        are modified at any time from now on, pass the current time to
+        ``last_updated_at``.
+
+        >>> from datetime import datetime
+        >>> vdt.virtualize.to_icechunk(  # doctest: +SKIP
+        ...     icechunkstore,
+        ...     last_updated_at=datetime.now(),
+        ... )
+        """
+        from virtualizarr.writers.icechunk import datatree_to_icechunk
+
+        datatree_to_icechunk(
+            self.dt,
+            store,
+            write_inherited_coords=write_inherited_coords,
+            last_updated_at=last_updated_at,
+        )

virtualizarr/tests/test_integration.py

@@ -1,6 +1,7 @@
+from collections.abc import Mapping
 from os.path import relpath
 from pathlib import Path
-from typing import Callable, Concatenate, TypeAlias
+from typing import Any, Callable, Concatenate, TypeAlias, overload
 
 import numpy as np
 import pytest
@@ -23,7 +24,9 @@
     dataset_from_kerchunk_refs,
 )
 
-RoundtripFunction: TypeAlias = Callable[Concatenate[xr.Dataset, Path, ...], xr.Dataset]
+RoundtripFunction: TypeAlias = Callable[
+    Concatenate[xr.Dataset | xr.DataTree, Path, ...], xr.Dataset | xr.DataTree
+]
 
 
 def test_kerchunk_roundtrip_in_memory_no_concat(array_v3_metadata):
@@ -111,7 +114,22 @@ def roundtrip_as_kerchunk_parquet(vds: xr.Dataset, tmpdir, **kwargs):
     return xr.open_dataset(f"{tmpdir}/refs.parquet", engine="kerchunk", **kwargs)
 
 
-def roundtrip_as_in_memory_icechunk(vds: xr.Dataset, tmpdir, **kwargs):
+@overload
+def roundtrip_as_in_memory_icechunk(
+    vdata: xr.Dataset, tmp_path: Path, **kwargs
+) -> xr.Dataset: ...
+@overload
+def roundtrip_as_in_memory_icechunk(
+    vdata: xr.DataTree, tmp_path: Path, **kwargs
+) -> xr.DataTree: ...
+
+
+def roundtrip_as_in_memory_icechunk(
+    vdata: xr.Dataset | xr.DataTree,
+    tmp_path: Path,
+    virtualize_kwargs: Mapping[str, Any] | None = None,
+    **kwargs,
+) -> xr.Dataset | xr.DataTree:
     from icechunk import Repository, Storage
 
     # create an in-memory icechunk store
@@ -120,7 +138,17 @@ def roundtrip_as_in_memory_icechunk(vds: xr.Dataset, tmpdir, **kwargs):
     session = repo.writable_session("main")
 
     # write those references to an icechunk store
-    vds.virtualize.to_icechunk(session.store)
+    vdata.virtualize.to_icechunk(session.store, **(virtualize_kwargs or {}))
+
+    if isinstance(vdata, xr.DataTree):
+        # read the dataset from icechunk
+        return xr.open_datatree(
+            session.store,  # type: ignore
+            engine="zarr",
+            zarr_format=3,
+            consolidated=False,
+            **kwargs,
+        )
 
     # read the dataset from icechunk
     return xr.open_zarr(session.store, zarr_format=3, consolidated=False, **kwargs)
@@ -219,16 +247,14 @@ def test_kerchunk_roundtrip_concat(
 
                 roundtrip = roundtrip_func(vds, tmp_path, decode_times=decode_times)
 
-                if decode_times is False:
-                    # assert all_close to original dataset
-                    xrt.assert_allclose(roundtrip, ds)
+                # assert all_close to original dataset
+                xrt.assert_allclose(roundtrip, ds)
 
-                    # assert coordinate attributes are maintained
-                    for coord in ds.coords:
-                        assert ds.coords[coord].attrs == roundtrip.coords[coord].attrs
-                else:
-                    # they are very very close! But assert_allclose doesn't seem to work on datetimes
-                    assert (roundtrip.time - ds.time).sum() == 0
+                # assert coordinate attributes are maintained
+                for coord in ds.coords:
+                    assert ds.coords[coord].attrs == roundtrip.coords[coord].attrs
+
+                if decode_times:
                     assert roundtrip.time.dtype == ds.time.dtype
                     assert roundtrip.time.encoding["units"] == ds.time.encoding["units"]
                     assert (
@@ -303,6 +329,102 @@ def test_datetime64_dtype_fill_value(
         assert roundtrip.a.attrs == vds.a.attrs
 
 
+@parametrize_over_hdf_backends
+@pytest.mark.parametrize(
+    "roundtrip_func", [roundtrip_as_in_memory_icechunk] if has_icechunk else []
+)
+@pytest.mark.parametrize("decode_times", (False, True))
+@pytest.mark.parametrize("time_vars", ([], ["time"]))
+@pytest.mark.parametrize("inherit", (False, True))
+def test_datatree_roundtrip(
+    tmp_path: Path,
+    roundtrip_func: RoundtripFunction,
+    hdf_backend: type[VirtualBackend],
+    decode_times: bool,
+    time_vars: list[str],
+    inherit: bool,
+):
+    # set up example xarray dataset
+    with xr.tutorial.open_dataset("air_temperature", decode_times=decode_times) as ds:
+        # split into two datasets
+        ds1 = ds.isel(time=slice(None, 1460))
+        ds2 = ds.isel(time=slice(1460, None))
+
+        # save it to disk as netCDF (in temporary directory)
+        air1_nc_path = tmp_path / "air1.nc"
+        air2_nc_path = tmp_path / "air2.nc"
+        ds1.to_netcdf(air1_nc_path)
+        ds2.to_netcdf(air2_nc_path)
+
+        # use open_dataset_via_kerchunk to read it as references
+        with (
+            open_virtual_dataset(
+                str(air1_nc_path),
+                loadable_variables=time_vars,
+                decode_times=decode_times,
+                backend=hdf_backend,
+            ) as vds1,
+            open_virtual_dataset(
+                str(air2_nc_path),
+                loadable_variables=time_vars,
+                decode_times=decode_times,
+                backend=hdf_backend,
+            ) as vds2,
+        ):
+            if not decode_times or not time_vars:
+                assert vds1.time.dtype == np.dtype("float32")
+                assert vds2.time.dtype == np.dtype("float32")
+            else:
+                assert vds1.time.dtype == np.dtype("<M8[ns]")
+                assert vds2.time.dtype == np.dtype("<M8[ns]")
+                assert "units" in vds1.time.encoding
+                assert "units" in vds2.time.encoding
+                assert "calendar" in vds1.time.encoding
+                assert "calendar" in vds2.time.encoding
+
+            vdt = xr.DataTree.from_dict({"/vds1": vds1, "/nested/vds2": vds2})
+
+            with roundtrip_func(
+                vdt,
+                tmp_path,
+                virtualize_kwargs=dict(write_inherited_coords=inherit),
+                decode_times=decode_times,
+            ) as roundtrip:
+                assert isinstance(roundtrip, xr.DataTree)
+
+                # assert all_close to original dataset
+                roundtrip_vds1 = roundtrip["/vds1"].to_dataset()
+                roundtrip_vds2 = roundtrip["/nested/vds2"].to_dataset()
+                xrt.assert_allclose(roundtrip_vds1, ds1)
+                xrt.assert_allclose(roundtrip_vds2, ds2)
+
+                # assert coordinate attributes are maintained
+                for coord in ds1.coords:
+                    assert ds1.coords[coord].attrs == roundtrip_vds1.coords[coord].attrs
+                for coord in ds2.coords:
+                    assert ds2.coords[coord].attrs == roundtrip_vds2.coords[coord].attrs
+
+                if decode_times:
+                    assert roundtrip_vds1.time.dtype == ds1.time.dtype
+                    assert roundtrip_vds2.time.dtype == ds2.time.dtype
+                    assert (
+                        roundtrip_vds1.time.encoding["units"]
+                        == ds1.time.encoding["units"]
+                    )
+                    assert (
+                        roundtrip_vds2.time.encoding["units"]
+                        == ds2.time.encoding["units"]
+                    )
+                    assert (
+                        roundtrip_vds1.time.encoding["calendar"]
+                        == ds1.time.encoding["calendar"]
+                    )
+                    assert (
+                        roundtrip_vds2.time.encoding["calendar"]
+                        == ds2.time.encoding["calendar"]
+                    )
+
+
 @parametrize_over_hdf_backends
 def test_open_scalar_variable(tmp_path: Path, hdf_backend: type[VirtualBackend]):
     # regression test for GH issue #100