encoding and decoding using custom coders

design_notes/variable_coders.md

@@ -0,0 +1,92 @@
+# dataset and variable coders
+
+## current design
+
+When opening datasets, xarray currently applies "coders" to the variables. The exact coders depend on the backend:
+
+Code paths:
+
+1. decode
+
+- `open_dataset`
+- `backend.open_dataset`
+- `StoreBackendEntrypoint` → `conventions.decode_cf_variables`
+
+2. encode
+
+- `to_*` → `dump_to_store`
+  → `encode_dataset_coordinates` (for cf-style coordinate storage)
+  → `store.store` → `store.encode` → `store.encode_variable` → `encode_zarr_variable` → `encode_cf_variable`
+
+All backends use `decode_cf_variable` (through the `StoreBackendEntrypoint`), while `zarr` is the only backend that also uses `encode_cf_variable`.
+
+Within `encode_cf_variable`:
+
+- cfdatetime coder
+- cftimedelta coder
+- cfscaleoffset coder
+- cf mask coder
+- native enum coder
+- nonstring coder
+- default fillvalue coder
+- boolean coder
+  Additionally, `decode_cf_variable` has:
+- characterarray coder
+- encoded string coder
+- objectvlen string coder
+- numpy2 string dtype coder
+- endian coder
+
+These coders try to apply their operations in a lazy way, such that the actual computation is only triggered when explicitly or implicitly requested.
+
+## new: custom coders
+
+### variable coders
+
+Variable coders will follow a protocol (not a ABC), with two methods:
+
+- `VariableCoder.encode(variable, *, **additional_metadata)`
+- `VariableCoder.decode(variable, *, **additional_metadata)`
+
+It will also need a heuristic to decide whether the coder should be applied. This could be:
+
+- a function that, given metadata (dtype, attributes, encoding), decides whether to apply the coder
+- the coder performs the check and returns `NotImplemented` if it doesn't fit.
+
+### dataset coders
+
+Dataset coders have a very similar API (still a protocol):
+
+- `DatasetCoder.encode(dataset, **additional_metadata)`
+- `DatasetCoder.decode(dataset, **additional_metadata)`
+
+Just like with variable coders it might make sense to have a function that determines whether a coder is applicable given dataset structure and attributes.
+
+### coder pipelines
+
+Coder pipelines describe a set of coding operations of the same type.
+
+TODO: defaults based on the existing attributes (for a CF coder variable and dataset pipeline).
+
+## processing steps
+
+### decoding
+
+The `StoreBackendEntrypoint.open_dataset` method will be split up into different parts (as functions):
+
+- load variables and attributes from the datastore
+- apply variable coders (given a ordered list of coders)
+- construct a backend dataset from variables, attrs, and the file object
+- apply dataset coders (by default contains a CF `coordinates` dataset coder)
+
+Where only the first two will stay in the `StoreBackendEntrypoint` (?).
+
+Then there will be a constructor / object that, given the cf coder settings, constructs a list of variable coders that need to be applied. Backends can then filter these coders to only select those that apply.
+
+### encoding
+
+The steps from `decoding` can be inverted:
+
+- apply dataset coders
+- apply variable coders
+- split dataset into variables and attrs

xarray/coding/api.py

