From 773e02393f72d84fb44a629aa074ab4036d0cc45 Mon Sep 17 00:00:00 2001 From: Nicola Masotti Date: Mon, 26 Oct 2026 18:29:33 +0100 Subject: [PATCH 1/2] Pad a request that comes without the first timestamp before applying the read selection --- tests/test_20_client_cdsapi.py | 131 +++++++++++++++++++++++++++++++++ xarray_ecmwf/client_cdsapi.py | 25 +++---- 2 files changed, 140 insertions(+), 16 deletions(-) create mode 100644 tests/test_20_client_cdsapi.py diff --git a/tests/test_20_client_cdsapi.py b/tests/test_20_client_cdsapi.py new file mode 100644 index 0000000..1e46414 --- /dev/null +++ b/tests/test_20_client_cdsapi.py @@ -0,0 +1,131 @@ +import hashlib +import json +from pathlib import Path +from typing import Any + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +# writing the schema sample with zarr 3 needs xarray >= 2025.6 +pytest.importorskip("xarray", minversion="2025.6") + +# CDS can return fewer time steps than requested for the first request of a +# dataset: ERA5-Land has no 1950-01-01 00:00. xarray-ecmwf must return NaN for +# the missing step and the right value for every other step. +DATE_REQUEST = { + "dataset": "dummy-dataset", + "variable": ["dummy_variable"], + "date": ["1950-01-01/1950-01-09"], + "time": ["00:00"], +} +YMD_REQUEST = { + "dataset": "dummy-dataset", + "variable": ["dummy_variable"], + "year": ["1950"], + "month": ["01", "02", "03"], + "day": ["01", "02"], + "time": ["00:00"], +} +# the first day of the fake dataset, like ERA5-Land +DATASET_START = pd.Timestamp("1950-01-02") +# the value of the first requested day is missing +EXPECTED = np.array([np.nan, 1, 2, 3, 4, 5, 6, 7, 8], dtype="float32") + + +# the request itself is the result, downloaded as JSON +class MissingFirstStepClient: + def __init__(self, client_kwargs: dict[str, Any]) -> None: + pass + + def submit_and_wait_on_result(self, request: dict[str, Any]) -> Any: + return request + + def get_filename(self, result: Any) -> str: + return hashlib.md5(str(result).encode("utf-8")).hexdigest() + ".json" + + def download(self, result: Any, target: str | None = None) -> str: + assert target is not None + Path(target).write_text(json.dumps(result)) + return target + + +def as_list(value: str | list[str]) -> list[str]: + return [value] if isinstance(value, str) else value + + +# one field per day, the value is the days since 1950-01-01 +def open_missing_first_step(path: str, **kwargs: Any) -> xr.Dataset: + request = json.loads(Path(path).read_text()) + if "date" in request: + start, stop = as_list(request["date"])[0].split("/") + time = pd.date_range(start, stop, freq="D") + else: + time = pd.DatetimeIndex( + [ + f"{year}-{month}-{day}" + for year in as_list(request["year"]) + for month in as_list(request["month"]) + for day in as_list(request["day"]) + ] + ) + time = time[time >= DATASET_START] + days = (time - pd.Timestamp("1950-01-01")).days.to_numpy(dtype="float32") + values = np.broadcast_to(days[:, None, None], (len(time), 2, 3)) + return xr.Dataset( + {"dummy": (("time", "latitude", "longitude"), values)}, + coords={"time": time, "latitude": [10.0, 0.0], "longitude": [0.0, 1.0, 2.0]}, + ) + + +def open_dummy_dataset( + request: dict[str, Any], request_chunks: dict[str, int], cache_folder: Path +) -> xr.Dataset: + return xr.open_dataset( + request, # type: ignore + engine="ecmwf", + request_chunks=request_chunks, + request_client_class=MissingFirstStepClient, + open_dataset=open_missing_first_step, + cache_kwargs={"cache_folder": str(cache_folder)}, + chunks={}, + ) + + +@pytest.mark.parametrize( + "time_key", + [ + slice(0, 3), # the whole first chunk + slice(1, 3), # first chunk, starting after the missing step + slice(0, 2), # first chunk, stopping before its end + 2, # a single step + ], +) +def test_first_chunk_read_returns_requested_steps( + tmp_path: Path, time_key: int | slice +) -> None: + ds = open_dummy_dataset(DATE_REQUEST, {"day": 3}, tmp_path) + + res = ds["dummy"].isel(time=time_key, latitude=0, longitude=0).values + + assert np.array_equal(res, EXPECTED[time_key], equal_nan=True) + + +def test_shifted_read_returns_next_step_values(tmp_path: Path) -> None: + ds = open_dummy_dataset(DATE_REQUEST, {"day": 3}, tmp_path) + # dask turns the lazy shift into a read of the first chunk starting at index 1 + da = ds["dummy"].shift(time=-1).isel(time=slice(None, -1)) + + res = da.isel(latitude=0, longitude=0).values + + assert np.array_equal(res, EXPECTED[1:], equal_nan=True) + + +def test_month_chunked_read_returns_requested_steps(tmp_path: Path) -> None: + # month chunks have a tuple of chunk sizes + ds = open_dummy_dataset(YMD_REQUEST, {"month": 1}, tmp_path) + + res = ds["dummy"].isel(latitude=0, longitude=0).values + + assert np.array_equal(res, [np.nan, 1, 31, 32, 59, 60], equal_nan=True) diff --git a/xarray_ecmwf/client_cdsapi.py b/xarray_ecmwf/client_cdsapi.py index aa50055..cc28bf3 100644 --- a/xarray_ecmwf/client_cdsapi.py +++ b/xarray_ecmwf/client_cdsapi.py @@ -270,22 +270,15 @@ def get_chunk_values( da = da.expand_dims(dims, axis=axis) - out = da.isel(selection).values - # horrible workaround for the crazy CDS / MARS convention to return # a short request at the start of a dataset (at least on ERA5 and ERA5 Land) if self.time_dim in indices and indices[self.time_dim] == 0: - if isinstance(self.chunks[self.time_dim], int): - time_chunk = self.chunks[self.time_dim] - else: - time_chunk = self.chunks[self.time_dim][0] # type: ignore - if da.coords[self.time_dim].size < time_chunk: - expected_time_dim_size = self.chunks[self.time_dim] - assert isinstance(expected_time_dim_size, int) - fixed = np.empty((expected_time_dim_size,) + out.shape[1:]) - offset = fixed.shape[0] - out.shape[0] - fixed[offset:] = out - fixed[:offset] = np.nan - out = fixed - - return out + time_chunk = self.chunks[self.time_dim] + if not isinstance(time_chunk, int): + time_chunk = time_chunk[0] + missing = time_chunk - da.sizes[self.time_dim] + if missing > 0: + # pad before selecting, as the selection refers to the full chunk + da = da.pad({self.time_dim: (missing, 0)}) + + return da.isel(selection).values From 03a53c8507ff9ace4bef00b780bc4aea2b2f1bc3 Mon Sep 17 00:00:00 2001 From: Nicola Masotti Date: Tue, 27 Oct 2026 11:02:29 +0100 Subject: [PATCH 2/2] Add more hours of day to tests --- tests/test_20_client_cdsapi.py | 49 +++++++++++++++++++--------------- 1 file changed, 27 insertions(+), 22 deletions(-) diff --git a/tests/test_20_client_cdsapi.py b/tests/test_20_client_cdsapi.py index 1e46414..5f818da 100644 --- a/tests/test_20_client_cdsapi.py +++ b/tests/test_20_client_cdsapi.py @@ -12,26 +12,26 @@ pytest.importorskip("xarray", minversion="2025.6") # CDS can return fewer time steps than requested for the first request of a -# dataset: ERA5-Land has no 1950-01-01 00:00. xarray-ecmwf must return NaN for -# the missing step and the right value for every other step. +# dataset: ERA5-Land hourly has no 1950-01-01 00:00. xarray-ecmwf must return +# NaN for the missing step and the right value for every other step. DATE_REQUEST = { "dataset": "dummy-dataset", "variable": ["dummy_variable"], - "date": ["1950-01-01/1950-01-09"], - "time": ["00:00"], + "date": ["1950-01-01/1950-01-04"], + "time": ["00:00", "01:00", "02:00"], } YMD_REQUEST = { "dataset": "dummy-dataset", "variable": ["dummy_variable"], "year": ["1950"], - "month": ["01", "02", "03"], + "month": ["01", "02"], "day": ["01", "02"], - "time": ["00:00"], + "time": ["00:00", "01:00", "02:00"], } -# the first day of the fake dataset, like ERA5-Land -DATASET_START = pd.Timestamp("1950-01-02") -# the value of the first requested day is missing -EXPECTED = np.array([np.nan, 1, 2, 3, 4, 5, 6, 7, 8], dtype="float32") +# the first step of the fake dataset, like ERA5-Land +DATASET_START = pd.Timestamp("1950-01-01 01:00") +# the value of the first requested hour is missing +EXPECTED = np.array([np.nan, 1, 2, 24, 25, 26, 48, 49, 50, 72, 73, 74]) # the request itself is the result, downloaded as JSON @@ -55,14 +55,14 @@ def as_list(value: str | list[str]) -> list[str]: return [value] if isinstance(value, str) else value -# one field per day, the value is the days since 1950-01-01 +# one field per hour, the value is the hours since 1950-01-01 00:00 def open_missing_first_step(path: str, **kwargs: Any) -> xr.Dataset: request = json.loads(Path(path).read_text()) if "date" in request: start, stop = as_list(request["date"])[0].split("/") - time = pd.date_range(start, stop, freq="D") + days = pd.date_range(start, stop, freq="D") else: - time = pd.DatetimeIndex( + days = pd.DatetimeIndex( [ f"{year}-{month}-{day}" for year in as_list(request["year"]) @@ -70,9 +70,12 @@ def open_missing_first_step(path: str, **kwargs: Any) -> xr.Dataset: for day in as_list(request["day"]) ] ) + time = pd.DatetimeIndex( + [f"{day:%Y-%m-%d} {hour}" for day in days for hour in request["time"]] + ) time = time[time >= DATASET_START] - days = (time - pd.Timestamp("1950-01-01")).days.to_numpy(dtype="float32") - values = np.broadcast_to(days[:, None, None], (len(time), 2, 3)) + hours = (time - pd.Timestamp("1950-01-01")) / pd.Timedelta("1h") + values = np.broadcast_to(hours.to_numpy()[:, None, None], (len(time), 2, 3)) return xr.Dataset( {"dummy": (("time", "latitude", "longitude"), values)}, coords={"time": time, "latitude": [10.0, 0.0], "longitude": [0.0, 1.0, 2.0]}, @@ -96,16 +99,16 @@ def open_dummy_dataset( @pytest.mark.parametrize( "time_key", [ - slice(0, 3), # the whole first chunk - slice(1, 3), # first chunk, starting after the missing step - slice(0, 2), # first chunk, stopping before its end - 2, # a single step + slice(0, 6), # the whole first chunk + slice(1, 6), # first chunk, starting after the missing step + slice(0, 4), # first chunk, stopping before its end + 4, # a single step ], ) def test_first_chunk_read_returns_requested_steps( tmp_path: Path, time_key: int | slice ) -> None: - ds = open_dummy_dataset(DATE_REQUEST, {"day": 3}, tmp_path) + ds = open_dummy_dataset(DATE_REQUEST, {"day": 2}, tmp_path) res = ds["dummy"].isel(time=time_key, latitude=0, longitude=0).values @@ -113,7 +116,7 @@ def test_first_chunk_read_returns_requested_steps( def test_shifted_read_returns_next_step_values(tmp_path: Path) -> None: - ds = open_dummy_dataset(DATE_REQUEST, {"day": 3}, tmp_path) + ds = open_dummy_dataset(DATE_REQUEST, {"day": 2}, tmp_path) # dask turns the lazy shift into a read of the first chunk starting at index 1 da = ds["dummy"].shift(time=-1).isel(time=slice(None, -1)) @@ -128,4 +131,6 @@ def test_month_chunked_read_returns_requested_steps(tmp_path: Path) -> None: res = ds["dummy"].isel(latitude=0, longitude=0).values - assert np.array_equal(res, [np.nan, 1, 31, 32, 59, 60], equal_nan=True) + # 1950-02-01 is 744 hours after 1950-01-01 + expected = [np.nan, 1, 2, 24, 25, 26, 744, 745, 746, 768, 769, 770] + assert np.array_equal(res, expected, equal_nan=True)