diff --git a/dataretrieval/waterdata/nearest.py b/dataretrieval/waterdata/nearest.py index 1e9b83c0..5d030923 100644 --- a/dataretrieval/waterdata/nearest.py +++ b/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: diff --git a/tests/waterdata_nearest_test.py b/tests/waterdata_nearest_test.py index 64deeccd..e16e2470 100644 --- a/tests/waterdata_nearest_test.py +++ b/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