@@ -0,0 +1,162 @@
+from collections.abc import Mapping
+from typing import Literal
+
+from xarray.coding.core import Coder
+from xarray.coding.times import CFDatetimeCoder, CFTimedeltaCoder
+from xarray.core.dataset import Dataset
+from xarray.core.variable import Variable
+
+
+def cf_coders(
+    mask_and_scale: bool | Mapping[str, bool] | None = None,
+    decode_times: (
+        bool | CFDatetimeCoder | Mapping[str, bool | CFDatetimeCoder] | None
+    ) = None,
+    decode_timedelta: (
+        bool | CFTimedeltaCoder | Mapping[str, bool | CFTimedeltaCoder] | None
+    ) = None,
+    concat_characters: bool | Mapping[str, bool] | None = None,
+    decode_coords: Literal["coordinates", "all"] | bool | None = None,
+) -> tuple[list[Coder[Dataset]], list[Coder[Variable]]]:
+    """Create the default cf coders
+
+    Parameters
+    ----------
+    mask_and_scale : bool or mapping of str to bool, optional
+        If True, replace array values equal to `_FillValue` with NA and scale
+        values according to the formula `original_values * scale_factor +
+        add_offset`, where `_FillValue`, `scale_factor` and `add_offset` are
+        taken from variable attributes (if they exist).  If the `_FillValue` or
+        `missing_value` attribute contains multiple values a warning will be
+        issued and all array values matching one of the multiple values will
+        be replaced by NA. Pass a mapping, e.g. ``{"my_variable": False}``,
+        to toggle this feature per-variable individually.
+    decode_times : bool, CFDatetimeCoder or dict-like, optional
+        If True, decode times encoded in the standard NetCDF datetime format
+        into datetime objects. Otherwise, use :py:class:`coders.CFDatetimeCoder` or leave them
+        encoded as numbers.
+        Pass a mapping, e.g. ``{"my_variable": False}``,
+        to toggle this feature per-variable individually.
+    decode_timedelta : bool, CFTimedeltaCoder, or dict-like, optional
+        If True, decode variables and coordinates with time units in
+        {"days", "hours", "minutes", "seconds", "milliseconds", "microseconds"}
+        into timedelta objects. If False, leave them encoded as numbers.
+        If None (default), assume the same value of ``decode_times``; if
+        ``decode_times`` is a :py:class:`coders.CFDatetimeCoder` instance, this
+        takes the form of a :py:class:`coders.CFTimedeltaCoder` instance with a
+        matching ``time_unit``.
+        Pass a mapping, e.g. ``{"my_variable": False}``,
+        to toggle this feature per-variable individually.
+    concat_characters : bool or dict-like, optional
+        If True, concatenate along the last dimension of character arrays to
+        form string arrays. Dimensions will only be concatenated over (and
+        removed) if they have no corresponding variable and if they are only
+        used as the last dimension of character arrays.
+        Pass a mapping, e.g. ``{"my_variable": False}``,
+        to toggle this feature per-variable individually.
+        This keyword may not be supported by all the backends.
+    decode_coords : bool or {"coordinates", "all"}, optional
+        Controls which variables are set as coordinate variables:
+
+        - "coordinates" or True: Set variables referred to in the
+          ``'coordinates'`` attribute of the datasets or individual variables
+          as coordinate variables.
+        - "all": Set variables referred to in  ``'grid_mapping'``, ``'bounds'`` and
+          other attributes as coordinate variables.
+
+        Only existing variables can be set as coordinates. Missing variables
+        will be silently ignored.
+
+    Returns
+    -------
+    dataset_coders : list of Coder
+        The constructed dataset coders.
+    variable_coders : list of Coder
+        The constructed variable coders.
+
+    See Also
+    --------
+    decode, encode
+    decode_cf
+    """
+
+
+def encode(
+    obj: Dataset,
+    *,
+    dataset_coders: list[Coder[Dataset]] | None = None,
+    variable_coders: list[Coder[Variable]] | None = None,
+) -> Dataset:
+    """Encode a dataset using the given coders
+
+    Parameters
+    ----------
+    obj : xarray.Dataset
+        The dataset to encode.
+    dataset_coders : list of Coder, optional
+        The dataset coders to apply.
+    variable_coders : list of Coder, optional
+        The variable coders to apply.
+
+    Returns
+    -------
+    xarray.Dataset
+        The encoded dataset after applying all coders.
+
+    Notes
+    -----
+    Coders that return `NotImplemented` will be skipped.
+
+    See Also
+    --------
+    decode, decode_cf
+    """
+    encoded = obj
+    for coder in dataset_coders:
+        encoded = coder.encode(encoded)
+
+    # by applying the dataset coders first there are no coordinates anymore (though we may want to check that)
+    for coder in variable_coders:
+        encoded = encoded.map(coder.encode)
+
+    return encoded
+
+
+def decode(
+    obj: Dataset,
+    *,
+    dataset_coders: list[Coder[Dataset]] | None = None,
+    variable_coders: list[Coder[Variable]] | None = None,
+) -> Dataset:
+    """Decode a dataset using the given coders
+
+    Parameters
+    ----------
+    obj : xarray.Dataset
+        The dataset to decode.
+    dataset_coders : list of Coder, optional
+        The dataset coders to apply.
+    variable_coders : list of Coder, optional
+        The variable coders to apply.
+
+    Returns
+    -------
+    xarray.Dataset
+        The decoded dataset after applying all coders.
+
+    Notes
+    -----
+    Coders that return `NotImplemented` will be skipped.
+
+    See Also
+    --------
+    encode, decode_cf
+    """
+    decoded = obj
+    for coder in variable_coders:
+        decoded = decoded.map(coder.decode)
+
+    for coder in dataset_coders:
+        decoded = coder.decode(decoded)
+
+    return decoded

