Support nested Zarr datasets

CHANGES.md

@@ -1,5 +1,8 @@
 ## Version 0.9.0 (in development)
 
+* Added support for nested Zarr structures represented by `xarray.DataTree`,
+  including atomic rollback of complete DataTree slices and optional per-group
+  variable encoding configuration. (#113)
 * Migrated to use `zarr >=3`. Supports writing Zarr v2 and v3 stores, with
   Zarr v3 as the recommended format for new configurations. The `zarr_version`
   setting is deprecated in favor of the now-required `zarr_format` and only

README.md

@@ -51,7 +51,8 @@ The `zappend` tool provides the following features:
   created, effectively preventing concurrent dataset modifications.
 * **Transaction-based dataset appends**: On failure during an append step, 
   the transaction is rolled back, so that the target dataset remains valid and 
-  preserves its integrity.
+  preserves its integrity. A nested `xarray.DataTree` slice is rolled back as
+  a whole if appending any of its groups fails.
 * **Filesystem transparency**: The target dataset may be generated and updated 
   in any writable filesystems supported by the 
   [fsspec](https://filesystem-spec.readthedocs.io/) package. 
@@ -64,7 +65,8 @@ The `zappend` tool provides the following features:
   [`zappend`](cli.md) command or from Python. When used from Python using the 
   [`zappend()`](api.md) function, slice datasets can be passed as local file 
   paths, URIs, as datasets of type 
-  [xarray.Dataset](https://docs.xarray.dev/en/stable/generated/xarray.Dataset.html), or as custom
+  [xarray.Dataset](https://docs.xarray.dev/en/stable/generated/xarray.Dataset.html),
+  [xarray.DataTree](https://docs.xarray.dev/en/stable/generated/xarray.DataTree.html), or as custom
   [slice sources](https://bcdev.github.io/zappend/guide/#slice-sources).
 
 

docs/cli.md

@@ -10,6 +10,8 @@ Usage: zappend [OPTIONS] [SLICES]...
 
   The zappend command concatenates the dataset SLICES along a given append
   dimension, e.g., "time" (the default) for geospatial satellite observations.
+  A nested Zarr slice is opened as an xarray.DataTree and its groups are
+  appended as one transaction.
   Each append step is atomic, that is, the append operation is a transaction
   that can be rolled back, in case the append operation fails. This ensures
   integrity of the target data cube given by TARGET or in CONFIG.

docs/config.md

@@ -113,6 +113,30 @@ Variable metadata.
   - `attrs`:
     Arbitrary variable metadata attributes.
 
+### `group_variables`
+
+Type _object_.
+Optional variable-configuration overrides for individual `xarray.DataTree`
+groups. Keys are absolute Zarr group paths such as `"/group_0"` or `"/group_1"`; each
+value uses the same syntax as [`variables`](#variables).
+
+For a group, its configuration is merged over the global `variables`
+configuration. If a group has no entry here, normal zappend behavior applies:
+the global `variables` setting is used if present, otherwise encoding such as
+chunk sizes is taken from that group's first slice.
+
+```yaml
+group_variables:
+  "/group_0":
+    var_0:
+      encoding:
+        chunks: [31, 450, 450]
+  "/group_1":
+    var_0:
+      encoding:
+        chunks: [31, 450, 450]
+```
+
 ### `attrs`
 
 Type _object_.
@@ -201,7 +225,7 @@ Must be one of the following:
 ### `slice_source`
 
 Type _string_.
-The fully qualified name of a class or function that receives a slice item as argument(s) and provides the slice dataset. If a class is given, it must be derived from `zappend.api.SliceSource`. If the function is a context manager, it must yield an `xarray.Dataset`. If a plain function is given, it must return any valid slice item type. Refer to the user guide for more information.
+The fully qualified name of a class or function that receives a slice item as argument(s) and provides a slice object. If a class is given, it must be derived from `zappend.api.SliceSource`. If the function is a context manager, it must yield an `xarray.Dataset` or `xarray.DataTree`. If a plain function is given, it must return any valid slice item type. Refer to the user guide for more information.
 
 ### `slice_source_kwargs`
 

docs/guide.md

@@ -566,6 +566,10 @@ immediately after a target dataset modification failed. The rollback include res
 all changed files and removing added files. After the rollback you can analyse what
 went wrong and try to continue appending slices at the point it failed.
 
+For an `xarray.DataTree` slice, this transaction spans all data-bearing Zarr
+groups in the tree. If appending one group fails, changes already made to the
+other groups for that same slice are rolled back too.
+
 To allow for rollbacks, a slice append operation is treated as a transaction,
 hence temporary files must be written, e.g., to record required rollback actions and to
 save backup files with the original data. The location of the temporary transaction 
@@ -644,8 +648,26 @@ to `zappend`. Slice datasets can be provided in various ways.
   of slice items or provide a Python generator that provides the slice items.
 
 Each slice item can be a local file path or URI of type `str` or `FileObj`, 
-a dataset of type `xarray.Dataset`, or a `SliceSource` object explained in more detail 
-below.
+a dataset of type `xarray.Dataset`, a nested dataset of type `xarray.DataTree`,
+or a `SliceSource` object explained in more detail below.
+
+#### Nested Zarr Groups and DataTree
+
+Nested Zarr inputs are represented by
+[`xarray.DataTree`](https://docs.xarray.dev/en/stable/generated/xarray.DataTree.html).
+Pass a DataTree directly, return one from a `SliceSource`, or pass a path to a
+nested Zarr store: zappend opens the latter as a DataTree automatically.
+
+All DataTree slices in one run must have the same group paths. DataTree and
+Dataset slices cannot be mixed in one run. Each data-bearing group is appended
+along the configured append dimension, and the complete DataTree slice is
+rolled back if any group append fails.
+
+Use `group_variables` to override variable encoding for individual absolute
+group paths. Group settings override global `variables`; groups without an
+override use global `variables`, or the encoding of their first slice when no
+global setting exists. See the [configuration reference](config.md#group_variables)
+and `examples/datatree.ipynb` for a complete example.
 
 #### Paths and URIs
 
@@ -695,14 +717,18 @@ An alternative to providing the slice dataset as path or URI is using the
 `zappend.api.FileObj` class, which combines a URI with dedicated filesystem 
 storage options.
 
-#### Dataset Objects
+#### Dataset and DataTree Objects
 
 You can also use dataset objects of type 
 [`xarray.Dataset`](https://docs.xarray.dev/en/stable/generated/xarray.Dataset.html) 
 as slice item. Such objects may originate from opening datasets from some storage, e.g., 
 `xarray.open_dataset(slice_store, ...)` or by composing, aggregating, resampling 
 slice datasets from other datasets and dataset variables.
 
+Nested inputs can instead be passed as
+[`xarray.DataTree`](https://docs.xarray.dev/en/stable/generated/xarray.DataTree.html)
+objects. The same resource-management and persistence considerations apply.
+
 !!! note "Datasets are not closed automatically"
     If you pass [`xarray.Dataset`](https://docs.xarray.dev/en/stable/generated/xarray.Dataset.html) objects to `zappend` they will not be
     automatically closed. This may become be issue, if you have many datasets 
@@ -735,8 +761,8 @@ at the cost of additional i/o. It therefore defaults to `false`.
 
 A slice source gives you full control about how a slice dataset is created, loaded, 
 or generated and how its bound resources, if any, are released. In its simplest form, 
-a slice source is a plain Python function that can take any arguments and returns 
-an `xarray.Dataset`:
+a slice source is a plain Python function that can take any arguments and returns
+an `xarray.Dataset` or `xarray.DataTree`:
 
 ```python
 import xarray as xr
@@ -774,8 +800,8 @@ def get_dataset(path: str) -> xr.Dataset:
 You can also implement your slice source as a class derived from the abstract 
 `zappend.api.SliceSource` class. Its interface methods are:
 
-* `get_dataset()`: a zero-argument method that returns the slice dataset of type 
-  `xarray.Dataset`. You must implement this abstract method. 
+* `get_dataset()`: a zero-argument method that returns the slice object of type
+  `xarray.Dataset` or `xarray.DataTree`. You must implement this abstract method.
 * `close()`: Optional method. Put your cleanup code here. 
   (in zappend < v0.7, the `close` method was called `dispose`).
 

docs/index.md

@@ -42,7 +42,8 @@ The `zappend` tool provides the following features:
   [`zappend`](cli.md) command or from Python. When used from Python using the 
   [`zappend()`](api.md) function, slice datasets can be passed as local file 
   paths, URIs, as datasets of type 
-  [xarray.Dataset](https://docs.xarray.dev/en/stable/generated/xarray.Dataset.html), or as custom
+  [xarray.Dataset](https://docs.xarray.dev/en/stable/generated/xarray.Dataset.html),
+  [xarray.DataTree](https://docs.xarray.dev/en/stable/generated/xarray.DataTree.html), or as custom
   [slice sources](guide.md#slice-sources).
 
 ## How It Works
@@ -51,6 +52,9 @@ At its core, `zappend` calls the [to_zarr()](https://docs.xarray.dev/en/stable/g
 for each dataset slice it receives and either creates the target dataset if it does 
 not exist yet or updates it with the current slice, if it already exists.
 
+For nested Zarr stores, it applies the same operation to each data-bearing
+`xarray.DataTree` group as one atomic slice transaction.
+
 If there is no target dataset yet, `zappend` does the following:
 
 * get target dataset outline and encoding from configuration and first slice dataset;

docs/requirements.md

@@ -7,11 +7,15 @@ of the `zappend` tool.
 
 * Create a target Zarr dataset by appending Zarr dataset slices along a 
   given *append dimension*, usually `time`.   
+* Create a target nested Zarr store by appending compatible
+  `xarray.DataTree` slices along the append dimension of every data-bearing
+  group.
 * The tool takes care of modifying the target dataset using the slices,
   but doesn't care how the slice datasets are created.
 * Slice datasets may be given as URIs with storage options or as 
   in-memory datasets of type 
   [xarray.Dataset](https://docs.xarray.dev/en/stable/generated/xarray.Dataset.html)
+  or [xarray.DataTree](https://docs.xarray.dev/en/stable/generated/xarray.DataTree.html)
   or 
   [xcube.core.mldataset.MultiLevelDataset](https://xcube.readthedocs.io/en/latest/mldatasets.html).
 * Target and slices are allowed to live in different filesystems.

examples/create_inputs.py

@@ -90,3 +90,45 @@ def make_netcdf_slices(n_slices: int) -> None:
             engine="h5netcdf",
             encoding=encoding,
         )
+
+
+def make_datatree_slices(n_slices: int) -> None:
+    """Create *n_slices* daily nested Zarr v2 inputs with two resolutions."""
+    times = START_TIME + np.arange(n_slices) * np.timedelta64(1, "D")
+    directory = INPUT_DIR / "datatree"
+    directory.mkdir(parents=True, exist_ok=True)
+    for index, time in enumerate(times):
+        tree = xr.DataTree.from_dict(
+            {
+                "/": xr.Dataset(
+                    attrs={
+                        "title": "Zappend nested example",
+                        "description": "Two spatial resolutions of daily data",
+                    }
+                ),
+                "/0": _make_datatree_dataset(time, index, 1800, 3600),
+                "/1": _make_datatree_dataset(time, index, 900, 1800),
+            }
+        )
+        tree.to_zarr(
+            directory / f"slice-{time!s}.zarr",
+            mode="w",
+            zarr_format=2,
+            consolidated=True,
+        )
+
+
+def _make_datatree_dataset(
+    time: np.datetime64, index: int, lat_size: int, lon_size: int
+) -> xr.Dataset:
+    lat = np.linspace(-90, 90, lat_size, endpoint=False, dtype="float32")
+    lon = np.linspace(-180, 180, lon_size, endpoint=False, dtype="float32")
+    data = np.full((1, lat_size, lon_size), index, dtype="uint8")
+    return xr.Dataset(
+        data_vars={
+            "temperature": (("time", "lat", "lon"), data.astype("float32")),
+            "quality_flag": (("time", "lat", "lon"), data),
+        },
+        coords={"time": [time], "lat": lat, "lon": lon},
+        attrs={"spatial_resolution": f"{360 / lon_size:g} degrees"},
+    ).chunk({"time": 1, "lat": min(450, lat_size), "lon": min(450, lon_size)})