Cut time coord for ymd chunking

README.md

@@ -12,6 +12,7 @@ xarray-ecmwf is a Python library and [Xarray](https://docs.xarray.dev) backend w
   - the [ECMWF Open data](https://www.ecmwf.int/en/forecasts/datasets/open-data) via [ecmwf-opendata](https://github.com/ecmwf/ecmwf-opendata): High resolution forecasts, ensemble forecast
 - allows lazy loading the data and well integrated with [Dask](https://www.dask.org) and [Dask.distributed](https://distributed.dask.org)
 - allows chunking the input request according to a configurable splitting strategy. Allowed strategies:
+  - by N years
   - by one month
   - by one day
 - supports requests returning a single GRIB file, via [cfgrib](https://github.com/ecmwf/cfgrib)
@@ -54,6 +55,25 @@ Attributes:
 
 ```
 
+## `force_end_date`
+
+Pass `force_end_date` in `request_chunker_kwargs`. The keyword uses underscores.
+
+For a request built from `year`, `month`, `day`, and `time`, the time coordinate is every combination of those lists. Timestamps after `force_end_date` are dropped. A timestamp equal to it stays. The dataset coordinate ends at the last remaining timestamp, also when only one chunk is left. With a monthly series on day `01` and `force_end_date` `1950-08-17`, the last time is `1950-08-01`.
+
+```python
+>>> ds = xr.open_dataset(
+...     request,
+...     engine="ecmwf",
+...     request_chunks={"year": 1},
+...     request_chunker_kwargs={"force_end_date": "1950-08-17"},
+... )
+```
+
+The option applies to requests that contain `year`. A request that uses `date` is left unchanged.
+
+The CDS request for a year still carries the original `month` list. The coordinate ends at the last timestamp that is not after `force_end_date`.
+
 ## License
 
 ```

tests/test_10_client_common.py

@@ -2,10 +2,11 @@
 import itertools
 from typing import Any
 
+import numpy as np
 import pandas as pd
 import pytest
 
-from xarray_ecmwf import client_common
+from xarray_ecmwf import client_cdsapi, client_common
 
 ALL_MONTHS = [f"{m:02}" for m in range(1, 13)]
 ALL_DAYS = [f"{d:02}" for d in range(1, 32)]
@@ -188,6 +189,141 @@ def test_build_chunk_ymd_year_requests_groups_years(
     assert time_chunk == tuple(len(group) for group in expected_year_groups)
 
 
+def _monthly_year_request() -> dict[str, Any]:
+    return {
+        "year": ["2024", "2025", "2026"],
+        "month": [f"{month:02}" for month in range(1, 13)],
+        "day": ["01"],
+        "time": ["00:00"],
+    }
+
+
+def test_force_end_date_trims_partial_last_year() -> None:
+    time, time_chunk, time_chunk_requests = client_common.build_time_chunk_requests(
+        _monthly_year_request(),
+        {"year": 1},
+        force_end_date="2026-07-01",
+    )
+    assert len(time) == 12 + 12 + 7
+    assert time[-1] == np.datetime64("2026-07-01T00:00", "ns")
+    assert time_chunk == (12, 12, 7)
+    assert [chunk_request["year"] for _, chunk_request in time_chunk_requests] == [
+        ["2024"],
+        ["2025"],
+        ["2026"],
+    ]
+
+
+def test_force_end_date_drops_years_after_the_cutoff() -> None:
+    time, time_chunk, time_chunk_requests = client_common.build_time_chunk_requests(
+        _monthly_year_request(),
+        {"year": 1},
+        force_end_date="2025-12-01",
+    )
+    assert len(time) == 24
+    assert time[-1] == np.datetime64("2025-12-01T00:00", "ns")
+    assert time_chunk == (12, 12)
+    assert [chunk_request["year"] for _, chunk_request in time_chunk_requests] == [
+        ["2024"],
+        ["2025"],
+    ]
+
+
+def test_force_end_date_shrinks_a_multi_year_chunk() -> None:
+    time, time_chunk, time_chunk_requests = client_common.build_time_chunk_requests(
+        _monthly_year_request(),
+        {"year": 2},
+        force_end_date="2025-06-01",
+    )
+    assert len(time) == 12 + 6
+    assert time[-1] == np.datetime64("2025-06-01T00:00", "ns")
+    assert time_chunk == (18,)
+    assert time_chunk_requests == [(0, {"year": ["2024", "2025"]})]
+
+
+def test_force_end_date_on_a_chunk_boundary_keeps_equal_day_chunks() -> None:
+    request = {
+        "year": ["2024"],
+        "month": ["01", "02"],
+        "day": ["01", "02"],
+        "time": ["00:00"],
+    }
+    time, time_chunk, time_chunk_requests = client_common.build_time_chunk_requests(
+        request,
+        {"day": 1},
+        force_end_date="2024-01-02",
+    )
+    assert list(time) == [
+        np.datetime64("2024-01-01T00:00", "ns"),
+        np.datetime64("2024-01-02T00:00", "ns"),
+    ]
+    assert time_chunk == 1
+    assert len(time_chunk_requests) == 2
+
+
+def test_force_end_date_shrinks_the_last_equal_sized_chunk() -> None:
+    request = {
+        "year": ["2024"],
+        "month": ["01"],
+        "day": ["01", "02"],
+        "time": ["00:00", "12:00"],
+    }
+    time, time_chunk, time_chunk_requests = client_common.build_time_chunk_requests(
+        request,
+        {"day": 1},
+        force_end_date="2024-01-02T00:00",
+    )
+    assert list(time) == [
+        np.datetime64("2024-01-01T00:00", "ns"),
+        np.datetime64("2024-01-01T12:00", "ns"),
+        np.datetime64("2024-01-02T00:00", "ns"),
+    ]
+    assert time_chunk == (2, 1)
+    assert [chunk_request["day"] for _, chunk_request in time_chunk_requests] == [
+        "01",
+        "02",
+    ]
+
+
+def test_force_end_date_past_the_last_timestamp_is_a_no_op() -> None:
+    full, full_chunk, full_requests = client_common.build_time_chunk_requests(
+        _monthly_year_request(), {"year": 1}
+    )
+    trimmed, trimmed_chunk, trimmed_requests = client_common.build_time_chunk_requests(
+        _monthly_year_request(), {"year": 1}, force_end_date="2027-01-01"
+    )
+    assert list(trimmed) == list(full)
+    assert trimmed_chunk == full_chunk
+    assert trimmed_requests == full_requests
+
+
+def test_force_end_date_stays_when_one_chunk_remains() -> None:
+    chunker = client_cdsapi.CdsapiRequestChunker(
+        request={
+            "year": ["1950", "1951"],
+            "month": [f"{month:02}" for month in range(1, 13)],
+            "day": ["01"],
+            "time": ["00:00"],
+        },
+        request_chunks={"year": 1},
+        time_dim="valid_time",
+        force_end_date="1950-08-17",
+    )
+    valid_time = chunker.compute_chunked_request_coords()["valid_time"]
+    assert valid_time.size == 8
+    assert valid_time.values[-1] == np.datetime64("1950-08-01T00:00", "ns")
+    assert chunker.first_chunk_request()["year"] == ["1950"]
+
+
+def test_force_end_date_before_the_first_timestamp_raises() -> None:
+    with pytest.raises(ValueError, match="before the first timestamp"):
+        client_common.build_time_chunk_requests(
+            _monthly_year_request(),
+            {"year": 1},
+            force_end_date="2020-01-01",
+        )
+
+
 def test_build_chunk_ymd_year_requests_invalid() -> None:
     request = {
         "year": ["2022"],

xarray_ecmwf/client_cdsapi.py

@@ -52,6 +52,7 @@ class CdsapiRequestChunker:
     merge_date_time: bool = True
     time_dim: str = "time"
     time_sep: str = "/"
+    force_end_date: str | None = None
 
     def get_request_dimensions(self) -> dict[str, list[Any]]:
         request_dimensions: dict[str, list[Any]] = {}
@@ -116,9 +117,12 @@ def compute_chunked_request_coords(self) -> dict[str, Any]:
                     time_chunk,
                     time_chunk_requests,
                 ) = client_common.build_time_chunk_requests(
-                    self.request | override_time, self.request_chunks, self.time_sep
+                    self.request | override_time,
+                    self.request_chunks,
+                    self.time_sep,
+                    force_end_date=self.force_end_date,
                 )
-                if len(time_chunk_requests) > 1:
+                if self.force_end_date is not None or len(time_chunk_requests) > 1:
                     self.chunks[self.time_dim] = time_chunk
                     self.chunk_requests[self.time_dim] = time_chunk_requests
                     self.chunked_coords[self.time_dim] = xr.IndexVariable(  # type: ignore

xarray_ecmwf/client_common.py

@@ -1,10 +1,13 @@
 import calendar
+import logging
 from typing import Any, ContextManager, Protocol
 
 import numpy as np
 import pandas as pd
 import xarray as xr
 
+LOGGER = logging.getLogger(__name__)
+
 
 class RequestClientProtocol(Protocol):
     def __init__(self, client_kwargs: dict[str, Any]) -> None: ...
@@ -263,8 +266,62 @@ def build_chunk_ymd_year_requests(
     return np.array(datetimes), tuple(chunks), chunk_requests
 
 
+def trim_ymd_time_to_force_end_date(
+    time: np.typing.NDArray[np.datetime64],
+    time_chunk: int | tuple[int, ...],
+    time_chunk_requests: list[tuple[int, dict[str, Any]]],
+    force_end_date: str,
+) -> tuple[
+    np.typing.NDArray[np.datetime64],
+    int | tuple[int, ...],
+    list[tuple[int, dict[str, Any]]],
+]:
+    """Cut a year/month/day time coordinate at ``force_end_date``.
+
+    ``build_chunk_ymd_requests`` builds every timestamp in the cartesian
+    product, so the last year is a full set of months even when CDS has not
+    published them yet. Timestamps after the cutoff are dropped, and the
+    chunk sizes are resized so they still add up to the coordinate that
+    remains. The cutoff is inclusive: a timestamp equal to
+    ``force_end_date`` stays. A chunk that lies entirely past the cutoff
+    is removed. A chunk that crosses it keeps the same CDS request and only
+    its length changes.
+    """
+    if len(time) == 0:
+        raise ValueError("cannot apply force_end_date to an empty time coordinate")
+
+    end = pd.Timestamp(force_end_date).to_datetime64().astype("datetime64[ns]")
+    cut = int(np.searchsorted(time, end, side="right"))
+    if cut == 0:
+        raise ValueError(
+            f"end_date {force_end_date} is before the first timestamp {time[0]}"
+        )
+    if cut >= len(time):
+        return time, time_chunk, time_chunk_requests
+
+    starts = [start for start, _ in time_chunk_requests]
+    new_requests: list[tuple[int, dict[str, Any]]] = []
+    new_sizes: list[int] = []
+    for index, (start, chunk_request) in enumerate(time_chunk_requests):
+        if start >= cut:
+            break
+        next_start = starts[index + 1] if index + 1 < len(starts) else len(time)
+        new_sizes.append(min(next_start, cut) - start)
+        new_requests.append((start, chunk_request))
+
+    trimmed = time[:cut]
+    if isinstance(time_chunk, int) and all(size == time_chunk for size in new_sizes):
+        return trimmed, time_chunk, new_requests
+    if isinstance(time_chunk, int) and len(new_sizes) == 1:
+        return trimmed, new_sizes[0], new_requests
+    return trimmed, tuple(new_sizes), new_requests
+
+
 def build_time_chunk_requests(
-    request: dict[str, Any], request_chunks: dict[str, int], sep: str = "/"
+    request: dict[str, Any],
+    request_chunks: dict[str, int],
+    sep: str = "/",
+    force_end_date: str | None = None,
 ) -> tuple[
     np.typing.NDArray[np.datetime64],
     int | tuple[int, ...],
@@ -274,10 +331,17 @@ def build_time_chunk_requests(
         time, time_chunk, time_chunk_requests = build_chunk_ymd_requests(
             request, request_chunks
         )
+        if force_end_date is not None:
+            time, time_chunk, time_chunk_requests = trim_ymd_time_to_force_end_date(
+                time, time_chunk, time_chunk_requests, force_end_date
+            )
     elif "date" in request:
         time, time_chunk, time_chunk_requests = build_chunk_date_requests(
             request, request_chunks, sep
         )
+        if force_end_date is not None:
+            msg = "forced_end_date param is ignored for date-interval requests."
+            LOGGER.warning(msg)
     else:
         raise ValueError("request must contain either 'year' or 'date'")