fix(nearest): preserve result shape across resume

dataretrieval/waterdata/nearest.py

@@ -6,11 +6,14 @@
 from __future__ import annotations
 
 from collections.abc import Iterable
-from typing import TYPE_CHECKING, Any, Literal, get_args
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Literal, Protocol, cast, get_args
 
+import httpx
 import pandas as pd
 
 from dataretrieval._validation import require_one_of
+from dataretrieval.interruptions import FanOutInterrupted
 from dataretrieval.waterdata.time_series import get_continuous
 
 if TYPE_CHECKING:
@@ -23,6 +26,80 @@
 _VALID_ON_TIE: tuple[OnTie, ...] = get_args(OnTie)
 
 
+class _ResumableCall(Protocol):
+    """Structural subset of a fan-out call needed by the outer decorator."""
+
+    @property
+    def partial_frame(self) -> pd.DataFrame: ...
+
+    @property
+    def partial_response(self) -> httpx.Response | None: ...
+
+    def resume(self) -> tuple[pd.DataFrame, BaseMetadata]: ...
+
+
+class _MutableInterruption(Protocol):
+    """Writable call slot exposed by fan-out interruption instances."""
+
+    call: _ResumableCall | None
+
+
+@dataclass(frozen=True, slots=True)
+class _NearestSelector:
+    """Immutable policy for selecting nearest rows from continuous data."""
+
+    targets: pd.DatetimeIndex
+    window: pd.Timedelta
+    on_tie: OnTie
+
+    def select(self, frame: pd.DataFrame) -> pd.DataFrame:
+        """Apply the public nearest-per-target shape to continuous rows."""
+        return _select_nearest_rows(frame, self.targets, self.window, self.on_tie)
+
+    def select_partial(self, frame: pd.DataFrame) -> pd.DataFrame:
+        """Select partial rows, including the no-completed-chunks shape."""
+        if frame.empty and "time" not in frame.columns:
+            return _empty_nearest_result(frame)
+        return self.select(frame)
+
+
+class _NearestCall:
+    """Preserve nearest-result semantics around an interrupted inner call."""
+
+    def __init__(self, inner: _ResumableCall, selector: _NearestSelector) -> None:
+        self._inner = inner
+        self._selector = selector
+
+    @property
+    def partial_frame(self) -> pd.DataFrame:
+        """Return the live partial rows in the outer getter's shape."""
+        return self._selector.select_partial(self._inner.partial_frame)
+
+    @property
+    def partial_response(self) -> httpx.Response | None:
+        """Pass through the inner call's live aggregate response."""
+        return self._inner.partial_response
+
+    def resume(self) -> tuple[pd.DataFrame, BaseMetadata]:
+        """Resume inner work and apply the outer getter's selection."""
+        try:
+            frame, metadata = self._inner.resume()
+        except FanOutInterrupted as exc:
+            _shape_interruption(exc, self._selector)
+            raise
+        return self._selector.select(frame), metadata
+
+
+def _shape_interruption(
+    exc: FanOutInterrupted,
+    selector: _NearestSelector,
+) -> None:
+    """Decorate one inner interruption with the outer getter's semantics."""
+    exc.partial_frame = selector.select_partial(exc.partial_frame)
+    if exc.call is not None:
+        cast("_MutableInterruption", exc).call = _NearestCall(exc.call, selector)
+
+
 def get_nearest_continuous(
     targets: Iterable[Any],
     monitoring_location_id: str | Iterable[str] | None = None,
@@ -98,6 +175,13 @@ def get_nearest_continuous(
     md : :class:`~dataretrieval.utils.BaseMetadata`
         Metadata from the underlying ``get_continuous`` call.
 
+    Raises
+    ------
+    FanOutInterrupted
+        If the underlying fan-out is interrupted. ``partial_frame`` and
+        ``call.partial_frame`` contain nearest-selected rows with
+        ``target_time``; ``call.resume()`` returns that same public shape.
+
     Notes
     -----
     *Window sizing and ties.* When ``window`` is exactly half the service
@@ -151,21 +235,36 @@ def get_nearest_continuous(
     if len(target_index) == 0:
         raise ValueError("targets must contain at least one timestamp")
 
+    selector = _NearestSelector(target_index, window_td, on_tie)
     filter_expr = _build_window_or_filter(target_index, window_td)
-    df, md = get_continuous(
-        monitoring_location_id=monitoring_location_id,
-        parameter_code=parameter_code,
-        filter=filter_expr,
-        filter_lang="cql-text",
-        **kwargs,
-    )
+    try:
+        df, md = get_continuous(
+            monitoring_location_id=monitoring_location_id,
+            parameter_code=parameter_code,
+            filter=filter_expr,
+            filter_lang="cql-text",
+            **kwargs,
+        )
+    except FanOutInterrupted as exc:
+        _shape_interruption(exc, selector)
+        raise
+    return selector.select(df), md
+
+
+def _select_nearest_rows(
+    df: pd.DataFrame,
+    targets: pd.DatetimeIndex,
+    window_td: pd.Timedelta,
+    on_tie: OnTie,
+) -> pd.DataFrame:
+    """Apply the public nearest-per-target shape to continuous rows."""
     if "time" not in df.columns:
         raise ValueError(
             "get_nearest_continuous requires a 'time' column in the response; "
             "if a `properties` kwarg was passed, include 'time' in it"
         )
     if df.empty:
-        return _empty_nearest_result(df), md
+        return _empty_nearest_result(df)
 
     df = df.assign(time=pd.to_datetime(df["time"], utc=True))
     site_groups = (
@@ -177,12 +276,12 @@ def get_nearest_continuous(
     selected = [
         row
         for _, site_df in site_groups
-        for target in target_index
+        for target in targets
         if (row := _pick_nearest_row(site_df, target, window_td, on_tie)) is not None
     ]
     if not selected:
-        return _empty_nearest_result(df), md
-    return pd.DataFrame(selected).reset_index(drop=True), md
+        return _empty_nearest_result(df)
+    return pd.DataFrame(selected).reset_index(drop=True)
 
 
 def _coerce_targets(targets: Any) -> pd.DatetimeIndex:

tests/waterdata_nearest_test.py

@@ -9,6 +9,7 @@
 import pandas as pd
 import pytest
 
+from dataretrieval.interruptions import QuotaExhausted, ServiceInterrupted
 from dataretrieval.waterdata.nearest import get_nearest_continuous
 
 
@@ -330,3 +331,87 @@ def test_missing_time_column_raises_helpful_error(patch_get_continuous):
             monitoring_location_id="USGS-02238500",
             properties=["value", "monitoring_location_id"],
         )
+
+
+def test_interruption_preserves_nearest_shape_for_partial_and_resumed_rows(
+    patch_get_continuous,
+):
+    """Partial and resumed results keep the outer getter's return shape."""
+    targets = pd.to_datetime(["2024-01-01T00:00:00Z", "2024-01-02T00:00:00Z"], utc=True)
+    raw = _fake_df(
+        [
+            {"time": "2024-01-01T00:03:00Z", "value": 8.1},
+            {"time": "2024-01-02T00:03:00Z", "value": 8.2},
+        ]
+    )
+    response = mock.Mock()
+    metadata = mock.Mock()
+
+    class FakeCall:
+        partial_frame = raw
+        partial_response = response
+
+        def resume(self):
+            return raw, metadata
+
+    patch_get_continuous.side_effect = QuotaExhausted(
+        completed_chunks=1, total_chunks=2, call=FakeCall()
+    )
+
+    with pytest.raises(QuotaExhausted) as excinfo:
+        get_nearest_continuous(
+            targets,
+            monitoring_location_id="USGS-02238500",
+            window="PT1H",
+        )
+
+    interrupted = excinfo.value
+    assert list(interrupted.partial_frame["target_time"]) == list(targets)
+    assert list(interrupted.call.partial_frame["target_time"]) == list(targets)
+    assert interrupted.partial_response is response
+
+    resumed, resumed_metadata = interrupted.call.resume()
+
+    assert list(resumed["target_time"]) == list(targets)
+    assert list(resumed["value"]) == [8.1, 8.2]
+    assert resumed_metadata is metadata
+
+
+def test_nearest_shape_survives_repeated_resume_interruptions(patch_get_continuous):
+    """A resume that is interrupted again returns another outer-shaped call."""
+    target = pd.to_datetime(["2024-01-01T00:00:00Z"], utc=True)
+    raw = _fake_df([{"time": "2024-01-01T00:03:00Z", "value": 8.1}])
+    metadata = mock.Mock()
+
+    class FakeCall:
+        partial_frame = raw
+        partial_response = mock.Mock()
+        attempts = 0
+
+        def resume(self):
+            self.attempts += 1
+            if self.attempts == 1:
+                raise ServiceInterrupted(completed_chunks=1, total_chunks=2, call=self)
+            return raw, metadata
+
+    call = FakeCall()
+    patch_get_continuous.side_effect = QuotaExhausted(
+        completed_chunks=1, total_chunks=2, call=call
+    )
+
+    with pytest.raises(QuotaExhausted) as first:
+        get_nearest_continuous(
+            target,
+            monitoring_location_id="USGS-02238500",
+            window="PT1H",
+        )
+    with pytest.raises(ServiceInterrupted) as second:
+        first.value.call.resume()
+
+    assert list(second.value.partial_frame["target_time"]) == list(target)
+    assert list(second.value.call.partial_frame["target_time"]) == list(target)
+
+    resumed, resumed_metadata = second.value.call.resume()
+
+    assert list(resumed["target_time"]) == list(target)
+    assert resumed_metadata is metadata