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/guide.md

@@ -31,7 +31,8 @@ exist.) Both the CLI command and the Python function can be run without any furt
 configuration beyond the required `zarr_format` setting. The target dataset path must
 point to a directory that will contain a Zarr
 group to be created and updated. The slice dataset paths may be provided as Zarr as 
-well or in other data formats supported by the [xarray.open_dataset()](https://docs.xarray.dev/en/stable/generated/xarray.open_dataset.html) 
+well or in other data formats supported by the
+[xarray.open_dataset()](https://docs.xarray.dev/en/stable/generated/xarray.open_dataset.html)
 function. Because we provided no additional configuration, the default append dimension
 `time` is used above.
 
@@ -45,7 +46,8 @@ in section [_Data I/O_](#data-io) below.
     New configurations must set `zarr_format`; use `zarr_format: 3` for the
     recommended [Zarr storage specification 3](https://zarr.readthedocs.io/en/stable/spec/v3.html),
     or `zarr_format: 2` for Zarr v2 compatibility. `zarr_version: 2` is a
-    deprecated legacy setting; it cannot select Zarr v3. The target's format is fixed when it is created; use the
+    deprecated legacy setting; it cannot select Zarr v3. The target's format is
+    fixed when it is created; use the
     same setting for all later appends.
 
 The tool takes care of generating the target dataset from slice datasets, but doesn't 
@@ -119,7 +121,7 @@ between the labels of a coordinate variable associated with the append dimension
 Its value can be a number for numerical labels or a timedelta value of the form
 `<count><unit>` for date/time labels. In the latter case `<count>` is an integer 
 and `<units>` is one of the possible
-[numpy datetime units](https://numpy.org/doc/stable/reference/arrays.datetime.html#datetime-units), 
+[numpy datetime units](https://numpy.org/doc/stable/reference/arrays.datetime.html#datetime-units),
 for example, `8h` (8 hours) or `2D` (two days). Numerical and timedelta values
 may be negative. `append_step` can also take the two special values `"+"` and 
 `"-"`. In this case it is just verified that the append labels are monotonically 
@@ -566,6 +568,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 +650,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,16 +719,21 @@ 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
+    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 
     and each one binds resources such as open file handles. Consider using a 
     _slice source_ then, see below.
@@ -735,8 +764,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 +803,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`).
 
@@ -981,7 +1010,8 @@ just provide its name:
 
 If you also want logging output in a file or using a different format or if you want to
 see logging output of other Python modules, you can configure Python's logging 
-system following the [logging dictionary schema](https://docs.python.org/3/library/logging.config.html#logging-config-dictschema).
+system following the [logging dictionary
+schema](https://docs.python.org/3/library/logging.config.html#logging-config-dictschema).
 Given here is an example that logs `zappend`'s output to the console using 
 the `INFO` level (same as `"logging": true`):
 

docs/howdoi.md

@@ -54,7 +54,8 @@ slice_ds = image_ds.expand_dims("time", axis=0)
 slice_ds.coords["time"] = xr.Dataset(np.array([slice_time]), dims="time")
 ```
 
-See also [How do I create datacubes from a directory of GeoTIFFs](#create-datacubes-from-a-directory-of-geotiffs) 
+See also [How do I create datacubes from a directory of
+GeoTIFFs](#create-datacubes-from-a-directory-of-geotiffs)
 above. 
 
 ## ... dynamically update global metadata attributes
@@ -125,4 +126,5 @@ configuration setting `force_new`.
 
 ### Error `Append dimension 'foo' not found in dataset`
 
-Refer to [How do I create datacubes from datasets without append dimension](#create-datacubes-from-datasets-without-append-dimension).
+Refer to [How do I create datacubes from datasets without append
+dimension](#create-datacubes-from-datasets-without-append-dimension).

docs/index.md

@@ -4,7 +4,8 @@
 
 `zappend` is a tool written in Python that is used for robustly creating and 
 updating Zarr datacubes from smaller dataset slices. It is built on top of the 
-awesome Python packages [xarray](https://docs.xarray.dev/) and [zarr](https://zarr.readthedocs.io/).
+awesome Python packages [xarray](https://docs.xarray.dev/) and
+[zarr](https://zarr.readthedocs.io/).
 
 ## Motivation
 
@@ -42,15 +43,22 @@ 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
 
-At its core, `zappend` calls the [to_zarr()](https://docs.xarray.dev/en/stable/generated/xarray.Dataset.to_zarr.html#xarray-dataset-to-zarr) method of [xarray.Dataset](https://docs.xarray.dev/en/stable/generated/xarray.Dataset.html) 
+At its core, `zappend` calls the
+[to_zarr()](https://docs.xarray.dev/en/stable/generated/xarray.Dataset.to_zarr.html#xarray-dataset-to-zarr)
+method of [xarray.Dataset](https://docs.xarray.dev/en/stable/generated/xarray.Dataset.html)
 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,13 +7,18 @@ 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).
+  [xcube.core.mldataset.MultiLevelDataset]
+  (https://xcube.readthedocs.io/en/latest/mldatasets.html).
 * Target and slices are allowed to live in different filesystems.
 * The tool is configurable. The configuration defines 
     - the append dimension;

docs/start.md

@@ -109,7 +109,8 @@ zappend(
 )
 ```
 
-For the details, please refer to the section [_Slice Sources_](guide.md#slice-sources) in the 
+For the details, please refer to the section [_Slice Sources_](guide.md#slice-sources)
+in the
 [User Guide](guide.md).
 
 

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)})