xarray/coding/cf.py

@@ -0,0 +1,66 @@
+import itertools
+from typing import ClassVar
+
+from xarray.coding.core import CoderKind
+from xarray.coding.variables import SerializationWarning
+from xarray.core.dataset import Dataset
+from xarray.core.utils import emit_user_level_warning
+
+
+class CFCoordinateCoder:
+    """CF coordinate coder
+
+    Allows for roundtripping variables as coordinates. Note that `xarray`
+    associates coordinates based on dimensions, so the association of
+    coordinates with specific coordinates is lost.
+    """
+
+    kind: ClassVar[CoderKind] = "dataset"
+
+    def decode(self, obj: Dataset) -> Dataset:
+        unparsed = [
+            var.attrs.pop("coordinates", None) for var in obj.variables.values()
+        ] + [obj.attrs.get("coordinates", None)]
+
+        dim_coords = [name for name in obj.variables if name in obj.dims]
+        non_dim_coords = list(
+            itertools.chain.from_iterable(
+                coordinates.split(" ")
+                for coordinates in unparsed
+                if isinstance(coordinates, str)
+            )
+        )
+
+        return obj.set_coords(dim_coords + non_dim_coords)
+
+    def encode(self, obj: Dataset) -> Dataset:
+        encoded = obj.copy(deep=False)
+        coords = dict(encoded.coords)
+        for name in list(coords):
+            if isinstance(name, str) and " " in name:
+                emit_user_level_warning(
+                    f"coordinate {name!r} has a space in its name, which means it "
+                    "cannot be marked as a coordinate on disk and will be "
+                    "saved as a data variable instead",
+                    category=SerializationWarning,
+                )
+                del coords[name]
+
+        covered = set()
+        for variable in encoded.values():
+            dims = set(variable.dims)
+            coordinates = [
+                name
+                for name, coord in coords.items()
+                if name not in encoded.dims and set(coord.dims).issubset(dims)
+            ]
+            covered.update(coordinates)
+            variable.attrs["coordinates"] = " ".join(map(str, coordinates))
+
+        uncovered = [
+            name for name in coords if name not in covered and name not in encoded.dims
+        ]
+        encoded.attrs["coordinates"] = " ".join(uncovered)
+
+        # TODO: compare with the algorithm in xarray.conventions._encode_coordinates
+        return encoded

xarray/coding/core.py

@@ -0,0 +1,12 @@
+from typing import ClassVar, Generic, Literal, Protocol, TypeVar
+
+CoderKind = Literal["variable", "dataset", "datatree"]
+T = TypeVar("T")
+
+
+class Coder(Protocol, Generic[T]):
+    kind: ClassVar[CoderKind]
+
+    def decode(self, obj: T) -> T: ...
+
+    def encode(self, obj: T) -> T: ...

xarray/coding/range.py

@@ -0,0 +1,48 @@
+import operator
+from functools import reduce
+from typing import ClassVar
+
+import xarray as xr
+from xarray.coding.core import CoderKind
+
+
+class RangeIndexCoder:
+    kind: ClassVar[CoderKind] = "dataset"
+
+    def decode(self, obj: xr.Dataset) -> xr.Dataset:
+        def _decode_index(coord, params):
+            return xr.indexes.RangeIndex.arange(**params, coord_name=coord)
+
+        encoded_ranges = obj.attrs.get("ranges")
+        if encoded_ranges is None:
+            return obj
+
+        indexes = [
+            _decode_index(name, params) for name, params in encoded_ranges.items()
+        ]
+        coords = reduce(
+            operator.or_, (xr.Coordinates.from_xindex(index) for index in indexes)
+        )
+
+        decoded = obj.assign_coords(coords)
+        del decoded.attrs["ranges"]
+
+        return decoded
+
+    def encode(self, obj: xr.Dataset) -> xr.Dataset:
+        def _encode_index(index):
+            return {
+                "start": index.start,
+                "stop": index.stop,
+                "step": index.step,
+                "dim": index.dim,
+            }
+
+        encoded_ranges = {
+            name: _encode_index(index)
+            for name, index in obj.xindexes.items()
+            if isinstance(index, xr.indexes.RangeIndex)
+        }
+        encoded = obj.drop_indexes(list(encoded_ranges)).drop_vars(list(encoded_ranges))
+        encoded.attrs["ranges"] = encoded_ranges
+        return encoded