From ed7d929a40e183849de29497af568aba85564f2f Mon Sep 17 00:00:00 2001
From: James Fulton <djamesfulton@yahoo.co.uk>
Date: Fri, 31 Jul 2026 08:49:28 +0000
Subject: [PATCH 1/6] Upgrade xarray-tensor and add tenorstore concat function

---
 README.md                                     |   2 +-
 pyproject.toml                                |   3 +-
 src/ocf_data_sampler/common/xr_tensorstore.py | 275 ++++++------------
 tests/common/test_xr_tensorstore.py           | 115 +++++---
 tests/datasets/pvnet/test_materialise.py      |   4 +-
 uv.lock                                       |  69 +----
 6 files changed, 179 insertions(+), 289 deletions(-)

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..aa902927 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",
 ]
 
diff --git a/src/ocf_data_sampler/common/xr_tensorstore.py b/src/ocf_data_sampler/common/xr_tensorstore.py
index aafdbc35..7fe01475 100644
--- a/src/ocf_data_sampler/common/xr_tensorstore.py
+++ b/src/ocf_data_sampler/common/xr_tensorstore.py
@@ -1,96 +1,107 @@
-"""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,
-)
+import xarray_tensorstore as xrt
 
-logger = logging.getLogger(__name__)
 
 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 _tensorstore_of(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)} != {sorted(first.coords)}")
+        if set(ds.data_vars) != set(first.data_vars):
+            raise ValueError(
+                f"dataset {i}: data_vars {sorted(ds.data_vars)} != {sorted(first.data_vars)}"
+            )
+
+        # 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.
+    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(
+                [_tensorstore_of(ds[name]) for ds in datasets],
+                axis=da.dims.index(concat_dim),
+            )
+            out[name] = xr.Variable(da.dims, xrt._TensorStoreAdapter(store), attrs=da.attrs)
+        else:
+            out[name] = da.variable  # attrs travel with the Variable; xarray copies on assign
+    return out
 
 
 def open_zarr_paths(zarr_path: ZarrSource, concat_dim: str | None = None) -> xr.Dataset:
@@ -103,7 +114,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 +123,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..982b6de7 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 (
+    _tensorstore_of,
+    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 = _tensorstore_of(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..4c803235 100644
--- a/uv.lock
+++ b/uv.lock
@@ -145,15 +145,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/fb/e2/79c688af8b210d232694e31e59da9f6ec747bae31c3f5946e4e9b98860d5/click-8.4.2-py3-none-any.whl", hash = "sha256:e6f9f66136c816745b9d65817da91d61d957fb16e02e4dcd0552553c5a197b76", size = 119243, upload-time = "2026-06-24T17:45:13.73Z" },
 ]
 
-[[package]]
-name = "cloudpickle"
-version = "3.1.2"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" },
-]
-
 [[package]]
 name = "colorama"
 version = "0.4.6"
@@ -245,25 +236,6 @@ nvtx = [
     { name = "nvidia-nvtx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
 ]
 
-[[package]]
-name = "dask"
-version = "2026.7.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "click" },
-    { name = "cloudpickle" },
-    { name = "fsspec" },
-    { name = "importlib-metadata", marker = "python_full_version < '3.12'" },
-    { name = "packaging" },
-    { name = "partd" },
-    { name = "pyyaml" },
-    { name = "toolz" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/d6/39/cbd21c9133d02b4e60899ed466ad5e876553ea68ebffee6209e7dcafc8d4/dask-2026.7.1.tar.gz", hash = "sha256:5727484427665f051e86bf87d021a64d6411141cdc8a20bfe3c1ad2968cc06b7", size = 11548794, upload-time = "2026-07-14T01:06:22.46Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/7e/5f/7c22733da92b3a6cc4dddcaa8731089d213c2790bbc997e51c429a4e8f8b/dask-2026.7.1-py3-none-any.whl", hash = "sha256:985ffd6c5e9d7979ede515e84ae8d39b647d6aa64f77600f15714ff65f578fe6", size = 1496882, upload-time = "2026-07-14T01:06:20.341Z" },
-]
-
 [[package]]
 name = "docstring-to-markdown"
 version = "0.17"
@@ -422,15 +394,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" },
 ]
 
-[[package]]
-name = "locket"
-version = "1.0.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/2f/83/97b29fe05cb6ae28d2dbd30b81e2e402a3eed5f460c26e9eaa5895ceacf5/locket-1.0.0.tar.gz", hash = "sha256:5c0d4c052a8bbbf750e056a8e65ccd309086f4f0f18a2eac306a8dfa4112a632", size = 4350, upload-time = "2022-04-20T22:04:44.312Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/db/bc/83e112abc66cd466c6b83f99118035867cecd41802f8d044638aa78a106e/locket-1.0.0-py2.py3-none-any.whl", hash = "sha256:b6c819a722f7b6bd955b80781788e4a66a55628b858d347536b7e81325a3a5e3", size = 4398, upload-time = "2022-04-20T22:04:42.23Z" },
-]
-
 [[package]]
 name = "lsprotocol"
 version = "2025.0.0"
@@ -840,7 +803,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'" },
@@ -868,7 +830,6 @@ dev = [
 
 [package.metadata]
 requires-dist = [
-    { name = "dask" },
     { name = "numcodecs" },
     { name = "numpy" },
     { name = "pandas" },
@@ -878,7 +839,7 @@ 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" },
 ]
 
@@ -955,19 +916,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/99/5d/8268b644392ee874ee82a635cd0df1773de230bde356c38de28e298392cc/parso-0.8.7-py2.py3-none-any.whl", hash = "sha256:a8926eb2a1b915486941fdbd31e86a4baf88fe8c210f25f2f35ecec5b574ca1c", size = 107025, upload-time = "2026-05-01T23:12:58.867Z" },
 ]
 
-[[package]]
-name = "partd"
-version = "1.4.2"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "locket" },
-    { name = "toolz" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/b2/3a/3f06f34820a31257ddcabdfafc2672c5816be79c7e353b02c1f318daa7d4/partd-1.4.2.tar.gz", hash = "sha256:d022c33afbdc8405c226621b015e8067888173d85f7f5ecebb3cafed9a20f02c", size = 21029, upload-time = "2024-05-06T19:51:41.945Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/71/e7/40fb618334dcdf7c5a316c0e7343c5cd82d3d866edc100d98e29bc945ecd/partd-1.4.2-py3-none-any.whl", hash = "sha256:978e4ac767ec4ba5b86c6eaa52e5a2a3bc748a2ca839e8cc798f1cc6ce6efb0f", size = 18905, upload-time = "2024-05-06T19:51:39.271Z" },
-]
-
 [[package]]
 name = "pathspec"
 version = "1.1.1"
@@ -1624,15 +1572,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/a5/12/97d8ad183e3130e168f2feb860edd68f1b72e57f29268d980f3b70e34cd0/tensorstore-0.1.84-cp313-cp313-win_amd64.whl", hash = "sha256:fe9bf1c7fef69884a91222179550f9b5ba6c1454f9534429221824d9b15c00ec", size = 13398858, upload-time = "2026-05-16T06:17:33.658Z" },
 ]
 
-[[package]]
-name = "toolz"
-version = "1.1.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/11/d6/114b492226588d6ff54579d95847662fc69196bdeec318eb45393b24c192/toolz-1.1.0.tar.gz", hash = "sha256:27a5c770d068c110d9ed9323f24f1543e83b2f300a687b7891c1a6d56b697b5b", size = 52613, upload-time = "2025-10-17T04:03:21.661Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/fb/12/5911ae3eeec47800503a238d971e51722ccea5feb8569b735184d5fcdbc0/toolz-1.1.0-py3-none-any.whl", hash = "sha256:15ccc861ac51c53696de0a5d6d4607f99c210739caf987b5d2054f3efed429d8", size = 58093, upload-time = "2025-10-17T04:03:20.435Z" },
-]
-
 [[package]]
 name = "torch"
 version = "2.13.0"
@@ -1823,7 +1762,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 +1772,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]]

From ef5c0552daba38a8ee817e982fb2f3ce63c52b17 Mon Sep 17 00:00:00 2001
From: James Fulton <djamesfulton@yahoo.co.uk>
Date: Fri, 31 Jul 2026 10:18:36 +0000
Subject: [PATCH 2/6] lint

---
 src/ocf_data_sampler/common/xr_tensorstore.py | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/src/ocf_data_sampler/common/xr_tensorstore.py b/src/ocf_data_sampler/common/xr_tensorstore.py
index 7fe01475..4d546011 100644
--- a/src/ocf_data_sampler/common/xr_tensorstore.py
+++ b/src/ocf_data_sampler/common/xr_tensorstore.py
@@ -13,7 +13,6 @@
 import xarray as xr
 import xarray_tensorstore as xrt
 
-
 ZarrSource: TypeAlias = str | list[str] | tuple[str, ...]
 
 
@@ -44,7 +43,7 @@ def _validate(datasets: Sequence[xr.Dataset], concat_dim: str) -> None:
         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:
@@ -54,7 +53,7 @@ def _validate(datasets: Sequence[xr.Dataset], concat_dim: str) -> None:
                 )
 
         # 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, 
+        # 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]):
@@ -66,7 +65,7 @@ def _validate(datasets: Sequence[xr.Dataset], concat_dim: str) -> None:
 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. 
+    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:
@@ -79,11 +78,11 @@ def concat_tensorstore(datasets: Sequence[xr.Dataset], concat_dim: str) -> xr.Da
     _validate(datasets, concat_dim)
     first = datasets[0]
 
-    # Create a new shell dataset which contains only the concatenated coords. We will handle the 
+    # 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 
+    # - 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.
     out = xr.concat(
         [ds.drop_vars(first.data_vars) for ds in datasets],

From 5aebec207c4b6ec3494af18a42f74b267465a366 Mon Sep 17 00:00:00 2001
From: James Fulton <djamesfulton@yahoo.co.uk>
Date: Fri, 31 Jul 2026 10:39:07 +0000
Subject: [PATCH 3/6] Fix type checking

---
 src/ocf_data_sampler/common/xr_tensorstore.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/ocf_data_sampler/common/xr_tensorstore.py b/src/ocf_data_sampler/common/xr_tensorstore.py
index 4d546011..b4348fe4 100644
--- a/src/ocf_data_sampler/common/xr_tensorstore.py
+++ b/src/ocf_data_sampler/common/xr_tensorstore.py
@@ -33,10 +33,14 @@ def _validate(datasets: Sequence[xr.Dataset], concat_dim: str) -> None:
 
         # 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)} != {sorted(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)} != {sorted(first.data_vars)}"
+                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

From 4d4968f58fe9bd0d523cd2622a5ddc4b3c826d03 Mon Sep 17 00:00:00 2001
From: James Fulton <djamesfulton@yahoo.co.uk>
Date: Fri, 31 Jul 2026 10:49:55 +0000
Subject: [PATCH 4/6] add dask as dev dependency

---
 pyproject.toml |  1 +
 uv.lock        | 61 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index aa902927..b744e2d7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,6 +39,7 @@ dev = [
     # Testing
     "pytest",
     "pvlib",
+    "dask",
     "ruff >= 0.9.2",
     "ty",
     # LSP Support
diff --git a/uv.lock b/uv.lock
index 4c803235..3ba661cd 100644
--- a/uv.lock
+++ b/uv.lock
@@ -145,6 +145,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/fb/e2/79c688af8b210d232694e31e59da9f6ec747bae31c3f5946e4e9b98860d5/click-8.4.2-py3-none-any.whl", hash = "sha256:e6f9f66136c816745b9d65817da91d61d957fb16e02e4dcd0552553c5a197b76", size = 119243, upload-time = "2026-06-24T17:45:13.73Z" },
 ]
 
+[[package]]
+name = "cloudpickle"
+version = "3.1.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" },
+]
+
 [[package]]
 name = "colorama"
 version = "0.4.6"
@@ -236,6 +245,25 @@ nvtx = [
     { name = "nvidia-nvtx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
 ]
 
+[[package]]
+name = "dask"
+version = "2026.7.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "click" },
+    { name = "cloudpickle" },
+    { name = "fsspec" },
+    { name = "importlib-metadata", marker = "python_full_version < '3.12'" },
+    { name = "packaging" },
+    { name = "partd" },
+    { name = "pyyaml" },
+    { name = "toolz" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/d6/39/cbd21c9133d02b4e60899ed466ad5e876553ea68ebffee6209e7dcafc8d4/dask-2026.7.1.tar.gz", hash = "sha256:5727484427665f051e86bf87d021a64d6411141cdc8a20bfe3c1ad2968cc06b7", size = 11548794, upload-time = "2026-07-14T01:06:22.46Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/7e/5f/7c22733da92b3a6cc4dddcaa8731089d213c2790bbc997e51c429a4e8f8b/dask-2026.7.1-py3-none-any.whl", hash = "sha256:985ffd6c5e9d7979ede515e84ae8d39b647d6aa64f77600f15714ff65f578fe6", size = 1496882, upload-time = "2026-07-14T01:06:20.341Z" },
+]
+
 [[package]]
 name = "docstring-to-markdown"
 version = "0.17"
@@ -394,6 +422,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" },
 ]
 
+[[package]]
+name = "locket"
+version = "1.0.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/2f/83/97b29fe05cb6ae28d2dbd30b81e2e402a3eed5f460c26e9eaa5895ceacf5/locket-1.0.0.tar.gz", hash = "sha256:5c0d4c052a8bbbf750e056a8e65ccd309086f4f0f18a2eac306a8dfa4112a632", size = 4350, upload-time = "2022-04-20T22:04:44.312Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/db/bc/83e112abc66cd466c6b83f99118035867cecd41802f8d044638aa78a106e/locket-1.0.0-py2.py3-none-any.whl", hash = "sha256:b6c819a722f7b6bd955b80781788e4a66a55628b858d347536b7e81325a3a5e3", size = 4398, upload-time = "2022-04-20T22:04:42.23Z" },
+]
+
 [[package]]
 name = "lsprotocol"
 version = "2025.0.0"
@@ -820,6 +857,7 @@ dependencies = [
 
 [package.dev-dependencies]
 dev = [
+    { name = "dask" },
     { name = "pvlib" },
     { name = "pytest" },
     { name = "python-lsp-ruff" },
@@ -845,6 +883,7 @@ requires-dist = [
 
 [package.metadata.requires-dev]
 dev = [
+    { name = "dask" },
     { name = "pvlib" },
     { name = "pytest" },
     { name = "python-lsp-ruff" },
@@ -916,6 +955,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/99/5d/8268b644392ee874ee82a635cd0df1773de230bde356c38de28e298392cc/parso-0.8.7-py2.py3-none-any.whl", hash = "sha256:a8926eb2a1b915486941fdbd31e86a4baf88fe8c210f25f2f35ecec5b574ca1c", size = 107025, upload-time = "2026-05-01T23:12:58.867Z" },
 ]
 
+[[package]]
+name = "partd"
+version = "1.4.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "locket" },
+    { name = "toolz" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/b2/3a/3f06f34820a31257ddcabdfafc2672c5816be79c7e353b02c1f318daa7d4/partd-1.4.2.tar.gz", hash = "sha256:d022c33afbdc8405c226621b015e8067888173d85f7f5ecebb3cafed9a20f02c", size = 21029, upload-time = "2024-05-06T19:51:41.945Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/71/e7/40fb618334dcdf7c5a316c0e7343c5cd82d3d866edc100d98e29bc945ecd/partd-1.4.2-py3-none-any.whl", hash = "sha256:978e4ac767ec4ba5b86c6eaa52e5a2a3bc748a2ca839e8cc798f1cc6ce6efb0f", size = 18905, upload-time = "2024-05-06T19:51:39.271Z" },
+]
+
 [[package]]
 name = "pathspec"
 version = "1.1.1"
@@ -1572,6 +1624,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/a5/12/97d8ad183e3130e168f2feb860edd68f1b72e57f29268d980f3b70e34cd0/tensorstore-0.1.84-cp313-cp313-win_amd64.whl", hash = "sha256:fe9bf1c7fef69884a91222179550f9b5ba6c1454f9534429221824d9b15c00ec", size = 13398858, upload-time = "2026-05-16T06:17:33.658Z" },
 ]
 
+[[package]]
+name = "toolz"
+version = "1.1.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/11/d6/114b492226588d6ff54579d95847662fc69196bdeec318eb45393b24c192/toolz-1.1.0.tar.gz", hash = "sha256:27a5c770d068c110d9ed9323f24f1543e83b2f300a687b7891c1a6d56b697b5b", size = 52613, upload-time = "2025-10-17T04:03:21.661Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/fb/12/5911ae3eeec47800503a238d971e51722ccea5feb8569b735184d5fcdbc0/toolz-1.1.0-py3-none-any.whl", hash = "sha256:15ccc861ac51c53696de0a5d6d4607f99c210739caf987b5d2054f3efed429d8", size = 58093, upload-time = "2025-10-17T04:03:20.435Z" },
+]
+
 [[package]]
 name = "torch"
 version = "2.13.0"

From f312560adae328ed402e9a5bee4ab8e3f8ffadc6 Mon Sep 17 00:00:00 2001
From: James Fulton <djamesfulton@yahoo.co.uk>
Date: Fri, 31 Jul 2026 11:09:49 +0000
Subject: [PATCH 5/6] rename

---
 src/ocf_data_sampler/common/xr_tensorstore.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/ocf_data_sampler/common/xr_tensorstore.py b/src/ocf_data_sampler/common/xr_tensorstore.py
index b4348fe4..8b815a57 100644
--- a/src/ocf_data_sampler/common/xr_tensorstore.py
+++ b/src/ocf_data_sampler/common/xr_tensorstore.py
@@ -88,7 +88,7 @@ def concat_tensorstore(datasets: Sequence[xr.Dataset], concat_dim: str) -> xr.Da
     #   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.
-    out = xr.concat(
+    ds_out = xr.concat(
         [ds.drop_vars(first.data_vars) for ds in datasets],
         dim=concat_dim,
         join="exact",
@@ -101,10 +101,10 @@ def concat_tensorstore(datasets: Sequence[xr.Dataset], concat_dim: str) -> xr.Da
                 [_tensorstore_of(ds[name]) for ds in datasets],
                 axis=da.dims.index(concat_dim),
             )
-            out[name] = xr.Variable(da.dims, xrt._TensorStoreAdapter(store), attrs=da.attrs)
+            ds_out[name] = xr.Variable(da.dims, xrt._TensorStoreAdapter(store), attrs=da.attrs)
         else:
-            out[name] = da.variable  # attrs travel with the Variable; xarray copies on assign
-    return out
+            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:

From 2986d2f7e763e29071a3b0826da089371f3e139a Mon Sep 17 00:00:00 2001
From: James Fulton <djamesfulton@yahoo.co.uk>
Date: Wed, 5 Aug 2026 09:57:10 +0000
Subject: [PATCH 6/6] rename function

---
 src/ocf_data_sampler/common/xr_tensorstore.py | 4 ++--
 tests/common/test_xr_tensorstore.py           | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/ocf_data_sampler/common/xr_tensorstore.py b/src/ocf_data_sampler/common/xr_tensorstore.py
index 8b815a57..b61ccdea 100644
--- a/src/ocf_data_sampler/common/xr_tensorstore.py
+++ b/src/ocf_data_sampler/common/xr_tensorstore.py
@@ -16,7 +16,7 @@
 ZarrSource: TypeAlias = str | list[str] | tuple[str, ...]
 
 
-def _tensorstore_of(da: xr.DataArray) -> ts.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):
@@ -98,7 +98,7 @@ def concat_tensorstore(datasets: Sequence[xr.Dataset], concat_dim: str) -> xr.Da
     for name, da in first.data_vars.items():
         if concat_dim in da.dims:
             store = ts.concat(
-                [_tensorstore_of(ds[name]) for ds in datasets],
+                [_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)
diff --git a/tests/common/test_xr_tensorstore.py b/tests/common/test_xr_tensorstore.py
index 982b6de7..02b6a671 100644
--- a/tests/common/test_xr_tensorstore.py
+++ b/tests/common/test_xr_tensorstore.py
@@ -3,7 +3,7 @@
 import xarray as xr
 
 from ocf_data_sampler.common.xr_tensorstore import (
-    _tensorstore_of,
+    _extract_tensorstore,
     concat_tensorstore,
     open_zarr_paths,
 )
@@ -81,7 +81,7 @@ def test_concat_stays_tensorstore_backed(tensorstore_datasets):
     ds_ts = concat_tensorstore(tensorstore_datasets, concat_dim="init_time_utc")
 
     # Raises TypeError if the variable has been materialised
-    store = _tensorstore_of(ds_ts["ECMWF_UK"])
+    store = _extract_tensorstore(ds_ts["ECMWF_UK"])
     assert store.shape == ds_ts["ECMWF_UK"].shape