Skip to content

get_nearest_continuous: resume() returns un-post-processed rows (no target_time) #382

Description

@thodson-usgs

Summary

When get_nearest_continuous is interrupted by a chunk failure, the documented recovery — exc.call.resume() — returns the raw get_continuous rows, not the nearest-per-target selection that get_nearest_continuous exists to produce. The resumed frame has no target_time column, so code following the documented resume pattern fails with KeyError: 'target_time'.

The same applies to exc.partial_frame.

Why

get_nearest_continuous fetches and then post-processes:

  • waterdata/nearest.py:149df, md = get_continuous(...) (the chunked call)
  • waterdata/nearest.py:157-174_pick_nearest_row per (monitoring_location_id, target), attaching target_time at line 236

A ChunkInterrupted raised by the fetch propagates out of get_nearest_continuous before the post-processing. The .call handle it carries is the handle for the inner get_continuous call, so resume() re-drives that and returns its combined frame. The outer function's own semantics are simply skipped.

This makes the resume contract shape-inconsistent for any getter that post-processes its chunked fetch: the caller gets a different frame from resume() than the same call would have returned had it not been interrupted.

ChunkInterrupted's docstring — "Call self.call.resume() to pick up where the failure stopped you" — reads as a promise that resuming yields the result the call would have produced. For get_nearest_continuous it does not.

Reproduction

Deterministic, no network — get_continuous is stubbed so the control flow is the real one:

import pandas as pd
import dataretrieval.waterdata.nearest as nearest
from dataretrieval.ogc.interruptions import QuotaExhausted

TARGETS = pd.to_datetime(["2024-01-01T00:00:00Z", "2024-01-02T00:00:00Z"])
RAW = pd.DataFrame({
    "monitoring_location_id": ["USGS-01018035"] * 2,
    "time": TARGETS + pd.Timedelta("3min"),
    "value": [8.1, 8.2],
})

class FakeChunkedCall:
    partial_frame = RAW
    partial_response = None
    def resume(self):
        return RAW, None

def fake_get_continuous(*args, **kwargs):
    raise QuotaExhausted(completed_chunks=0, total_chunks=8, call=FakeChunkedCall())

nearest.get_continuous = fake_get_continuous

try:
    nearest.get_nearest_continuous(
        targets=TARGETS,
        monitoring_location_id=["USGS-01018035"],
        parameter_code="00400",
        window="PT1H",
    )
except QuotaExhausted as exc:
    resumed, _ = exc.call.resume()
    print(sorted(resumed.columns))          # ['monitoring_location_id', 'time', 'value']
    print("target_time" in resumed.columns) # False
    resumed["target_time"]                  # KeyError: 'target_time'

How it showed up

Probing ~7,400 USGS monitoring locations for record gaps, one get_nearest_continuous call per location. Of 291 location failures, 283 were KeyError: 'target_time', and every one occurred immediately after a successful resume — the quota handling waited out the window, got its data, and then could not read it. Only 8 failures were genuine rate-limit exhaustion.

The failure is quiet in the worst way: the resumed frame is a perfectly valid DataFrame, just of the inner call's shape, so it only fails when you index a column the outer function would have added.

Suggested fix

Options, roughly in order of preference:

  1. Have get_nearest_continuous catch ChunkInterrupted, and re-raise an exception whose call wraps resume() so that the post-processing is applied to the resumed frame. The caller then gets the same shape either way.
  2. Failing that, document explicitly on get_nearest_continuous (and any other post-processing getter) that resume() and partial_frame return un-post-processed rows, and say which columns are absent.

Option 1 keeps the resume contract meaning one thing across the API, which is what makes it safe to write generic retry code around.

Workaround

Match the targets against the raw time column yourself, using the same window passed to the query:

if "target_time" in df.columns:
    answered = set(pd.to_datetime(df["target_time"], utc=True))
else:                                   # a resumed / partial frame
    obs = pd.DatetimeIndex(pd.to_datetime(df["time"], utc=True)).sort_values()
    answered = {t for t in targets if _within(obs, t, window)}

Environment

dataretrieval 1.2.0, Python 3.14, macOS. Authenticated via API_USGS_PAT.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions