diff --git a/CHANGES.md b/CHANGES.md index 502910e..2dec497 100644 --- a/CHANGES.md +++ b/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 diff --git a/README.md b/README.md index 6b4d53d..e3eee65 100644 --- a/README.md +++ b/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). diff --git a/docs/cli.md b/docs/cli.md index 243bc5b..fe8a70f 100644 --- a/docs/cli.md +++ b/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. diff --git a/docs/config.md b/docs/config.md index 47af003..b6a5f8a 100644 --- a/docs/config.md +++ b/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` diff --git a/docs/guide.md b/docs/guide.md index 104414a..6423b32 100644 --- a/docs/guide.md +++ b/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,7 +717,7 @@ 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) @@ -703,6 +725,10 @@ as slice item. Such objects may originate from opening datasets from some storag `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`). diff --git a/docs/index.md b/docs/index.md index 450fe8a..ed96536 100644 --- a/docs/index.md +++ b/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; diff --git a/docs/requirements.md b/docs/requirements.md index 88199c5..d458337 100644 --- a/docs/requirements.md +++ b/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. diff --git a/examples/create_inputs.py b/examples/create_inputs.py index 3809ed3..27b63a3 100644 --- a/examples/create_inputs.py +++ b/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)}) diff --git a/examples/datatree.ipynb b/examples/datatree.ipynb new file mode 100644 index 0000000..750e636 --- /dev/null +++ b/examples/datatree.ipynb @@ -0,0 +1,1789 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "af9386c9aea0a587", + "metadata": {}, + "source": [ + "# Appending nested Zarr groups with DataTree\n", + "\n", + "This notebook generates daily nested Zarr v2 inputs with two spatial-resolution groups, `/0` and `/1`. `zappend` opens each input as an `xarray.DataTree` and appends all groups as one transactional slice: if one group fails, the complete daily tree is rolled back." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "9f9e2f0a0d05b730", + "metadata": { + "ExecuteTime": { + "end_time": "2026-10-30T11:23:55.980707308Z", + "start_time": "2026-10-30T11:23:55.001287664Z" + } + }, + "outputs": [], + "source": [ + "import os\n", + "import shutil\n", + "from pathlib import Path\n", + "\n", + "import xarray as xr\n", + "from create_inputs import make_datatree_slices\n", + "\n", + "from zappend.api import zappend" + ] + }, + { + "cell_type": "markdown", + "id": "e784207fd4840b95", + "metadata": {}, + "source": [ + "## Generate and collect nested daily inputs\n", + "\n", + "Each input has a metadata-only root group and two data-bearing child groups. Both carry the same daily `time` coordinate, but `/0` is 180 by 360 pixels and `/1` is 90 by 180 pixels." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "819ffbe9089aee7b", + "metadata": { + "ExecuteTime": { + "end_time": "2026-10-30T11:24:06.073416043Z", + "start_time": "2026-10-30T11:23:55.988596686Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "CPU times: user 10.8 s, sys: 2.52 s, total: 13.4 s\n", + "Wall time: 10.1 s\n" + ] + }, + { + "data": { + "text/plain": [ + "['inputs/datatree/slice-2024-01-01.zarr',\n", + " 'inputs/datatree/slice-2024-01-02.zarr',\n", + " 'inputs/datatree/slice-2024-01-03.zarr']" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "%%time\n", + "make_datatree_slices(31)\n", + "\n", + "input_dir = Path('inputs/datatree')\n", + "output_dir = Path('outputs')\n", + "output_dir.mkdir(exist_ok=True)\n", + "slices = sorted(str(path) for path in input_dir.glob('*.zarr'))\n", + "assert len(slices) == 31\n", + "slices[:3]" + ] + }, + { + "cell_type": "markdown", + "id": "3cc4d9f1cf156567", + "metadata": {}, + "source": [ + "## Default nested layout\n", + "\n", + "The first target keeps each group's encoding from its first slice. Because the inputs use one day per time chunk, this layout is suited to maps and single-day access." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "cfad5c1fa9a2b5d6", + "metadata": { + "ExecuteTime": { + "end_time": "2026-10-30T11:24:20.959529751Z", + "start_time": "2026-10-30T11:24:06.089167470Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2026-10-30 12:24:06,186 WARNING Setting 'force_new' is enabled. This will permanently delete existing targets (no rollback).\n", + "2026-10-30 12:24:06,188 WARNING Permanently deleting outputs/datatree_example_space_opt.zarr\n", + "2026-10-30 12:24:06,274 INFO Opening slice dataset from inputs/datatree/slice-2024-01-01.zarr\n", + "2026-10-30 12:24:06,319 INFO Creating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:06,547 INFO Creating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:06,688 INFO Transaction completed.\n", + "2026-10-30 12:24:06,689 INFO Slice dataset inputs/datatree/slice-2024-01-01.zarr closed\n", + "2026-10-30 12:24:06,697 INFO Opening slice dataset from inputs/datatree/slice-2024-01-02.zarr\n", + "2026-10-30 12:24:06,743 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:07,031 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:07,148 INFO Transaction completed.\n", + "2026-10-30 12:24:07,149 INFO Slice dataset inputs/datatree/slice-2024-01-02.zarr closed\n", + "2026-10-30 12:24:07,154 INFO Opening slice dataset from inputs/datatree/slice-2024-01-03.zarr\n", + "2026-10-30 12:24:07,191 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:07,573 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:07,740 INFO Transaction completed.\n", + "2026-10-30 12:24:07,741 INFO Slice dataset inputs/datatree/slice-2024-01-03.zarr closed\n", + "2026-10-30 12:24:07,748 INFO Opening slice dataset from inputs/datatree/slice-2024-01-04.zarr\n", + "2026-10-30 12:24:07,796 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:08,046 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:08,173 INFO Transaction completed.\n", + "2026-10-30 12:24:08,174 INFO Slice dataset inputs/datatree/slice-2024-01-04.zarr closed\n", + "2026-10-30 12:24:08,180 INFO Opening slice dataset from inputs/datatree/slice-2024-01-05.zarr\n", + "2026-10-30 12:24:08,224 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:08,494 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:08,618 INFO Transaction completed.\n", + "2026-10-30 12:24:08,619 INFO Slice dataset inputs/datatree/slice-2024-01-05.zarr closed\n", + "2026-10-30 12:24:08,624 INFO Opening slice dataset from inputs/datatree/slice-2024-01-06.zarr\n", + "2026-10-30 12:24:08,664 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:08,986 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:09,116 INFO Transaction completed.\n", + "2026-10-30 12:24:09,116 INFO Slice dataset inputs/datatree/slice-2024-01-06.zarr closed\n", + "2026-10-30 12:24:09,121 INFO Opening slice dataset from inputs/datatree/slice-2024-01-07.zarr\n", + "2026-10-30 12:24:09,167 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:09,432 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:09,552 INFO Transaction completed.\n", + "2026-10-30 12:24:09,553 INFO Slice dataset inputs/datatree/slice-2024-01-07.zarr closed\n", + "2026-10-30 12:24:09,558 INFO Opening slice dataset from inputs/datatree/slice-2024-01-08.zarr\n", + "2026-10-30 12:24:09,601 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:09,955 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:10,115 INFO Transaction completed.\n", + "2026-10-30 12:24:10,115 INFO Slice dataset inputs/datatree/slice-2024-01-08.zarr closed\n", + "2026-10-30 12:24:10,121 INFO Opening slice dataset from inputs/datatree/slice-2024-01-09.zarr\n", + "2026-10-30 12:24:10,160 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:10,433 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:10,561 INFO Transaction completed.\n", + "2026-10-30 12:24:10,561 INFO Slice dataset inputs/datatree/slice-2024-01-09.zarr closed\n", + "2026-10-30 12:24:10,567 INFO Opening slice dataset from inputs/datatree/slice-2024-01-10.zarr\n", + "2026-10-30 12:24:10,610 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:10,886 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:11,008 INFO Transaction completed.\n", + "2026-10-30 12:24:11,009 INFO Slice dataset inputs/datatree/slice-2024-01-10.zarr closed\n", + "2026-10-30 12:24:11,014 INFO Opening slice dataset from inputs/datatree/slice-2024-01-11.zarr\n", + "2026-10-30 12:24:11,060 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:11,464 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:11,581 INFO Transaction completed.\n", + "2026-10-30 12:24:11,582 INFO Slice dataset inputs/datatree/slice-2024-01-11.zarr closed\n", + "2026-10-30 12:24:11,587 INFO Opening slice dataset from inputs/datatree/slice-2024-01-12.zarr\n", + "2026-10-30 12:24:11,632 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:11,908 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:12,027 INFO Transaction completed.\n", + "2026-10-30 12:24:12,027 INFO Slice dataset inputs/datatree/slice-2024-01-12.zarr closed\n", + "2026-10-30 12:24:12,032 INFO Opening slice dataset from inputs/datatree/slice-2024-01-13.zarr\n", + "2026-10-30 12:24:12,077 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:12,379 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:12,502 INFO Transaction completed.\n", + "2026-10-30 12:24:12,503 INFO Slice dataset inputs/datatree/slice-2024-01-13.zarr closed\n", + "2026-10-30 12:24:12,508 INFO Opening slice dataset from inputs/datatree/slice-2024-01-14.zarr\n", + "2026-10-30 12:24:12,553 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:12,816 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:12,936 INFO Transaction completed.\n", + "2026-10-30 12:24:12,936 INFO Slice dataset inputs/datatree/slice-2024-01-14.zarr closed\n", + "2026-10-30 12:24:12,941 INFO Opening slice dataset from inputs/datatree/slice-2024-01-15.zarr\n", + "2026-10-30 12:24:12,984 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:13,232 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:13,383 INFO Transaction completed.\n", + "2026-10-30 12:24:13,383 INFO Slice dataset inputs/datatree/slice-2024-01-15.zarr closed\n", + "2026-10-30 12:24:13,389 INFO Opening slice dataset from inputs/datatree/slice-2024-01-16.zarr\n", + "2026-10-30 12:24:13,440 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:13,684 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:13,838 INFO Transaction completed.\n", + "2026-10-30 12:24:13,839 INFO Slice dataset inputs/datatree/slice-2024-01-16.zarr closed\n", + "2026-10-30 12:24:13,845 INFO Opening slice dataset from inputs/datatree/slice-2024-01-17.zarr\n", + "2026-10-30 12:24:13,898 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:14,158 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:14,311 INFO Transaction completed.\n", + "2026-10-30 12:24:14,312 INFO Slice dataset inputs/datatree/slice-2024-01-17.zarr closed\n", + "2026-10-30 12:24:14,321 INFO Opening slice dataset from inputs/datatree/slice-2024-01-18.zarr\n", + "2026-10-30 12:24:14,391 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:14,646 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:14,770 INFO Transaction completed.\n", + "2026-10-30 12:24:14,771 INFO Slice dataset inputs/datatree/slice-2024-01-18.zarr closed\n", + "2026-10-30 12:24:14,777 INFO Opening slice dataset from inputs/datatree/slice-2024-01-19.zarr\n", + "2026-10-30 12:24:14,841 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:15,123 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:15,237 INFO Transaction completed.\n", + "2026-10-30 12:24:15,237 INFO Slice dataset inputs/datatree/slice-2024-01-19.zarr closed\n", + "2026-10-30 12:24:15,242 INFO Opening slice dataset from inputs/datatree/slice-2024-01-20.zarr\n", + "2026-10-30 12:24:15,301 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:15,567 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:15,687 INFO Transaction completed.\n", + "2026-10-30 12:24:15,688 INFO Slice dataset inputs/datatree/slice-2024-01-20.zarr closed\n", + "2026-10-30 12:24:15,694 INFO Opening slice dataset from inputs/datatree/slice-2024-01-21.zarr\n", + "2026-10-30 12:24:15,745 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:16,038 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:16,157 INFO Transaction completed.\n", + "2026-10-30 12:24:16,158 INFO Slice dataset inputs/datatree/slice-2024-01-21.zarr closed\n", + "2026-10-30 12:24:16,163 INFO Opening slice dataset from inputs/datatree/slice-2024-01-22.zarr\n", + "2026-10-30 12:24:16,214 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:16,494 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:16,624 INFO Transaction completed.\n", + "2026-10-30 12:24:16,624 INFO Slice dataset inputs/datatree/slice-2024-01-22.zarr closed\n", + "2026-10-30 12:24:16,629 INFO Opening slice dataset from inputs/datatree/slice-2024-01-23.zarr\n", + "2026-10-30 12:24:16,682 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:16,971 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:17,092 INFO Transaction completed.\n", + "2026-10-30 12:24:17,093 INFO Slice dataset inputs/datatree/slice-2024-01-23.zarr closed\n", + "2026-10-30 12:24:17,098 INFO Opening slice dataset from inputs/datatree/slice-2024-01-24.zarr\n", + "2026-10-30 12:24:17,153 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:17,443 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:17,564 INFO Transaction completed.\n", + "2026-10-30 12:24:17,565 INFO Slice dataset inputs/datatree/slice-2024-01-24.zarr closed\n", + "2026-10-30 12:24:17,570 INFO Opening slice dataset from inputs/datatree/slice-2024-01-25.zarr\n", + "2026-10-30 12:24:17,623 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:17,913 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:18,036 INFO Transaction completed.\n", + "2026-10-30 12:24:18,036 INFO Slice dataset inputs/datatree/slice-2024-01-25.zarr closed\n", + "2026-10-30 12:24:18,041 INFO Opening slice dataset from inputs/datatree/slice-2024-01-26.zarr\n", + "2026-10-30 12:24:18,095 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:18,430 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:18,570 INFO Transaction completed.\n", + "2026-10-30 12:24:18,571 INFO Slice dataset inputs/datatree/slice-2024-01-26.zarr closed\n", + "2026-10-30 12:24:18,576 INFO Opening slice dataset from inputs/datatree/slice-2024-01-27.zarr\n", + "2026-10-30 12:24:18,623 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:18,958 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:19,092 INFO Transaction completed.\n", + "2026-10-30 12:24:19,093 INFO Slice dataset inputs/datatree/slice-2024-01-27.zarr closed\n", + "2026-10-30 12:24:19,099 INFO Opening slice dataset from inputs/datatree/slice-2024-01-28.zarr\n", + "2026-10-30 12:24:19,146 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:19,419 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:19,544 INFO Transaction completed.\n", + "2026-10-30 12:24:19,544 INFO Slice dataset inputs/datatree/slice-2024-01-28.zarr closed\n", + "2026-10-30 12:24:19,550 INFO Opening slice dataset from inputs/datatree/slice-2024-01-29.zarr\n", + "2026-10-30 12:24:19,601 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:19,873 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:20,003 INFO Transaction completed.\n", + "2026-10-30 12:24:20,004 INFO Slice dataset inputs/datatree/slice-2024-01-29.zarr closed\n", + "2026-10-30 12:24:20,010 INFO Opening slice dataset from inputs/datatree/slice-2024-01-30.zarr\n", + "2026-10-30 12:24:20,069 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:20,369 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:20,487 INFO Transaction completed.\n", + "2026-10-30 12:24:20,488 INFO Slice dataset inputs/datatree/slice-2024-01-30.zarr closed\n", + "2026-10-30 12:24:20,493 INFO Opening slice dataset from inputs/datatree/slice-2024-01-31.zarr\n", + "2026-10-30 12:24:20,552 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/0\n", + "2026-10-30 12:24:20,817 INFO Updating target dataset outputs/datatree_example_space_opt.zarr/1\n", + "2026-10-30 12:24:20,947 INFO Transaction completed.\n", + "2026-10-30 12:24:20,948 INFO Slice dataset inputs/datatree/slice-2024-01-31.zarr closed\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "CPU times: user 15.5 s, sys: 3.2 s, total: 18.7 s\n", + "Wall time: 14.9 s\n" + ] + } + ], + "source": [ + "%%time\n", + "target = output_dir / 'datatree_example_space_opt.zarr'\n", + "count = zappend(\n", + " slices,\n", + " target_dir=str(target),\n", + " force_new=True,\n", + " logging='INFO',\n", + ")\n", + "assert count == len(slices)" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "a2e0e79dfe38f97a", + "metadata": { + "ExecuteTime": { + "end_time": "2026-10-30T11:24:21.085093795Z", + "start_time": "2026-10-30T11:24:20.960792517Z" + } + }, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "
<xarray.DataTree>\n",
+       "Group: /\n",
+       "│   Attributes:\n",
+       "│       title:        Zappend nested example\n",
+       "│       description:  Two spatial resolutions of daily data\n",
+       "├── Group: /0\n",
+       "│       Dimensions:       (time: 31, lat: 1800, lon: 3600)\n",
+       "│       Coordinates:\n",
+       "│         * time          (time) datetime64[ns] 248B 2024-01-01 ... 2024-01-31\n",
+       "│         * lat           (lat) float32 7kB -90.0 -89.9 -89.8 -89.7 ... 89.7 89.8 89.9\n",
+       "│         * lon           (lon) float32 14kB -180.0 -179.9 -179.8 ... 179.7 179.8 179.9\n",
+       "│       Data variables:\n",
+       "│           quality_flag  (time, lat, lon) uint8 201MB dask.array<chunksize=(1, 450, 450), meta=np.ndarray>\n",
+       "│           temperature   (time, lat, lon) float32 804MB dask.array<chunksize=(1, 450, 450), meta=np.ndarray>\n",
+       "│       Attributes:\n",
+       "│           spatial_resolution:  0.1 degrees\n",
+       "└── Group: /1\n",
+       "        Dimensions:       (time: 31, lat: 900, lon: 1800)\n",
+       "        Coordinates:\n",
+       "          * time          (time) datetime64[ns] 248B 2024-01-01 ... 2024-01-31\n",
+       "          * lat           (lat) float32 4kB -90.0 -89.8 -89.6 -89.4 ... 89.4 89.6 89.8\n",
+       "          * lon           (lon) float32 7kB -180.0 -179.8 -179.6 ... 179.4 179.6 179.8\n",
+       "        Data variables:\n",
+       "            quality_flag  (time, lat, lon) uint8 50MB dask.array<chunksize=(1, 450, 450), meta=np.ndarray>\n",
+       "            temperature   (time, lat, lon) float32 201MB dask.array<chunksize=(1, 450, 450), meta=np.ndarray>\n",
+       "        Attributes:\n",
+       "            spatial_resolution:  0.2 degrees
" + ], + "text/plain": [ + "\n", + "Group: /\n", + "│ Attributes:\n", + "│ title: Zappend nested example\n", + "│ description: Two spatial resolutions of daily data\n", + "├── Group: /0\n", + "│ Dimensions: (time: 31, lat: 1800, lon: 3600)\n", + "│ Coordinates:\n", + "│ * time (time) datetime64[ns] 248B 2024-01-01 ... 2024-01-31\n", + "│ * lat (lat) float32 7kB -90.0 -89.9 -89.8 -89.7 ... 89.7 89.8 89.9\n", + "│ * lon (lon) float32 14kB -180.0 -179.9 -179.8 ... 179.7 179.8 179.9\n", + "│ Data variables:\n", + "│ quality_flag (time, lat, lon) uint8 201MB dask.array\n", + "│ temperature (time, lat, lon) float32 804MB dask.array\n", + "│ Attributes:\n", + "│ spatial_resolution: 0.1 degrees\n", + "└── Group: /1\n", + " Dimensions: (time: 31, lat: 900, lon: 1800)\n", + " Coordinates:\n", + " * time (time) datetime64[ns] 248B 2024-01-01 ... 2024-01-31\n", + " * lat (lat) float32 4kB -90.0 -89.8 -89.6 -89.4 ... 89.4 89.6 89.8\n", + " * lon (lon) float32 7kB -180.0 -179.8 -179.6 ... 179.4 179.6 179.8\n", + " Data variables:\n", + " quality_flag (time, lat, lon) uint8 50MB dask.array\n", + " temperature (time, lat, lon) float32 201MB dask.array\n", + " Attributes:\n", + " spatial_resolution: 0.2 degrees" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "tree = xr.open_datatree(target, engine='zarr', chunks={})\n", + "tree" + ] + }, + { + "cell_type": "markdown", + "id": "a1f311789cccc583", + "metadata": {}, + "source": [ + "## Per-group time-series chunks\n", + "\n", + "`group_variables` supplies optional overrides keyed by absolute group path. Here both resolutions get a full-month time chunk, while their spatial chunk shapes remain appropriate for their own grids. Any group omitted from `group_variables` would retain the shared `variables` configuration, or its first slice's encoding when no shared configuration is given." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "45465b532a72ebe1", + "metadata": { + "ExecuteTime": { + "end_time": "2026-10-30T11:29:11.180228870Z", + "start_time": "2026-10-30T11:26:25.776613956Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2026-10-30 12:26:25,846 WARNING Setting 'force_new' is enabled. This will permanently delete existing targets (no rollback).\n", + "2026-10-30 12:26:25,847 WARNING Permanently deleting outputs/datatree_example_time_opt.zarr\n", + "2026-10-30 12:26:25,987 INFO Opening slice dataset from inputs/datatree/slice-2024-01-01.zarr\n", + "2026-10-30 12:26:26,048 INFO Creating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:26:30,854 INFO Creating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:26:32,105 INFO Transaction completed.\n", + "2026-10-30 12:26:32,106 INFO Slice dataset inputs/datatree/slice-2024-01-01.zarr closed\n", + "2026-10-30 12:26:32,113 INFO Opening slice dataset from inputs/datatree/slice-2024-01-02.zarr\n", + "2026-10-30 12:26:32,166 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:26:39,677 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:26:41,479 INFO Transaction completed.\n", + "2026-10-30 12:26:41,480 INFO Slice dataset inputs/datatree/slice-2024-01-02.zarr closed\n", + "2026-10-30 12:26:41,486 INFO Opening slice dataset from inputs/datatree/slice-2024-01-03.zarr\n", + "2026-10-30 12:26:41,573 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:26:54,757 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:26:56,964 INFO Transaction completed.\n", + "2026-10-30 12:26:56,966 INFO Slice dataset inputs/datatree/slice-2024-01-03.zarr closed\n", + "2026-10-30 12:26:56,972 INFO Opening slice dataset from inputs/datatree/slice-2024-01-04.zarr\n", + "2026-10-30 12:26:57,033 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:27:06,864 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:27:09,171 INFO Transaction completed.\n", + "2026-10-30 12:27:09,172 INFO Slice dataset inputs/datatree/slice-2024-01-04.zarr closed\n", + "2026-10-30 12:27:09,179 INFO Opening slice dataset from inputs/datatree/slice-2024-01-05.zarr\n", + "2026-10-30 12:27:09,237 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:27:13,444 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:27:14,395 INFO Transaction completed.\n", + "2026-10-30 12:27:14,395 INFO Slice dataset inputs/datatree/slice-2024-01-05.zarr closed\n", + "2026-10-30 12:27:14,399 INFO Opening slice dataset from inputs/datatree/slice-2024-01-06.zarr\n", + "2026-10-30 12:27:14,425 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:27:18,435 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:27:19,335 INFO Transaction completed.\n", + "2026-10-30 12:27:19,336 INFO Slice dataset inputs/datatree/slice-2024-01-06.zarr closed\n", + "2026-10-30 12:27:19,339 INFO Opening slice dataset from inputs/datatree/slice-2024-01-07.zarr\n", + "2026-10-30 12:27:19,364 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:27:22,935 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:27:24,011 INFO Transaction completed.\n", + "2026-10-30 12:27:24,012 INFO Slice dataset inputs/datatree/slice-2024-01-07.zarr closed\n", + "2026-10-30 12:27:24,015 INFO Opening slice dataset from inputs/datatree/slice-2024-01-08.zarr\n", + "2026-10-30 12:27:24,040 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:27:27,320 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:27:28,208 INFO Transaction completed.\n", + "2026-10-30 12:27:28,209 INFO Slice dataset inputs/datatree/slice-2024-01-08.zarr closed\n", + "2026-10-30 12:27:28,212 INFO Opening slice dataset from inputs/datatree/slice-2024-01-09.zarr\n", + "2026-10-30 12:27:28,238 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:27:31,579 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:27:32,966 INFO Transaction completed.\n", + "2026-10-30 12:27:32,967 INFO Slice dataset inputs/datatree/slice-2024-01-09.zarr closed\n", + "2026-10-30 12:27:32,973 INFO Opening slice dataset from inputs/datatree/slice-2024-01-10.zarr\n", + "2026-10-30 12:27:33,012 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:27:36,595 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:27:37,471 INFO Transaction completed.\n", + "2026-10-30 12:27:37,471 INFO Slice dataset inputs/datatree/slice-2024-01-10.zarr closed\n", + "2026-10-30 12:27:37,475 INFO Opening slice dataset from inputs/datatree/slice-2024-01-11.zarr\n", + "2026-10-30 12:27:37,499 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:27:40,771 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:27:41,636 INFO Transaction completed.\n", + "2026-10-30 12:27:41,637 INFO Slice dataset inputs/datatree/slice-2024-01-11.zarr closed\n", + "2026-10-30 12:27:41,640 INFO Opening slice dataset from inputs/datatree/slice-2024-01-12.zarr\n", + "2026-10-30 12:27:41,684 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:27:44,938 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:27:45,811 INFO Transaction completed.\n", + "2026-10-30 12:27:45,811 INFO Slice dataset inputs/datatree/slice-2024-01-12.zarr closed\n", + "2026-10-30 12:27:45,815 INFO Opening slice dataset from inputs/datatree/slice-2024-01-13.zarr\n", + "2026-10-30 12:27:45,840 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:27:49,907 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:27:50,786 INFO Transaction completed.\n", + "2026-10-30 12:27:50,787 INFO Slice dataset inputs/datatree/slice-2024-01-13.zarr closed\n", + "2026-10-30 12:27:50,790 INFO Opening slice dataset from inputs/datatree/slice-2024-01-14.zarr\n", + "2026-10-30 12:27:50,816 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:27:54,148 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:27:55,034 INFO Transaction completed.\n", + "2026-10-30 12:27:55,035 INFO Slice dataset inputs/datatree/slice-2024-01-14.zarr closed\n", + "2026-10-30 12:27:55,038 INFO Opening slice dataset from inputs/datatree/slice-2024-01-15.zarr\n", + "2026-10-30 12:27:55,064 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:27:58,475 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:27:59,355 INFO Transaction completed.\n", + "2026-10-30 12:27:59,356 INFO Slice dataset inputs/datatree/slice-2024-01-15.zarr closed\n", + "2026-10-30 12:27:59,359 INFO Opening slice dataset from inputs/datatree/slice-2024-01-16.zarr\n", + "2026-10-30 12:27:59,386 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:28:02,672 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:28:03,559 INFO Transaction completed.\n", + "2026-10-30 12:28:03,560 INFO Slice dataset inputs/datatree/slice-2024-01-16.zarr closed\n", + "2026-10-30 12:28:03,564 INFO Opening slice dataset from inputs/datatree/slice-2024-01-17.zarr\n", + "2026-10-30 12:28:03,591 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:28:07,436 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:28:08,313 INFO Transaction completed.\n", + "2026-10-30 12:28:08,313 INFO Slice dataset inputs/datatree/slice-2024-01-17.zarr closed\n", + "2026-10-30 12:28:08,317 INFO Opening slice dataset from inputs/datatree/slice-2024-01-18.zarr\n", + "2026-10-30 12:28:08,342 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:28:11,763 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:28:12,626 INFO Transaction completed.\n", + "2026-10-30 12:28:12,627 INFO Slice dataset inputs/datatree/slice-2024-01-18.zarr closed\n", + "2026-10-30 12:28:12,630 INFO Opening slice dataset from inputs/datatree/slice-2024-01-19.zarr\n", + "2026-10-30 12:28:12,657 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:28:15,967 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:28:16,874 INFO Transaction completed.\n", + "2026-10-30 12:28:16,875 INFO Slice dataset inputs/datatree/slice-2024-01-19.zarr closed\n", + "2026-10-30 12:28:16,878 INFO Opening slice dataset from inputs/datatree/slice-2024-01-20.zarr\n", + "2026-10-30 12:28:16,924 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:28:20,667 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:28:21,677 INFO Transaction completed.\n", + "2026-10-30 12:28:21,678 INFO Slice dataset inputs/datatree/slice-2024-01-20.zarr closed\n", + "2026-10-30 12:28:21,681 INFO Opening slice dataset from inputs/datatree/slice-2024-01-21.zarr\n", + "2026-10-30 12:28:21,711 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:28:25,100 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:28:26,005 INFO Transaction completed.\n", + "2026-10-30 12:28:26,005 INFO Slice dataset inputs/datatree/slice-2024-01-21.zarr closed\n", + "2026-10-30 12:28:26,009 INFO Opening slice dataset from inputs/datatree/slice-2024-01-22.zarr\n", + "2026-10-30 12:28:26,040 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:28:29,282 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:28:30,170 INFO Transaction completed.\n", + "2026-10-30 12:28:30,171 INFO Slice dataset inputs/datatree/slice-2024-01-22.zarr closed\n", + "2026-10-30 12:28:30,174 INFO Opening slice dataset from inputs/datatree/slice-2024-01-23.zarr\n", + "2026-10-30 12:28:30,207 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:28:33,636 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:28:34,536 INFO Transaction completed.\n", + "2026-10-30 12:28:34,536 INFO Slice dataset inputs/datatree/slice-2024-01-23.zarr closed\n", + "2026-10-30 12:28:34,540 INFO Opening slice dataset from inputs/datatree/slice-2024-01-24.zarr\n", + "2026-10-30 12:28:34,586 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:28:38,587 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:28:39,457 INFO Transaction completed.\n", + "2026-10-30 12:28:39,458 INFO Slice dataset inputs/datatree/slice-2024-01-24.zarr closed\n", + "2026-10-30 12:28:39,461 INFO Opening slice dataset from inputs/datatree/slice-2024-01-25.zarr\n", + "2026-10-30 12:28:39,507 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:28:42,815 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:28:43,697 INFO Transaction completed.\n", + "2026-10-30 12:28:43,697 INFO Slice dataset inputs/datatree/slice-2024-01-25.zarr closed\n", + "2026-10-30 12:28:43,700 INFO Opening slice dataset from inputs/datatree/slice-2024-01-26.zarr\n", + "2026-10-30 12:28:43,730 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:28:47,101 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:28:47,958 INFO Transaction completed.\n", + "2026-10-30 12:28:47,958 INFO Slice dataset inputs/datatree/slice-2024-01-26.zarr closed\n", + "2026-10-30 12:28:47,961 INFO Opening slice dataset from inputs/datatree/slice-2024-01-27.zarr\n", + "2026-10-30 12:28:48,008 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:28:51,320 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:28:52,600 INFO Transaction completed.\n", + "2026-10-30 12:28:52,601 INFO Slice dataset inputs/datatree/slice-2024-01-27.zarr closed\n", + "2026-10-30 12:28:52,606 INFO Opening slice dataset from inputs/datatree/slice-2024-01-28.zarr\n", + "2026-10-30 12:28:52,656 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:28:56,509 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:28:57,408 INFO Transaction completed.\n", + "2026-10-30 12:28:57,409 INFO Slice dataset inputs/datatree/slice-2024-01-28.zarr closed\n", + "2026-10-30 12:28:57,412 INFO Opening slice dataset from inputs/datatree/slice-2024-01-29.zarr\n", + "2026-10-30 12:28:57,451 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:29:00,870 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:29:01,735 INFO Transaction completed.\n", + "2026-10-30 12:29:01,735 INFO Slice dataset inputs/datatree/slice-2024-01-29.zarr closed\n", + "2026-10-30 12:29:01,738 INFO Opening slice dataset from inputs/datatree/slice-2024-01-30.zarr\n", + "2026-10-30 12:29:01,770 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:29:05,068 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:29:05,950 INFO Transaction completed.\n", + "2026-10-30 12:29:05,951 INFO Slice dataset inputs/datatree/slice-2024-01-30.zarr closed\n", + "2026-10-30 12:29:05,955 INFO Opening slice dataset from inputs/datatree/slice-2024-01-31.zarr\n", + "2026-10-30 12:29:05,988 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/0\n", + "2026-10-30 12:29:10,276 INFO Updating target dataset outputs/datatree_example_time_opt.zarr/1\n", + "2026-10-30 12:29:11,166 INFO Transaction completed.\n", + "2026-10-30 12:29:11,167 INFO Slice dataset inputs/datatree/slice-2024-01-31.zarr closed\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "CPU times: user 3min 26s, sys: 49.4 s, total: 4min 16s\n", + "Wall time: 2min 45s\n" + ] + } + ], + "source": [ + "%%time\n", + "full_time_chunk = len(slices)\n", + "target = output_dir / 'datatree_example_time_opt.zarr'\n", + "\n", + "def variable_chunks(chunks):\n", + " return {\n", + " name: {'encoding': {'chunks': chunks}}\n", + " for name in ('temperature', 'quality_flag')\n", + " }\n", + "\n", + "count = zappend(\n", + " slices,\n", + " target_dir=str(target),\n", + " force_new=True,\n", + " group_variables={\n", + " '/0': variable_chunks([full_time_chunk, 90, 90]),\n", + " '/1': variable_chunks([full_time_chunk, 90, 90]),\n", + " },\n", + ")\n", + "assert count == full_time_chunk" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "361423d39986a366", + "metadata": { + "ExecuteTime": { + "end_time": "2026-10-30T11:29:11.295519683Z", + "start_time": "2026-10-30T11:29:11.181605322Z" + } + }, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "
<xarray.DataTree>\n",
+       "Group: /\n",
+       "│   Attributes:\n",
+       "│       title:        Zappend nested example\n",
+       "│       description:  Two spatial resolutions of daily data\n",
+       "├── Group: /0\n",
+       "│       Dimensions:       (time: 31, lat: 1800, lon: 3600)\n",
+       "│       Coordinates:\n",
+       "│         * time          (time) datetime64[ns] 248B 2024-01-01 ... 2024-01-31\n",
+       "│         * lat           (lat) float32 7kB -90.0 -89.9 -89.8 -89.7 ... 89.7 89.8 89.9\n",
+       "│         * lon           (lon) float32 14kB -180.0 -179.9 -179.8 ... 179.7 179.8 179.9\n",
+       "│       Data variables:\n",
+       "│           quality_flag  (time, lat, lon) uint8 201MB dask.array<chunksize=(31, 90, 90), meta=np.ndarray>\n",
+       "│           temperature   (time, lat, lon) float32 804MB dask.array<chunksize=(31, 90, 90), meta=np.ndarray>\n",
+       "│       Attributes:\n",
+       "│           spatial_resolution:  0.1 degrees\n",
+       "└── Group: /1\n",
+       "        Dimensions:       (time: 31, lat: 900, lon: 1800)\n",
+       "        Coordinates:\n",
+       "          * time          (time) datetime64[ns] 248B 2024-01-01 ... 2024-01-31\n",
+       "          * lat           (lat) float32 4kB -90.0 -89.8 -89.6 -89.4 ... 89.4 89.6 89.8\n",
+       "          * lon           (lon) float32 7kB -180.0 -179.8 -179.6 ... 179.4 179.6 179.8\n",
+       "        Data variables:\n",
+       "            quality_flag  (time, lat, lon) uint8 50MB dask.array<chunksize=(31, 90, 90), meta=np.ndarray>\n",
+       "            temperature   (time, lat, lon) float32 201MB dask.array<chunksize=(31, 90, 90), meta=np.ndarray>\n",
+       "        Attributes:\n",
+       "            spatial_resolution:  0.2 degrees
" + ], + "text/plain": [ + "\n", + "Group: /\n", + "│ Attributes:\n", + "│ title: Zappend nested example\n", + "│ description: Two spatial resolutions of daily data\n", + "├── Group: /0\n", + "│ Dimensions: (time: 31, lat: 1800, lon: 3600)\n", + "│ Coordinates:\n", + "│ * time (time) datetime64[ns] 248B 2024-01-01 ... 2024-01-31\n", + "│ * lat (lat) float32 7kB -90.0 -89.9 -89.8 -89.7 ... 89.7 89.8 89.9\n", + "│ * lon (lon) float32 14kB -180.0 -179.9 -179.8 ... 179.7 179.8 179.9\n", + "│ Data variables:\n", + "│ quality_flag (time, lat, lon) uint8 201MB dask.array\n", + "│ temperature (time, lat, lon) float32 804MB dask.array\n", + "│ Attributes:\n", + "│ spatial_resolution: 0.1 degrees\n", + "└── Group: /1\n", + " Dimensions: (time: 31, lat: 900, lon: 1800)\n", + " Coordinates:\n", + " * time (time) datetime64[ns] 248B 2024-01-01 ... 2024-01-31\n", + " * lat (lat) float32 4kB -90.0 -89.8 -89.6 -89.4 ... 89.4 89.6 89.8\n", + " * lon (lon) float32 7kB -180.0 -179.8 -179.6 ... 179.4 179.6 179.8\n", + " Data variables:\n", + " quality_flag (time, lat, lon) uint8 50MB dask.array\n", + " temperature (time, lat, lon) float32 201MB dask.array\n", + " Attributes:\n", + " spatial_resolution: 0.2 degrees" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "tree = xr.open_datatree(target, engine='zarr', chunks={})\n", + "tree" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "33fc36dead6e191c", + "metadata": { + "ExecuteTime": { + "end_time": "2026-10-30T11:29:11.425926988Z", + "start_time": "2026-10-30T11:29:11.299494739Z" + } + }, + "outputs": [], + "source": [ + "tree.close()\n", + "if os.path.exists(input_dir):\n", + " shutil.rmtree(input_dir)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python (zappend)", + "language": "python", + "name": "zappend" + }, + "language_info": { + "name": "python", + "version": "3.15" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tests/config/test_schema.py b/tests/config/test_schema.py index 33a35d3..52bd9a5 100644 --- a/tests/config/test_schema.py +++ b/tests/config/test_schema.py @@ -24,6 +24,7 @@ def test_get_config_schema(self): "extra", "force_new", "fixed_dims", + "group_variables", "included_variables", "logging", "permit_eval", diff --git a/tests/config/test_validate.py b/tests/config/test_validate.py index 1e3ce46..0e1673b 100644 --- a/tests/config/test_validate.py +++ b/tests/config/test_validate.py @@ -48,6 +48,15 @@ def test_validate_shards(self): } self.assertIs(config, validate_config(config)) + def test_validate_group_variables(self): + config = { + "zarr_format": 2, + "group_variables": { + "/0": {"v": {"encoding": {"chunks": [2, 4]}}}, + } + } + self.assertIs(config, validate_config(config)) + # noinspection PyMethodMayBeStatic def test_validate_variable_fail(self): config = { diff --git a/tests/contrib/test_datatree.py b/tests/contrib/test_datatree.py new file mode 100644 index 0000000..28a353e --- /dev/null +++ b/tests/contrib/test_datatree.py @@ -0,0 +1,236 @@ +# Copyright © 2024-2026 Brockmann Consult and contributors +# Permissions are hereby granted under the terms of the MIT License: +# https://opensource.org/licenses/MIT. + +import pytest +import xarray as xr + +from tests.helpers import clear_memory_fs, make_test_dataset +from zappend.api import zappend as _zappend +from zappend.context import Context +from zappend.fsutil import FileObj + +TARGET = "memory://datatree-target.zarr" + + +def zappend(*args, **kwargs): + kwargs.setdefault("zarr_format", 2) + return _zappend(*args, **kwargs) + + +def make_test_datatree( + time_index: int, + *, + include_group_1: bool = True, + group_1_shape: tuple[int, int] = (3, 2), +) -> xr.DataTree: + groups = { + "/": xr.Dataset(attrs={"title": "nested test data", "source": "test"}), + "/0": make_test_dataset( + shape=(1, 2, 3), chunks=(1, 2, 3), index=time_index + ), + } + if include_group_1: + groups["/1"] = make_test_dataset( + shape=(1, *group_1_shape), chunks=(1, *group_1_shape), index=time_index + ) + return xr.DataTree.from_dict(groups) + + +@pytest.fixture(autouse=True) +def clear_memory(): + clear_memory_fs() + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_zappend_appends_datatrees_to_nested_zarr_groups(zarr_format): + trees = [make_test_datatree(0), make_test_datatree(1)] + + assert zappend(trees, target_dir=TARGET, zarr_format=zarr_format) == 2 + + for group in ("0", "1"): + with xr.open_zarr( + TARGET, + group=group, + zarr_format=zarr_format, + consolidated=zarr_format == 2, + ) as dataset: + assert dataset.sizes["time"] == 2 + + with xr.open_zarr( + TARGET, zarr_format=zarr_format, consolidated=zarr_format == 2 + ) as root: + assert root.attrs == {"title": "nested test data", "source": "test"} + + +def test_zappend_rejects_datatrees_with_different_group_paths(): + with pytest.raises(ValueError, match="different group paths"): + zappend( + [make_test_datatree(0), make_test_datatree(1, include_group_1=False)], + target_dir=TARGET, + ) + + + +def test_zappend_rejects_unknown_group_variable_configuration(): + with pytest.raises(ValueError, match="group_variables contains paths"): + zappend( + [make_test_datatree(0)], + target_dir=TARGET, + group_variables={"/unknown": {}}, + ) + + +def test_zappend_uses_optional_group_variable_chunk_configuration(): + group_0_variables = { + variable: {"encoding": {"chunks": [2, 2, 3]}} + for variable in ("chl", "tsm") + } + zappend( + [make_test_datatree(0), make_test_datatree(1)], + target_dir=TARGET, + group_variables={"/0": group_0_variables}, + ) + + with xr.open_zarr(TARGET, group="0", zarr_format=2) as dataset: + assert dataset.chl.encoding["chunks"] == (2, 2, 3) + with xr.open_zarr(TARGET, group="1", zarr_format=2) as dataset: + # No group configuration: retain the first slice's existing chunking. + assert dataset.chl.encoding["chunks"] == (1, 3, 2) + + +@pytest.mark.parametrize("attrs_update_mode", ["replace", "update", "ignore"]) +def test_zappend_applies_root_attribute_update_modes(attrs_update_mode): + zappend( + [make_test_datatree(0)], + target_dir=TARGET, + attrs_update_mode=attrs_update_mode, + attrs={"configured": attrs_update_mode}, + ) + + with xr.open_zarr(TARGET, zarr_format=2) as root: + assert root.attrs["configured"] == attrs_update_mode + if attrs_update_mode == "ignore": + assert "title" not in root.attrs + else: + assert root.attrs["title"] == "nested test data" + + +def test_context_rejects_a_non_appendable_dataset_target(): + xr.Dataset({"not_time": ("x", [1])}).to_zarr(TARGET, zarr_format=2) + + with pytest.raises(ValueError, match="Append dimension 'time' not found"): + Context({"target_dir": TARGET, "zarr_format": 2}) + + +def test_zappend_rejects_an_invalid_slice_source_result(): + class InvalidContextManager: + def __enter__(self): + return object() + + def __exit__(self, *args): + pass + + with pytest.raises(TypeError, match="slice source must yield"): + zappend( + [None], target_dir=TARGET, slice_source=lambda _: InvalidContextManager() + ) + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_zappend_rolls_back_the_complete_failing_datatree_slice(zarr_format): + zappend([make_test_datatree(0)], target_dir=TARGET, zarr_format=zarr_format) + + with pytest.raises(ValueError, match="Wrong size for dimension"): + zappend( + [make_test_datatree(1, group_1_shape=(4, 2))], + target_dir=TARGET, + zarr_format=zarr_format, + ) + + # Group /0 was updated before incompatible group /1 was encountered. The + # root transaction must nevertheless restore both groups to time index 0. + for group in ("0", "1"): + with xr.open_zarr( + TARGET, + group=group, + zarr_format=zarr_format, + consolidated=zarr_format == 2, + ) as dataset: + assert dataset.sizes["time"] == 1 + assert dataset.time.values[0] == make_test_datatree(0)[group].ds.time.values[0] + + +def test_zappend_removes_a_new_tree_when_a_later_group_fails(): + tree = make_test_datatree(0) + tree["/1"].ds = tree["/1"].ds.isel(time=0, drop=True) + + with pytest.raises(ValueError, match="Append dimension 'time' not found"): + zappend([tree], target_dir=TARGET) + + assert not FileObj(TARGET).exists() + + +def test_zappend_accepts_datatrees_from_a_slice_source(): + def make_slice(index: int) -> xr.DataTree: + return make_test_datatree(index) + + assert zappend( + [0, 1], + target_dir=TARGET, + slice_source=make_slice, + persist_mem_slices=True, + zarr_format=3, + ) == 2 + + for group in ("0", "1"): + with xr.open_zarr(TARGET, group=group, zarr_format=3) as dataset: + assert dataset.sizes["time"] == 2 + + +def test_zappend_opens_a_nested_zarr_slice_as_a_datatree(): + source = "memory://nested-source.zarr" + make_test_datatree(0).to_zarr(source, zarr_format=2) + + assert zappend([source], target_dir=TARGET) == 1 + + for group in ("0", "1"): + with xr.open_zarr(TARGET, group=group, zarr_format=2) as dataset: + assert dataset.sizes["time"] == 1 + + +@pytest.mark.parametrize( + "slices", + [ + lambda: [make_test_datatree(0), make_test_dataset(shape=(1, 2, 3))], + lambda: [make_test_dataset(shape=(1, 2, 3)), make_test_datatree(0)], + ], +) +def test_zappend_rejects_mixed_dataset_and_datatree_slices(slices): + with pytest.raises(TypeError, match="Cannot mix xarray.Dataset and xarray.DataTree"): + zappend(slices(), target_dir=TARGET) + + +def test_zappend_also_supports_a_data_bearing_root(): + trees = [ + xr.DataTree(make_test_dataset(shape=(1, 2, 3), index=index)) + for index in (0, 1) + ] + + zappend(trees, target_dir=TARGET) + + with xr.open_zarr(TARGET, zarr_format=2) as dataset: + assert dataset.sizes["time"] == 2 + + +def test_zappend_force_new_replaces_complete_tree(): + zappend( + [make_test_datatree(0), make_test_datatree(1)], target_dir=TARGET + ) + + zappend([make_test_datatree(3)], target_dir=TARGET, force_new=True) + + for group in ("0", "1"): + with xr.open_zarr(TARGET, group=group, zarr_format=2) as dataset: + assert dataset.sizes["time"] == 1 + assert dataset.time.values[0] == make_test_datatree(3)[group].ds.time.values[0] diff --git a/zappend/api.py b/zappend/api.py index e29a798..4d013f1 100644 --- a/zappend/api.py +++ b/zappend/api.py @@ -49,7 +49,7 @@ def zappend( is configured or not (setting `slice_source`). If no slice source is configured, a slice item must be an object of type - `str`, `FileObj`, `xarray.Dataset`, or `SliceSource`. + `str`, `FileObj`, `xarray.Dataset`, `xarray.DataTree`, or `SliceSource`. If `str` or `FileObj` are used, they are interpreted as local dataset path or dataset URI. If a URI is used, protocol-specific parameters apply, given by the configuration parameter `slice_storage_options`. @@ -57,6 +57,12 @@ def zappend( If a slice source is configured, a slice item represents the argument(s) passed to that slice source. Multiple positional arguments can be passed as `list`, multiple keyword arguments as `dict`, and both as a `tuple` of `list` and `dict`. + A slice source may return an `xarray.DataTree` for a nested Zarr store. + zappend appends its data-bearing nodes to their corresponding Zarr groups. + A DataTree slice is transactional as a whole: if appending any group fails, + all changes made while appending that DataTree slice are rolled back. + The first opened slice establishes whether the run uses Dataset or DataTree + slices; the two container types cannot be mixed in one run. Args: slices: An iterable that yields slice items. diff --git a/zappend/config/config.py b/zappend/config/config.py index 5c05aaf..5f6e242 100644 --- a/zappend/config/config.py +++ b/zappend/config/config.py @@ -114,6 +114,11 @@ def variables(self) -> dict[str, Any]: """Variable definitions.""" return self._config.get("variables") or {} + @property + def group_variables(self) -> dict[str, dict[str, Any]]: + """Variable-definition overrides keyed by absolute DataTree group path.""" + return self._config.get("group_variables") or {} + @property def attrs(self) -> dict[str, Any]: """Global dataset attributes. May include dynamically computed @@ -156,10 +161,11 @@ def slice_engine(self) -> str | None: @property def slice_source(self) -> Callable[[...], Any] | None: """A class or function that receives a - slice item as argument(s) and provides the slice dataset. + 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`. + * 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. diff --git a/zappend/config/schema.py b/zappend/config/schema.py index 114d355..450dc01 100644 --- a/zappend/config/schema.py +++ b/zappend/config/schema.py @@ -214,6 +214,16 @@ }, } +GROUP_VARIABLES_SCHEMA = { + "category": "Target Outline", + "description": ( + "Defines variable configuration overrides for individual absolute Zarr " + "group paths in a DataTree target." + ), + "type": "object", + "additionalProperties": VARIABLES_SCHEMA, +} + LOG_REF_URL = ( "https://docs.python.org/3/library/logging.config.html#logging-config-dictschema" ) @@ -558,6 +568,7 @@ "items": {"type": "string", "minLength": 1}, }, "variables": VARIABLES_SCHEMA, + "group_variables": GROUP_VARIABLES_SCHEMA, "attrs": { "category": "Target Outline", "description": ( @@ -678,11 +689,11 @@ "category": "Data I/O - Slices", "description": ( "The fully qualified name of a class or function that receives a" - " slice item as argument(s) and provides the slice dataset." + " 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`." + " 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." diff --git a/zappend/context.py b/zappend/context.py index 3420410..b4c2377 100644 --- a/zappend/context.py +++ b/zappend/context.py @@ -34,7 +34,18 @@ def __init__(self, config: dict[str, Any] | Config): append_var = target_dataset.get(_config.append_dim) if append_var is not None and append_var.size > 0: last_append_label = append_var[-1] - target_metadata = DatasetMetadata.from_dataset(target_dataset, _config) + try: + target_metadata = DatasetMetadata.from_dataset( + target_dataset, _config + ) + except ValueError: + # A DataTree may have a metadata-only root node and its + # appendable datasets in child Zarr groups. Such a root is + # not a Dataset target itself, so its metadata is not used + # by the Dataset processing path. + if target_dataset.data_vars or _config.append_dim in target_dataset.sizes: + raise + target_metadata = None except FileNotFoundError: target_metadata = None diff --git a/zappend/processor.py b/zappend/processor.py index 67b01bb..3cc90b3 100644 --- a/zappend/processor.py +++ b/zappend/processor.py @@ -17,6 +17,7 @@ exclude_from_config, get_dyn_config_attrs_env, has_dyn_config_attrs, + merge_configs, normalize_config, validate_config, ) @@ -56,6 +57,11 @@ def __init__(self, config: ConfigLike = None, **kwargs): ) delete_target_permanently(_config) self._config = _config + # Keep the normalized configuration so DataTree child groups can use + # the same settings with only their target path changed. + self._config_dict = dict(config) + self._slice_container_type: type[xr.Dataset | xr.DataTree] | None = None + self._datatree_group_paths: frozenset[str] | None = None self._profiler = Profiler(_config.profiling) def process_slices(self, slices: Iterable[SliceItem]) -> int: @@ -100,27 +106,215 @@ def process_slice(self, slice_item: SliceItem, slice_index: int = 0): ctx = Context(self._config) - with open_slice_dataset(ctx, slice_item, slice_index) as slice_dataset: - slice_metadata = ctx.get_dataset_metadata(slice_dataset) - if ctx.target_metadata is None: - ctx.target_metadata = slice_metadata + with open_slice_dataset(ctx, slice_item, slice_index) as slice_data: + if isinstance(slice_data, xr.Dataset): + self._assert_slice_container_type(xr.Dataset) + self.process_dataset_slice(ctx, slice_data) + elif isinstance(slice_data, xr.DataTree): + self._assert_slice_container_type(xr.DataTree) + self.process_datatree_slice(ctx, slice_data) else: - ctx.target_metadata.assert_compatible_slice( - slice_metadata, ctx.config.append_dim + raise TypeError( + "slice source must yield an xarray.Dataset or xarray.DataTree, " + f"but yielded {type(slice_data).__name__!r}" ) - verify_append_labels(ctx, slice_dataset) + def _assert_slice_container_type( + self, container_type: type[xr.Dataset | xr.DataTree] + ): + """Ensure a streaming run does not mix Dataset and DataTree slices.""" + if self._slice_container_type is None: + self._slice_container_type = container_type + elif container_type is not self._slice_container_type: + raise TypeError( + "Cannot mix xarray.Dataset and xarray.DataTree slices in one " + f"zappend run; expected {self._slice_container_type.__name__}, " + f"got {container_type.__name__}" + ) + + def process_datatree_slice(self, ctx: Context, tree: xr.DataTree): + """Append one DataTree slice in a transaction spanning all groups. + + This is deliberately part of :class:`Processor`: DataTrees can now be + returned by any configured slice source and use the same polling, + profiling, force-new, and configuration flow as Dataset slices. + Each child group's rollback actions are registered relative to the + DataTree root, so an error restores the complete tree to its state + before this slice began. + """ + group_paths = frozenset(tree.groups) + unknown_group_paths = set(ctx.config.group_variables) - group_paths + if unknown_group_paths: + raise ValueError( + "group_variables contains paths that are not present in the " + f"DataTree: {sorted(unknown_group_paths)!r}" + ) + if self._datatree_group_paths is None: + self._datatree_group_paths = group_paths + elif group_paths != self._datatree_group_paths: + raise ValueError( + "DataTree has different group paths; expected " + f"{sorted(self._datatree_group_paths)!r}, got {sorted(group_paths)!r}" + ) + + transaction = Transaction( + ctx.config.target_dir, + ctx.config.temp_dir, + disable_rollback=ctx.config.disable_rollback, + ) + with transaction as rollback_callback: + self._process_datatree_slice(ctx, tree, rollback_callback) + + def _process_datatree_slice( + self, ctx: Context, tree: xr.DataTree, rollback_callback: RollbackCallback + ): + """Process a DataTree while recording all paths relative to its root.""" + + data_paths = tuple( + sorted( + (node.path for node in tree.subtree if node.ds.data_vars), + key=lambda path: (path.count("/"), path), + ) + ) + root_has_data = "/" in data_paths + if ( + not root_has_data + and not ctx.config.dry_run + and not ctx.config.target_dir.exists() + ): + # Child group transactions require their parent directory. Avoid + # consolidating here, as it would be stale after child writes. + rollback_callback("delete_dir", "", None) + xr.Dataset().to_zarr( + ctx.config.target_dir.uri, + mode="w", + zarr_format=ctx.config.zarr_format, + consolidated=False, + storage_options=ctx.config.target_dir.storage_options, + ) + + for group_path in data_paths: + group_variables = ctx.config.group_variables.get(group_path) + if group_path == "/" and group_variables is None: + group_ctx = ctx + else: + group_target = ( + ctx.config.target_dir + if group_path == "/" + else ctx.config.target_dir / group_path.lstrip("/") + ) + group_config = dict(self._config_dict) + group_config["target_dir"] = group_target.uri + group_config["variables"] = merge_configs( + self._config_dict.get("variables") or {}, group_variables or {} + ) + # force_new was already applied once to the complete root in + # Processor.__init__, never once per child group. + group_config["force_new"] = False + group_ctx = Context(group_config) + self.process_dataset_slice( + group_ctx, + tree[group_path].ds, + rollback_callback=prefix_rollback_callback( + rollback_callback, group_path + ), + ) + + if not root_has_data and not ctx.config.dry_run: + update_datatree_root_attrs(ctx, tree, rollback_callback) + + if not ctx.config.dry_run and ctx.config.zarr_format == 2: + zarr.consolidate_metadata( + make_rollback_store( + ctx.config.target_dir.fs.get_mapper(root=ctx.config.target_dir.path), + rollback_callback, + ) + ) + @staticmethod + def process_dataset_slice( + ctx: Context, + slice_dataset: xr.Dataset, + rollback_callback: RollbackCallback | None = None, + ): + """Process one Dataset slice using the normal transactional append path.""" + slice_metadata = ctx.get_dataset_metadata(slice_dataset) + if ctx.target_metadata is None: + ctx.target_metadata = slice_metadata + else: + ctx.target_metadata.assert_compatible_slice( + slice_metadata, ctx.config.append_dim + ) + + verify_append_labels(ctx, slice_dataset) + + if rollback_callback is None: transaction = Transaction( ctx.config.target_dir, ctx.config.temp_dir, disable_rollback=ctx.config.disable_rollback, ) - with transaction as rollback_callback: - if ctx.target_metadata is slice_metadata: - create_target_from_slice(ctx, slice_dataset, rollback_callback) - else: - update_target_from_slice(ctx, slice_dataset, rollback_callback) + with transaction as transaction_rollback_callback: + process_dataset_slice_with_rollback( + ctx, slice_dataset, slice_metadata, transaction_rollback_callback + ) + else: + process_dataset_slice_with_rollback( + ctx, slice_dataset, slice_metadata, rollback_callback + ) + + +def process_dataset_slice_with_rollback( + ctx: Context, + slice_dataset: xr.Dataset, + slice_metadata: Any, + rollback_callback: RollbackCallback, +): + if ctx.target_metadata is slice_metadata: + create_target_from_slice(ctx, slice_dataset, rollback_callback) + else: + update_target_from_slice(ctx, slice_dataset, rollback_callback) + + +def prefix_rollback_callback( + rollback_callback: RollbackCallback, group_path: str +) -> RollbackCallback: + """Adapt child-group paths to the root transaction's relative paths.""" + prefix = group_path.strip("/") + if not prefix: + return rollback_callback + + def prefixed_callback(action, target_path, data): + path = prefix if not target_path else f"{prefix}/{target_path}" + rollback_callback(action, path, data) + + return prefixed_callback + + +def update_datatree_root_attrs( + ctx: Context, tree: xr.DataTree, rollback_callback: RollbackCallback +): + """Apply Dataset attribute semantics to a metadata-only DataTree root.""" + root_group = zarr.open_group( + make_rollback_store( + ctx.config.target_dir.fs.get_mapper(root=ctx.config.target_dir.path), + rollback_callback, + ), + mode="a", + zarr_format=ctx.config.zarr_format, + ) + attrs = dict(tree.ds.attrs) + attrs.update(ctx.config.attrs) + if ctx.config.attrs_update_mode == "replace": + root_group.attrs.clear() + root_group.attrs.update(attrs) + elif ctx.config.attrs_update_mode == "update": + root_group.attrs.update(attrs) + elif ctx.config.attrs_update_mode == "keep": + for key, value in attrs.items(): + root_group.attrs.setdefault(key, value) + elif ctx.config.attrs_update_mode == "ignore": + root_group.attrs.update(ctx.config.attrs) def delete_target_permanently(config: Config): diff --git a/zappend/slice/cm.py b/zappend/slice/cm.py index baa8785..085eaf1 100644 --- a/zappend/slice/cm.py +++ b/zappend/slice/cm.py @@ -27,7 +27,7 @@ def __init__(self, slice_source: SliceSource): def slice_source(self) -> SliceSource: return self._slice_source - def __enter__(self) -> xr.Dataset: + def __enter__(self) -> xr.Dataset | xr.DataTree: return self._slice_source.get_dataset() def __exit__(self, *exception_args): @@ -39,7 +39,7 @@ def open_slice_dataset( ctx: Context, slice_item: Any, slice_index: int = 0, -) -> AbstractContextManager[xr.Dataset]: +) -> AbstractContextManager[xr.Dataset | xr.DataTree]: """Open the slice source for given slice item `slice_item`. The intended and only use of the returned slice source is as context @@ -64,7 +64,7 @@ class derived from `zappend.slice.SliceSource` or a function that returns * `str`: A local file path or URI pointing to a dataset file such as a Zarr or NetCDF. If it is a URI, the `ctx.slice_storage_options` apply. - * `xarray.Dataset`: An in-memory xarray dataset instance. + * `xarray.Dataset` or `xarray.DataTree`: An in-memory xarray object. * `zappend.api.FileObj`: A file object instance pointing to a dataset file in a predefined filesystem. * `zappend.api.SliceSource`: A slice source. Returned as-is. diff --git a/zappend/slice/source.py b/zappend/slice/source.py index b07de3c..8322e88 100644 --- a/zappend/slice/source.py +++ b/zappend/slice/source.py @@ -29,16 +29,14 @@ class SliceSource(ABC): """ @abstractmethod - def get_dataset(self) -> xr.Dataset: - """Open this slice source, do some processing and return a dataset of type - [xarray.Dataset](https://docs.xarray.dev/en/stable/generated/xarray.Dataset.html) - as result. + def get_dataset(self) -> xr.Dataset | xr.DataTree: + """Open this slice source and return a dataset or a data tree. This method is not intended to be called directly and is called exactly once for each instance of this class. - It should return a dataset that is compatible with - target dataset: + It should return a dataset, or a data tree whose data-bearing nodes + are compatible with their corresponding target Zarr groups: * slice must have same fixed dimensions; * append dimension must exist in slice. @@ -70,7 +68,12 @@ def dispose(self): SliceItem: TypeAlias = ( - str | FileObj | xr.Dataset | AbstractContextManager[xr.Dataset] | SliceSource + str + | FileObj + | xr.Dataset + | xr.DataTree + | AbstractContextManager[xr.Dataset | xr.DataTree] + | SliceSource ) """The possible types that can represent a slice dataset.""" @@ -106,7 +109,7 @@ def to_slice_source( return PersistentSliceSource(ctx, slice_file) if isinstance(slice_item, FileObj): return PersistentSliceSource(ctx, slice_item) - if isinstance(slice_item, xr.Dataset): + if isinstance(slice_item, (xr.Dataset, xr.DataTree)): if ctx.config.persist_mem_slices: return TemporarySliceSource(ctx, slice_item, slice_index) else: diff --git a/zappend/slice/sources/memory.py b/zappend/slice/sources/memory.py index 7b4ae47..dc480e2 100644 --- a/zappend/slice/sources/memory.py +++ b/zappend/slice/sources/memory.py @@ -17,11 +17,11 @@ class MemorySliceSource(SliceSource): slice_index: An index for slice identification (logging only) """ - def __init__(self, slice_ds: xr.Dataset, slice_index: int): + def __init__(self, slice_ds: xr.Dataset | xr.DataTree, slice_index: int): self._slice_ds = slice_ds self._slice_index = slice_index - def get_dataset(self) -> xr.Dataset: + def get_dataset(self) -> xr.Dataset | xr.DataTree: logger.info(f"Processing in-memory slice dataset #{self._slice_index}") return self._slice_ds diff --git a/zappend/slice/sources/persistent.py b/zappend/slice/sources/persistent.py index 7644a24..74c23a5 100644 --- a/zappend/slice/sources/persistent.py +++ b/zappend/slice/sources/persistent.py @@ -26,9 +26,9 @@ class PersistentSliceSource(SliceSource): def __init__(self, ctx: Context, slice_file: FileObj): self._config = ctx.config self._slice_file = slice_file - self._slice_ds: xr.Dataset | None = None + self._slice_ds: xr.Dataset | xr.DataTree | None = None - def get_dataset(self) -> xr.Dataset: + def get_dataset(self) -> xr.Dataset | xr.DataTree: logger.info(f"Opening slice dataset from {self._slice_file.uri}") self._slice_ds = self._wait_for_slice_dataset() return self._slice_ds @@ -39,7 +39,7 @@ def close(self): self._slice_ds = None logger.info(f"Slice dataset {self._slice_file.uri} closed") - def _wait_for_slice_dataset(self) -> xr.Dataset: + def _wait_for_slice_dataset(self) -> xr.Dataset | xr.DataTree: interval, timeout = self._config.slice_polling if timeout is None: return self._open_slice_dataset() @@ -59,7 +59,7 @@ def _wait_for_slice_dataset(self) -> xr.Dataset: ) time.sleep(interval) - def _open_slice_dataset(self) -> xr.Dataset: + def _open_slice_dataset(self) -> xr.Dataset | xr.DataTree: engine = self._config.slice_engine if engine is None and ( self._slice_file.path.endswith(".zarr") @@ -68,6 +68,17 @@ def _open_slice_dataset(self) -> xr.Dataset: engine = "zarr" if engine == "zarr": storage_options = self._config.slice_storage_options + tree = xr.open_datatree( + self._slice_file.uri, + engine="zarr", + storage_options=storage_options, + consolidated=self._config.zarr_format == 2, + ) + # Preserve the existing Dataset behaviour for ordinary Zarr + # stores, but expose actual nested stores as DataTrees. + if len(tree.groups) > 1: + return tree + tree.close() return xr.open_zarr( self._slice_file.uri, storage_options=storage_options, diff --git a/zappend/slice/sources/temporary.py b/zappend/slice/sources/temporary.py index c2f7b91..cd86ee4 100644 --- a/zappend/slice/sources/temporary.py +++ b/zappend/slice/sources/temporary.py @@ -21,29 +21,40 @@ class TemporarySliceSource(MemorySliceSource): slice_index: An index for slice identification """ - def __init__(self, ctx: Context, slice_ds: xr.Dataset, slice_index: int): + def __init__( + self, ctx: Context, slice_ds: xr.Dataset | xr.DataTree, slice_index: int + ): super().__init__(slice_ds, slice_index) self._config = ctx.config self._temp_slice_dir: FileObj | None = None - self._temp_slice_ds: xr.Dataset | None = None + self._temp_slice_ds: xr.Dataset | xr.DataTree | None = None - def get_dataset(self) -> xr.Dataset: + def get_dataset(self) -> xr.Dataset | xr.DataTree: slice_index = self._slice_index temp_slice_dir = self._config.temp_dir / f"slice-{slice_index}.zarr" self._temp_slice_dir = temp_slice_dir logger.info( f"Persisting in-memory slice dataset #{slice_index} to {temp_slice_dir.uri}" ) + is_datatree = isinstance(self._slice_ds, xr.DataTree) self._slice_ds.to_zarr( temp_slice_dir.uri, zarr_format=self._config.zarr_format, ) self._slice_ds = None - self._temp_slice_ds = xr.open_zarr( - temp_slice_dir.uri, - zarr_format=self._config.zarr_format, - consolidated=self._config.zarr_format == 2, - ) + if is_datatree: + self._temp_slice_ds = xr.open_datatree( + temp_slice_dir.uri, + engine="zarr", + zarr_format=self._config.zarr_format, + consolidated=self._config.zarr_format == 2, + ) + else: + self._temp_slice_ds = xr.open_zarr( + temp_slice_dir.uri, + zarr_format=self._config.zarr_format, + consolidated=self._config.zarr_format == 2, + ) return self._temp_slice_ds def close(self):