diff --git a/README.md b/README.md index 2e5d60c4..b25405de 100644 --- a/README.md +++ b/README.md @@ -18,7 +18,7 @@ We are currently migrating to this repo from [ocf_datapipes](https://github.com/ This project is primarily licensed under the MIT License (see LICENSE). -It includes and adapts internal functions from the Google xarray-tensorstore project, licensed under the Apache License, Version 2.0. +It includes and adapts the `ephemeris()` function from the [pvlib-python](https://github.com/pvlib/pvlib-python) project, licensed under the BSD 3-Clause License. The original notice is retained in `src/ocf_data_sampler/features/solar.py`. ## Documentation diff --git a/pyproject.toml b/pyproject.toml index 46a07957..b744e2d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,12 +26,11 @@ dependencies = [ "pandas", "xarray", "numcodecs", - "dask", "pydantic>=2.11.8", "pyproj", "pyaml_env", "pyresample", - "xarray-tensorstore==0.1.5", + "xarray-tensorstore==0.3.0", "zarr>=3", ] @@ -40,6 +39,7 @@ dev = [ # Testing "pytest", "pvlib", + "dask", "ruff >= 0.9.2", "ty", # LSP Support diff --git a/src/ocf_data_sampler/common/xr_tensorstore.py b/src/ocf_data_sampler/common/xr_tensorstore.py index aafdbc35..b61ccdea 100644 --- a/src/ocf_data_sampler/common/xr_tensorstore.py +++ b/src/ocf_data_sampler/common/xr_tensorstore.py @@ -1,96 +1,110 @@ -"""Utilities for loading TensorStore data into Xarray. +"""Utilities for opening and lazily concatenating TensorStore-backed Xarray datasets. -This module uses and adapts internal functions from the Google xarray-tensorstore project [1], -licensed under the Apache License, Version 2.0. See [2] for details. - -Modifications copyright 2025 Open Climate Fix. Licensed under the MIT License. - -Modifications from the original include: -- Adding support for opening multiple zarr files as a single xarray object -- Support for zarr 3 -> https://github.com/google/xarray-tensorstore/pull/22 - -References: - [1] https://github.com/google/xarray-tensorstore - [2] https://www.apache.org/licenses/LICENSE-2.0 +Note: this module relies on `xarray_tensorstore` internals (`_TensorStoreAdapter`) to reach +the underlying TensorStore without materialising data. The dependency is pinned in +pyproject.toml; upgrades need checking against this. """ -import logging -import os -import re +from collections.abc import Sequence from glob import glob, has_magic -from typing import Any, TypeAlias, cast +from typing import TypeAlias import tensorstore as ts import xarray as xr -import zarr -from xarray_tensorstore import ( - _DEFAULT_STORAGE_DRIVER, - _raise_if_mask_and_scale_used_for_data_vars, - _TensorStoreAdapter, -) - -logger = logging.getLogger(__name__) +import xarray_tensorstore as xrt ZarrSource: TypeAlias = str | list[str] | tuple[str, ...] -def _zarr_spec_from_path(path: str, zarr_format: int) -> dict[str, Any]: - if re.match(r"\w+\://", path): # path is a URI - kv_store: str | dict[str, str] = path - else: - kv_store = {"driver": _DEFAULT_STORAGE_DRIVER, "path": path} - return {"driver": f"zarr{zarr_format}", "kvstore": kv_store} - - -def _get_data_variable_array_futures( - path: str, - context: ts.Context | None, - variables: list[str], -) -> dict[str, ts.Future[ts.TensorStore]]: - """Open all data variables in a zarr group and return futures. - - Args: - path: path or URI to zarr group to open. - context: TensorStore configuration options to use when opening arrays. - variables: The variables in the zarr groupto open. - """ - zarr_format = zarr.open(path).metadata.zarr_format - specs = {k: _zarr_spec_from_path(os.path.join(path, k), zarr_format) for k in variables} - return {k: ts.open(spec, read=True, write=False, context=context) for k, spec in specs.items()} - - -def _tensorstore_open_zarrs( - paths: list[str], - data_vars: list[str], - concat_axes: list[int], - context: ts.Context, -) -> dict[str, ts.TensorStore]: - """Open multiple zarrs with TensorStore. +def _extract_tensorstore(da: xr.DataArray) -> ts.TensorStore: + """Extract the backing TensorStore, or fail with a message that says why.""" + data = da.variable._data + if not isinstance(data, xrt._TensorStoreAdapter): + raise TypeError(f"{da.name!r} is backed by {type(data).__name__}, expected TensorStore.") + return data.array + + +def _validate(datasets: Sequence[xr.Dataset], concat_dim: str) -> None: + first, *rest = datasets + if concat_dim not in first.dims: + raise ValueError(f"{concat_dim!r} is not a dimension: {tuple(first.dims)}") + + for i, ds in enumerate(rest, start=1): + + # All coords and data_vars must be present in all datasets + if set(ds.coords) != set(first.coords): + raise ValueError( + f"dataset {i}: coords {sorted(ds.coords, key=str)} " + f"!= {sorted(first.coords, key=str)}" + ) + if set(ds.data_vars) != set(first.data_vars): + raise ValueError( + f"dataset {i}: data_vars {sorted(ds.data_vars, key=str)} " + f"!= {sorted(first.data_vars, key=str)}" + ) + + # All dims except concat_dim must match in size + for dim, size in first.sizes.items(): + if dim != concat_dim and ds.sizes.get(dim) != size: + raise ValueError(f"dataset {i}: {dim}={ds.sizes.get(dim)}, expected {size}") + + # All data_vars must have the same dims and dtype + for name in first.data_vars: + if ds[name].dims != first[name].dims or ds[name].dtype != first[name].dtype: + raise ValueError( + f"dataset {i}: {name!r} is {ds[name].dims}/{ds[name].dtype}, " + f"expected {first[name].dims}/{first[name].dtype}" + ) + + # All coords and data_vars which don't contain the concat_dim dimension must be identical + # Note: `.equals()` reads lazy data into memory. This is fine for coords and static vars, + # which should be small + for name in [*first.coords, *first.data_vars]: + if concat_dim not in first[name].dims and not ds[name].equals(first[name]): + raise ValueError( + f"dataset {i}: {name!r} does not span {concat_dim!r} but differs" + ) + + +def concat_tensorstore(datasets: Sequence[xr.Dataset], concat_dim: str) -> xr.Dataset: + """Concatenate tensorstore-backed Datasets along an existing dimension, lazily. + + Data variables containing the `concat_dim` dimension are concatenated lazily using TensorStore. + Everything else must match across datasets and is taken from the first dataset, as are attrs. Args: - paths: List of paths to zarr stores. - data_vars: List of data variable names to open. - concat_axes: List of axes along which to concatenate the data variables. - context: TensorStore context. + datasets: Sequence of Datasets to concatenate. + concat_dim: Dimension along which to concatenate. """ - # Open all the variables from all the datasets - returned as futures - array_futures_list: list[dict[str, ts.Future[ts.TensorStore]]] = [] - for path in paths: - array_futures_list.append(_get_data_variable_array_futures(path, context, data_vars)) - - # Wait for the async open operations - arrays_list: list[dict[str, ts.TensorStore]] = [ - {k: future.result() for k, future in array_futures.items()} - for array_futures in array_futures_list - ] - - # Concatenate each of the variables along the required axis - arrays: dict[str, ts.TensorStore] = {} - for k, axis in zip(data_vars, concat_axes, strict=True): - variable_arrays = [d[k] for d in arrays_list] - arrays[k] = ts.concat(variable_arrays, axis=axis) - - return arrays + datasets = list(datasets) + if len(datasets) < 2: + raise ValueError("need at least two datasets") + _validate(datasets, concat_dim) + first = datasets[0] + + # Create a new shell dataset which contains only the concatenated coords. We will handle the + # data_vars separately so we can lazily concatenate them with tensorstore. + # - combine_attrs="override" keeps the attrs of the first dataset, which is the behaviour we + # copy for the data_vars below. + # - join="exact" ensures that the coords are identical across datasets, which is a bakstop for + # the _validate() check above. + ds_out = xr.concat( + [ds.drop_vars(first.data_vars) for ds in datasets], + dim=concat_dim, + join="exact", + combine_attrs="override", + ) + + for name, da in first.data_vars.items(): + if concat_dim in da.dims: + store = ts.concat( + [_extract_tensorstore(ds[name]) for ds in datasets], + axis=da.dims.index(concat_dim), + ) + ds_out[name] = xr.Variable(da.dims, xrt._TensorStoreAdapter(store), attrs=da.attrs) + else: + ds_out[name] = da.variable # attrs travel with the Variable; xarray copies on assign + return ds_out def open_zarr_paths(zarr_path: ZarrSource, concat_dim: str | None = None) -> xr.Dataset: @@ -103,7 +117,7 @@ def open_zarr_paths(zarr_path: ZarrSource, concat_dim: str | None = None) -> xr. if isinstance(zarr_path, str): path = zarr_path if not has_magic(path): - return _open_single_zarr(path) + return xrt.open_zarr(path) paths = sorted(glob(path)) else: paths = list(zarr_path) @@ -112,111 +126,9 @@ def open_zarr_paths(zarr_path: ZarrSource, concat_dim: str | None = None) -> xr. raise ValueError(f"No Zarr stores found for {zarr_path!r}") if len(paths) == 1: - return _open_single_zarr(paths[0]) + return xrt.open_zarr(paths[0]) if concat_dim is None: raise ValueError("`concat_dim` must be specified when opening multiple Zarr stores") - return _open_and_concat_zarrs(paths, concat_dim) - - -def _open_single_zarr( - path: str, - context: ts.Context | None = None, - mask_and_scale: bool = True, -) -> xr.Dataset: - """Open an xarray.Dataset from zarr using TensorStore. - - Args: - path: path or URI to zarr group to open. - context: TensorStore configuration options to use when opening arrays. - mask_and_scale: if True (default), attempt to apply masking and scaling like - xarray.open_zarr(). This is only supported for coordinate variables and - otherwise will raise an error. - - Returns: - Dataset with all data variables opened via TensorStore. - """ - if context is None: - context = ts.Context() - - # Avoid using dask by settung `chunks=None` - ds = xr.open_zarr(path, chunks=None, mask_and_scale=mask_and_scale, consolidated=False) - - if mask_and_scale: - _raise_if_mask_and_scale_used_for_data_vars(ds) - - # Open all data variables using tensorstore - returned as futures - data_vars = list(ds.data_vars) - array_futures = _get_data_variable_array_futures(path, context, data_vars) - - # Wait for the async open operations - arrays = {k: future.result() for k, future in array_futures.items()} - - # Adapt the tensorstore arrays and plug them into the xarray object - new_data = {k: _TensorStoreAdapter(v) for k, v in arrays.items()} - - return cast("xr.Dataset", ds.copy(data=new_data)) - - -def _open_and_concat_zarrs( - paths: list[str], - concat_dim: str, - context: ts.Context | None = None, - mask_and_scale: bool = True, -) -> xr.Dataset: - """Open multiple zarrs with TensorStore. - - Args: - paths: List of paths to zarr stores. - concat_dim: Dimension along which to concatenate the data variables. - context: TensorStore context. - mask_and_scale: Whether to mask and scale the data. - - Returns: - Concatenated Dataset with all data variables opened via TensorStore. - """ - if context is None: - context = ts.Context() - - ds_list = [ - xr.open_zarr(p, mask_and_scale=mask_and_scale, decode_timedelta=True, consolidated=False) - for p in paths - ] - try: - ds = xr.concat( - ds_list, - dim=concat_dim, - data_vars="minimal", - compat="equals", - combine_attrs="drop_conflicts", - join="exact", - ) - except ValueError: - logger.warning( - f"Coordinate mismatch found when opening paths {paths}. Opening with `join='override'` " - "to ignore coordinate mismatches. THIS MAY CAUSE UNEXPECTED BEHAVIOUR.", - ) - ds = xr.concat( - ds_list, - dim=concat_dim, - data_vars="minimal", - compat="equals", - combine_attrs="drop_conflicts", - join="override", - ) - - if mask_and_scale: - _raise_if_mask_and_scale_used_for_data_vars(ds) - - # Find the axis along which each data array must be concatenated - data_vars = list(ds.data_vars) - concat_axes = [ds[v].dims.index(concat_dim) for v in data_vars] - - # Open and concat all zarrs so each variables is a single TensorStore array - arrays = _tensorstore_open_zarrs(paths, data_vars, concat_axes, context) - - # Plug the arrays into the xarray object - new_data = {k: _TensorStoreAdapter(v) for k, v in arrays.items()} - - return cast("xr.Dataset", ds.copy(data=new_data)) + return concat_tensorstore([xrt.open_zarr(p) for p in paths], concat_dim) diff --git a/tests/common/test_xr_tensorstore.py b/tests/common/test_xr_tensorstore.py index be6aedac..02b6a671 100644 --- a/tests/common/test_xr_tensorstore.py +++ b/tests/common/test_xr_tensorstore.py @@ -2,11 +2,17 @@ import pytest import xarray as xr -from ocf_data_sampler.common.xr_tensorstore import open_zarr_paths +from ocf_data_sampler.common.xr_tensorstore import ( + _extract_tensorstore, + concat_tensorstore, + open_zarr_paths, +) + +ZARR_FORMATS = [2, 3] @pytest.fixture(scope="module") -def concatable_nwp_like_data(ds_nwp_ecmwf): +def concatable_nwp_like_data(ds_nwp_ecmwf) -> tuple[xr.Dataset, xr.Dataset]: """Create two NWP datasets with consecutive init times for concatenation""" ds_2 = ds_nwp_ecmwf.copy(deep=True) ds_2["init_time_utc"] = pd.date_range( @@ -17,51 +23,88 @@ def concatable_nwp_like_data(ds_nwp_ecmwf): return ds_nwp_ecmwf, ds_2 -def _save_nwp_zarr(session_tmp_path, datasets, zarr_format): - """Save NWP datasets to zarr with specified format""" - paths = [f"{session_tmp_path}/nwp_like_data_{n}.zarr{zarr_format}" - for n in range(len(datasets))] - for ds, path in zip(datasets, paths, strict=True): - ds.to_zarr(path, zarr_format=zarr_format) +@pytest.fixture(scope="module") +def zarr_paths(session_tmp_path, concatable_nwp_like_data) -> dict[int, list[str]]: + """Save the NWP datasets to zarr, keyed by zarr format""" + paths: dict[int, list[str]] = {} + for zarr_format in (2, 3): + paths[zarr_format] = [ + f"{session_tmp_path}/nwp_like_data_{n}.zarr{zarr_format}" + for n in range(len(concatable_nwp_like_data)) + ] + for ds, path in zip(concatable_nwp_like_data, paths[zarr_format], strict=True): + ds.to_zarr(path, zarr_format=zarr_format) return paths @pytest.fixture(scope="module") -def nwp_like_zarr2_paths(session_tmp_path, concatable_nwp_like_data): - """Save NWP datasets as zarr format 2""" - return _save_nwp_zarr(session_tmp_path, concatable_nwp_like_data, 2) +def tensorstore_datasets(zarr_paths) -> list[xr.Dataset]: + """The two consecutive datasets, opened as tensorstore-backed datasets""" + return [open_zarr_paths(path) for path in zarr_paths[2]] -@pytest.fixture(scope="module") -def nwp_like_zarr3_paths(session_tmp_path, concatable_nwp_like_data): - """Save NWP datasets as zarr format 3""" - return _save_nwp_zarr(session_tmp_path, concatable_nwp_like_data, 3) +@pytest.mark.parametrize("zarr_format", ZARR_FORMATS) +def test_open_single_zarr(zarr_paths, zarr_format): + path = zarr_paths[zarr_format][0] + # Check tensorstore version returns same results as dask version + assert open_zarr_paths(path).compute().equals(xr.open_zarr(path).compute()) -def test_open_single_zarr(nwp_like_zarr2_paths, nwp_like_zarr3_paths): - # Check function can open zarr2 - ds_ts = open_zarr_paths(nwp_like_zarr2_paths[0]) +@pytest.mark.parametrize("zarr_format", ZARR_FORMATS) +def test_open_multi_zarrs(zarr_paths, zarr_format): + paths = zarr_paths[zarr_format] + ds_ts = open_zarr_paths(paths, concat_dim="init_time_utc") # Check tensorstore version returns same results as dask version - ds_dask = xr.open_zarr(nwp_like_zarr2_paths[0]) + ds_dask = xr.open_mfdataset( + paths, concat_dim="init_time_utc", combine="nested", engine="zarr", + ) assert ds_ts.compute().equals(ds_dask.compute()) - # Check function can open zarr3 - ds_ts = open_zarr_paths(nwp_like_zarr3_paths[0]) - # Check tensorstore version returns same results as dask version - ds_dask = xr.open_zarr(nwp_like_zarr3_paths[0]) - assert ds_ts.compute().equals(ds_dask.compute()) +def test_open_multi_zarrs_requires_concat_dim(zarr_paths): + with pytest.raises(ValueError, match="`concat_dim` must be specified"): + open_zarr_paths(zarr_paths[2]) -def test_open_multi_zarrs(nwp_like_zarr2_paths, nwp_like_zarr3_paths): - # Check function can open zarr2 - ds_ts = open_zarr_paths(nwp_like_zarr2_paths, concat_dim="init_time_utc") - # Check tensorstore version returns same results as dask version - kwargs = {"concat_dim": "init_time_utc", "combine": "nested", "engine": "zarr"} - ds_dask = xr.open_mfdataset(nwp_like_zarr2_paths, **kwargs) - assert ds_ts.compute().equals(ds_dask.compute()) - # Check function can open zarr3 - ds_ts = open_zarr_paths(nwp_like_zarr3_paths, concat_dim="init_time_utc") - # Check tensorstore version returns same results as dask version - ds_dask = xr.open_mfdataset(nwp_like_zarr3_paths, **kwargs) - assert ds_ts.compute().equals(ds_dask.compute()) +def test_concat_matches_xr_concat(tensorstore_datasets, zarr_paths): + ds_ts = concat_tensorstore(tensorstore_datasets, concat_dim="init_time_utc") + + ds_numpy = xr.concat( + [xr.open_zarr(path, chunks=None) for path in zarr_paths[2]], + dim="init_time_utc", + ) + assert ds_ts.compute().equals(ds_numpy) + + +def test_concat_stays_tensorstore_backed(tensorstore_datasets): + """The concatenated data variables should still be lazy tensorstores""" + ds_ts = concat_tensorstore(tensorstore_datasets, concat_dim="init_time_utc") + + # Raises TypeError if the variable has been materialised + store = _extract_tensorstore(ds_ts["ECMWF_UK"]) + assert store.shape == ds_ts["ECMWF_UK"].shape + + +def test_concat_requires_multiple_datasets(tensorstore_datasets): + with pytest.raises(ValueError, match="need at least two datasets"): + concat_tensorstore(tensorstore_datasets[:1], concat_dim="init_time_utc") + + +def test_concat_requires_valid_concat_dim(tensorstore_datasets): + with pytest.raises(ValueError, match="'not_a_dim' is not a dimension"): + concat_tensorstore(tensorstore_datasets, concat_dim="not_a_dim") + + +def test_concat_rejects_mismatched_datasets(tensorstore_datasets): + ds_1, ds_2 = tensorstore_datasets + with pytest.raises(ValueError, match="data_vars"): + concat_tensorstore( + [ds_1, ds_2.rename({"ECMWF_UK": "other"})], + concat_dim="init_time_utc", + ) + + +def test_concat_rejects_non_tensorstore_datasets(zarr_paths): + datasets = [xr.open_zarr(path) for path in zarr_paths[2]] + with pytest.raises(TypeError, match="expected TensorStore"): + concat_tensorstore(datasets, concat_dim="init_time_utc") diff --git a/tests/datasets/pvnet/test_materialise.py b/tests/datasets/pvnet/test_materialise.py index 9c4cc7b5..56cae83b 100644 --- a/tests/datasets/pvnet/test_materialise.py +++ b/tests/datasets/pvnet/test_materialise.py @@ -2,7 +2,7 @@ import numpy as np import xarray as xr -from ocf_data_sampler.common.xr_tensorstore import _open_single_zarr +from ocf_data_sampler.common.xr_tensorstore import open_zarr_paths from ocf_data_sampler.datasets.pvnet.materialise import load, load_data_dict @@ -30,7 +30,7 @@ def test_load_data_dict(tmp_path): da_dask.to_dataset(name="dummy_array").to_zarr(tmp_path) # Re-open with tensorstore - da_ts = _open_single_zarr(str(tmp_path)).dummy_array + da_ts = open_zarr_paths(str(tmp_path)).dummy_array # Create a nested dictionary with tensorstore arrays lazy_data_dict = { diff --git a/uv.lock b/uv.lock index 305cf62f..3ba661cd 100644 --- a/uv.lock +++ b/uv.lock @@ -840,7 +840,6 @@ wheels = [ name = "ocf-data-sampler" source = { editable = "." } dependencies = [ - { name = "dask" }, { name = "numcodecs" }, { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, @@ -858,6 +857,7 @@ dependencies = [ [package.dev-dependencies] dev = [ + { name = "dask" }, { name = "pvlib" }, { name = "pytest" }, { name = "python-lsp-ruff" }, @@ -868,7 +868,6 @@ dev = [ [package.metadata] requires-dist = [ - { name = "dask" }, { name = "numcodecs" }, { name = "numpy" }, { name = "pandas" }, @@ -878,12 +877,13 @@ requires-dist = [ { name = "pyresample" }, { name = "torch" }, { name = "xarray" }, - { name = "xarray-tensorstore", specifier = "==0.1.5" }, + { name = "xarray-tensorstore", specifier = "==0.3.0" }, { name = "zarr", specifier = ">=3" }, ] [package.metadata.requires-dev] dev = [ + { name = "dask" }, { name = "pvlib" }, { name = "pytest" }, { name = "python-lsp-ruff" }, @@ -1823,7 +1823,7 @@ wheels = [ [[package]] name = "xarray-tensorstore" -version = "0.1.5" +version = "0.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, @@ -1833,9 +1833,9 @@ dependencies = [ { name = "zarr", version = "3.1.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "zarr", version = "3.2.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/62/a2/6c1dbf9562b59ee5520118d56e169e58a1130685006baf573dfda01d2ae3/xarray-tensorstore-0.1.5.tar.gz", hash = "sha256:a8e17c08df6b32875e50f7b9278a64bbd9156b660d2245153abba4b2a9b30e9c", size = 9258, upload-time = "2024-11-08T23:12:45.023Z" } +sdist = { url = "https://files.pythonhosted.org/packages/58/64/64d401494b1cd8943960b1f9197ab4605433a4b648187881f36140c0b8d2/xarray_tensorstore-0.3.0.tar.gz", hash = "sha256:3f01a8182b2d09b6e9a365b206a648fd0bbbb0dd6b676b3b86864f8e52d5495e", size = 10560, upload-time = "2025-10-31T21:35:11.586Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3e/79/76a418041aed181dc950ac9c2d8843763f8ae6fcbe13f9eb1d09e36be749/xarray_tensorstore-0.1.5-py3-none-any.whl", hash = "sha256:e641714cb05a4540da04b5f1c259e9617a17e7855c29aeb7d72af5de57ab648d", size = 9375, upload-time = "2024-11-08T23:12:43.27Z" }, + { url = "https://files.pythonhosted.org/packages/08/81/5f57b39b1661ddf9a6e1403ca7289c1a6dab6028f8623c484cb693489579/xarray_tensorstore-0.3.0-py3-none-any.whl", hash = "sha256:f8415bda9f4dd52857c62ee280a6d19c7e5c596bdc8be7818635ea4bc38d6a97", size = 10552, upload-time = "2025-10-31T21:35:09.976Z" }, ] [[package]]