From 858f4872d40d0c03de22dea44e2a2d628322879f Mon Sep 17 00:00:00 2001 From: thodson-usgs Date: Sun, 16 Aug 2026 19:06:42 -0500 Subject: [PATCH 1/2] fix(nearest): preserve result shape across resume --- dataretrieval/waterdata/nearest.py | 136 ++++++++++++++++++++++++++--- tests/waterdata_nearest_test.py | 85 ++++++++++++++++++ 2 files changed, 209 insertions(+), 12 deletions(-) diff --git a/dataretrieval/waterdata/nearest.py b/dataretrieval/waterdata/nearest.py index 1e9b83c0..082a1523 100644 --- a/dataretrieval/waterdata/nearest.py +++ b/dataretrieval/waterdata/nearest.py @@ -6,11 +6,13 @@ from __future__ import annotations from collections.abc import Iterable -from typing import TYPE_CHECKING, Any, Literal, get_args +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 +25,95 @@ _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 + + +class _NearestCall: + """Preserve nearest-result semantics around an interrupted inner call.""" + + def __init__( + self, + inner: _ResumableCall, + targets: pd.DatetimeIndex, + window_td: pd.Timedelta, + on_tie: OnTie, + ) -> None: + self._inner = inner + self._targets = targets + self._window_td = window_td + self._on_tie = on_tie + + @property + def partial_frame(self) -> pd.DataFrame: + """Return the live partial rows in the outer getter's shape.""" + return _select_nearest_partial( + self._inner.partial_frame, + self._targets, + self._window_td, + self._on_tie, + ) + + @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._targets, self._window_td, self._on_tie) + raise + return ( + _select_nearest_rows(frame, self._targets, self._window_td, self._on_tie), + metadata, + ) + + +def _shape_interruption( + exc: FanOutInterrupted, + targets: pd.DatetimeIndex, + window_td: pd.Timedelta, + on_tie: OnTie, +) -> None: + """Decorate one inner interruption with the outer getter's semantics.""" + exc.partial_frame = _select_nearest_partial( + exc.partial_frame, targets, window_td, on_tie + ) + if exc.call is not None: + cast("_MutableInterruption", exc).call = _NearestCall( + exc.call, targets, window_td, on_tie + ) + + +def _select_nearest_partial( + frame: pd.DataFrame, + targets: pd.DatetimeIndex, + window_td: pd.Timedelta, + on_tie: OnTie, +) -> pd.DataFrame: + """Select partial rows, including the no-completed-chunks empty shape.""" + if frame.empty and "time" not in frame.columns: + return _empty_nearest_result(frame) + return _select_nearest_rows(frame, targets, window_td, on_tie) + + def get_nearest_continuous( targets: Iterable[Any], monitoring_location_id: str | Iterable[str] | None = None, @@ -98,6 +189,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 @@ -152,20 +250,34 @@ def get_nearest_continuous( raise ValueError("targets must contain at least one timestamp") 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, target_index, window_td, on_tie) + raise + return _select_nearest_rows(df, target_index, window_td, on_tie), 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 +289,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 From 21c2995efddd072ab5c55107f7e996be44360e80 Mon Sep 17 00:00:00 2001 From: thodson-usgs Date: Sun, 16 Aug 2026 19:19:36 -0500 Subject: [PATCH 2/2] refactor(nearest): bundle selection policy --- dataretrieval/waterdata/nearest.py | 75 ++++++++++++------------------ 1 file changed, 31 insertions(+), 44 deletions(-) diff --git a/dataretrieval/waterdata/nearest.py b/dataretrieval/waterdata/nearest.py index 082a1523..5d030923 100644 --- a/dataretrieval/waterdata/nearest.py +++ b/dataretrieval/waterdata/nearest.py @@ -6,6 +6,7 @@ from __future__ import annotations from collections.abc import Iterable +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal, Protocol, cast, get_args import httpx @@ -43,30 +44,36 @@ class _MutableInterruption(Protocol): 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, - targets: pd.DatetimeIndex, - window_td: pd.Timedelta, - on_tie: OnTie, - ) -> None: + def __init__(self, inner: _ResumableCall, selector: _NearestSelector) -> None: self._inner = inner - self._targets = targets - self._window_td = window_td - self._on_tie = on_tie + self._selector = selector @property def partial_frame(self) -> pd.DataFrame: """Return the live partial rows in the outer getter's shape.""" - return _select_nearest_partial( - self._inner.partial_frame, - self._targets, - self._window_td, - self._on_tie, - ) + return self._selector.select_partial(self._inner.partial_frame) @property def partial_response(self) -> httpx.Response | None: @@ -78,40 +85,19 @@ def resume(self) -> tuple[pd.DataFrame, BaseMetadata]: try: frame, metadata = self._inner.resume() except FanOutInterrupted as exc: - _shape_interruption(exc, self._targets, self._window_td, self._on_tie) + _shape_interruption(exc, self._selector) raise - return ( - _select_nearest_rows(frame, self._targets, self._window_td, self._on_tie), - metadata, - ) + return self._selector.select(frame), metadata def _shape_interruption( exc: FanOutInterrupted, - targets: pd.DatetimeIndex, - window_td: pd.Timedelta, - on_tie: OnTie, + selector: _NearestSelector, ) -> None: """Decorate one inner interruption with the outer getter's semantics.""" - exc.partial_frame = _select_nearest_partial( - exc.partial_frame, targets, window_td, on_tie - ) + exc.partial_frame = selector.select_partial(exc.partial_frame) if exc.call is not None: - cast("_MutableInterruption", exc).call = _NearestCall( - exc.call, targets, window_td, on_tie - ) - - -def _select_nearest_partial( - frame: pd.DataFrame, - targets: pd.DatetimeIndex, - window_td: pd.Timedelta, - on_tie: OnTie, -) -> pd.DataFrame: - """Select partial rows, including the no-completed-chunks empty shape.""" - if frame.empty and "time" not in frame.columns: - return _empty_nearest_result(frame) - return _select_nearest_rows(frame, targets, window_td, on_tie) + cast("_MutableInterruption", exc).call = _NearestCall(exc.call, selector) def get_nearest_continuous( @@ -249,6 +235,7 @@ 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) try: df, md = get_continuous( @@ -259,9 +246,9 @@ def get_nearest_continuous( **kwargs, ) except FanOutInterrupted as exc: - _shape_interruption(exc, target_index, window_td, on_tie) + _shape_interruption(exc, selector) raise - return _select_nearest_rows(df, target_index, window_td, on_tie), md + return selector.select(df), md def _select_nearest_rows(