diff --git a/design_notes/variable_coders.md b/design_notes/variable_coders.md new file mode 100644 index 00000000000..e2321952261 --- /dev/null +++ b/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 diff --git a/xarray/coding/api.py b/xarray/coding/api.py new file mode 100644 index 00000000000..da834c5f2c0 --- /dev/null +++ b/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 diff --git a/xarray/coding/cf.py b/xarray/coding/cf.py new file mode 100644 index 00000000000..c590c4f9d75 --- /dev/null +++ b/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 diff --git a/xarray/coding/core.py b/xarray/coding/core.py new file mode 100644 index 00000000000..f40a334bd1d --- /dev/null +++ b/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: ... diff --git a/xarray/coding/range.py b/xarray/coding/range.py new file mode 100644 index 00000000000..8e52168a938 --- /dev/null +++ b/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 diff --git a/xarray/testing/assertions.py b/xarray/testing/assertions.py index 8b17b94735f..a8d51798c3e 100644 --- a/xarray/testing/assertions.py +++ b/xarray/testing/assertions.py @@ -372,7 +372,15 @@ def _assert_indexes_invariants_checks( for k, v in possible_coord_variables.items() if isinstance(v, IndexVariable) } - assert indexes.keys() <= index_vars, (set(indexes), index_vars) + only_default_indexes = { + name: index + for name, index in indexes.items() + if isinstance(index, PandasIndex) + } + assert only_default_indexes.keys() <= index_vars, ( + set(only_default_indexes), + index_vars, + ) assert all( k in index_vars for k, v in possible_coord_variables.items() @@ -416,9 +424,17 @@ def _assert_indexes_invariants_checks( if check_default: defaults = default_indexes(possible_coord_variables, dims) - assert indexes.keys() == defaults.keys(), (set(indexes), set(defaults)) - assert all(v.equals(defaults[k]) for k, v in indexes.items()), ( - indexes, + only_default_indexes = { + name: index + for name, index in indexes.items() + if isinstance(index, PandasIndex) + } + assert only_default_indexes.keys() == defaults.keys(), ( + set(indexes), + set(defaults), + ) + assert all(v.equals(defaults[k]) for k, v in only_default_indexes.items()), ( + only_default_indexes, defaults, ) diff --git a/xarray/tests/test_coders_cf.py b/xarray/tests/test_coders_cf.py new file mode 100644 index 00000000000..3f85a5980b9 --- /dev/null +++ b/xarray/tests/test_coders_cf.py @@ -0,0 +1,60 @@ +import numpy as np + +import xarray as xr +from xarray.coding import cf + + +class TestCFCoordinateCoder: + def test_init(self): + coder = cf.CFCoordinateCoder() + + assert coder.kind == "dataset" + + def test_decode(self): + encoded = xr.Dataset( + { + "var1": ("x", [1, 2, 3], {"coordinates": "coord1 coord2"}), + "var2": ( + ("x", "y"), + [[9, 4], [3, 6], [87, 7]], + {"coordinates": "coord1 coord2 coord3 coord4"}, + ), + "x": ("x", [0, 1, 2]), + "y": ("y", [-1, 1]), + "coord1": ("x", [0, 2, 4]), + "coord2": ("x", [1, 3, 5]), + "coord3": ("y", [7, 6]), + "coord4": (["x", "y"], np.ones((3, 2))), + "coord5": ("z", [0]), + }, + coords=xr.Coordinates(), + attrs={"coordinates": "coord5"}, + ) + + coder = cf.CFCoordinateCoder() + decoded = coder.decode(encoded) + + expected = {"x", "y", "coord1", "coord2", "coord3", "coord4", "coord5"} + + assert set(decoded.coords) == expected + + def test_encode(self): + decoded = xr.Dataset( + {"var1": ("x", [1, 2, 3]), "var2": (["x", "y"], [[9, 4], [3, 6], [87, 7]])}, + coords={ + "x": [0, 1, 2], + "y": [-1, 1], + "coord1": ("x", [0, 2, 4]), + "coord2": ("x", [1, 3, 5]), + "coord3": ("y", [7, 6]), + "coord4": (["x", "y"], np.ones((3, 2))), + "coord5": ("z", [0]), + }, + ) + + coder = cf.CFCoordinateCoder() + encoded = coder.encode(decoded) + + assert encoded.attrs == {"coordinates": "coord5"} + assert encoded["var1"].attrs == {"coordinates": "coord1 coord2"} + assert encoded["var2"].attrs == {"coordinates": "coord1 coord2 coord3 coord4"} diff --git a/xarray/tests/test_coders_range.py b/xarray/tests/test_coders_range.py new file mode 100644 index 00000000000..48c86319ae6 --- /dev/null +++ b/xarray/tests/test_coders_range.py @@ -0,0 +1,40 @@ +import numpy as np + +import xarray as xr +from xarray.coding.range import RangeIndexCoder +from xarray.indexes import RangeIndex +from xarray.tests import assert_identical + + +def test_encode() -> None: + index = RangeIndex.arange(-10, 10, 1, dim="x", coord_name="l") + ds = xr.Dataset( + {"a": ("x", np.arange(20))}, coords=xr.Coordinates.from_xindex(index) + ) + + coder = RangeIndexCoder() + encoded = coder.encode(ds) + + expected = {"l": {"start": -10, "stop": 10, "step": 1, "dim": "x"}} + + assert set(encoded.variables) == {"a"} # x is gone + assert not set(encoded.xindexes) # no indexes + assert encoded.attrs["ranges"] == expected + + +def test_decode() -> None: + ds = xr.Dataset( + {"a": ("x", np.arange(10))}, + attrs={"ranges": {"l": {"start": -5, "stop": 0, "step": 0.5, "dim": "x"}}}, + ) + + coder = RangeIndexCoder() + actual = coder.decode(ds) + + expected = xr.Dataset( + {"a": ("x", np.arange(10))}, + coords=xr.Coordinates.from_xindex( + RangeIndex.arange(-5, 0, 0.5, dim="x", coord_name="l") + ), + ) + assert_identical(actual, expected, check_indexes=True)