fix(ogc): preserve chunking result contracts

Keep feature-ID deduplication across every chunking axis and retain nested GeoJSON property flattening. Centralize transient classification, preserve resumability for future typed transients, and tighten httpx typing and cross-module documentation.

Author: thodson-usgs

dataretrieval/ogc/chunking.py

@@ -454,7 +454,7 @@ def _combine_raw(self) -> tuple[pd.DataFrame, httpx.Response]:
         frames = [self._chunks[i][0] for i in sorted(self._chunks)]
         responses = [response for _, response in self._chunks.values()]
         return (
-            _combine_chunk_frames(frames, dedup=self.plan.has_filter_axis),
+            _combine_chunk_frames(frames),
             _combine_chunk_responses(responses, self.plan.canonical_url),
         )
 

dataretrieval/ogc/combining.py

@@ -37,37 +37,36 @@ def _safe_elapsed(response: httpx.Response) -> timedelta:
     treat the missing attribute as zero elapsed.
     """
     try:
-        return response.elapsed
+        elapsed: object = response.elapsed
     except RuntimeError:
         return timedelta(0)
+    return elapsed if isinstance(elapsed, timedelta) else timedelta(0)
 
 
 def _set_response_url(response: httpx.Response, url: str | httpx.URL) -> None:
     """
     Overwrite the URL surfaced by a response without back-propagating
     the change into any aliased original.
 
-    Try the direct assignment first: on lightweight test mocks ``.url``
-    is a plain writable attribute. On real ``httpx.Response`` it's
-    read-only (it resolves through the bound request), so swap in a
-    fresh :class:`httpx.Request` carrying the new URL — mutating the
-    existing one would leak through any shallow copy that shares the
-    same ``.request``.
+    Lightweight test doubles expose ``.url`` as a writable attribute. Real
+    :class:`httpx.Response` objects resolve it through a bound request, so swap
+    in a fresh request carrying the new URL; mutating the existing request would
+    leak through any shallow copy that shares it.
     """
+    if not isinstance(response, httpx.Response):
+        # Lightweight test doubles expose ``url`` as a writable attribute.
+        response.url = url
+        return
+
+    target = httpx.URL(str(url))
     try:
-        response.url = url  # type: ignore[misc, assignment]
-    except AttributeError:
-        target = httpx.URL(str(url))
-        try:
-            old = response.request
-        except RuntimeError:
-            # No request bound (some hand-built httpx.Response fixtures);
-            # synthesize a minimal one to hold the URL.
-            response.request = httpx.Request("GET", target)
-            return
-        response.request = httpx.Request(
-            method=old.method, url=target, headers=old.headers
-        )
+        old = response.request
+    except RuntimeError:
+        # No request bound (some hand-built httpx.Response fixtures);
+        # synthesize a minimal one to hold the URL.
+        response.request = httpx.Request("GET", target)
+        return
+    response.request = httpx.Request(method=old.method, url=target, headers=old.headers)
 
 
 def _lowest_remaining(responses: list[httpx.Response]) -> httpx.Response:
@@ -104,8 +103,9 @@ def _merge_response(
     given.  ``base`` and ``headers_from`` are never mutated, and the fresh
     ``httpx.Headers`` means downstream mutations don't back-propagate into any
     underlying response — so callers may re-fold idempotently.  This is the one
-    low-level merge behind both pagination (:func:`_paginate`) and the chunked /
-    fan-out aggregation (:func:`_combine_chunk_responses`)."""
+    low-level merge behind both pagination
+    (:func:`~dataretrieval.ogc.engine._paginate`) and the chunked / fan-out
+    aggregation (:func:`_combine_chunk_responses`)."""
     merged = copy.copy(base)
     merged.headers = httpx.Headers(headers_from.headers)
     merged.elapsed = elapsed
@@ -114,77 +114,37 @@ def _merge_response(
     return merged
 
 
-def _combine_chunk_frames(
-    frames: list[pd.DataFrame], *, dedup: bool = True
-) -> pd.DataFrame:
-    """
-    Concatenate per-chunk frames, dropping empties and optionally deduping
-    by ``id``.
+def _combine_chunk_frames(frames: list[pd.DataFrame]) -> pd.DataFrame:
+    """Concatenate per-chunk frames and deduplicate non-null feature IDs.
 
-    Parameters
-    ----------
-    frames : list[pandas.DataFrame]
-        One frame per completed sub-request.
-    dedup : bool, default True
-        Whether to drop duplicate rows keyed on the ``id`` column.
-        List-axis chunks (multi-value list parameters like
-        ``monitoring_location_id``) produce non-overlapping partitions,
-        so dedup is unnecessary and skipping it saves ~2 ms on a 50 k-row
-        result.  Filter-axis chunks *can* overlap (a feature matching
-        multiple OR-clauses appears in each clause's chunk), so dedup is
-        required there.  Callers that know their chunks don't overlap
-        (e.g. :meth:`ChunkedCall._combine_raw` when the plan has no
-        filter axis) pass ``dedup=False``.
+    Empty frames are ignored before concatenation so an empty plain
+    :class:`pandas.DataFrame` cannot downgrade a real ``GeoDataFrame`` and
+    strip its geometry or CRS. When every frame is empty, the first frame is
+    returned to preserve its concrete type.
 
-    Returns
-    -------
-    pandas.DataFrame
-        The concatenated (and optionally deduplicated) result. Empty when
-        every input frame is empty.
-
-    Notes
-    -----
-    An empty chunk can be a plain ``pd.DataFrame()`` (no geopandas);
-    concatenating it with real ``GeoDataFrame``s downgrades the result
-    to plain ``DataFrame`` and strips geometry/CRS, so empties are
-    dropped first.  Dedup on the pre-rename feature ``id`` keeps
-    overlapping user OR-clauses from producing duplicate rows across
-    chunks.
-
-    Dedup is restricted to rows whose ``id`` is non-null. ``pandas``
-    treats NaN==NaN as a duplicate for ``drop_duplicates``, so a
-    blanket call would collapse every id-less row into a single one —
-    silent data loss if any chunk emits features without an
-    ``id`` field.
+    Deduplication is unconditional because overlap is possible on every plan
+    axis. Filter clauses can match the same feature, and list inputs can contain
+    repeated values or otherwise select overlapping records. Rows without an
+    ``id`` are preserved verbatim: pandas treats null values as duplicates, so
+    applying ``drop_duplicates`` to those rows would silently lose data.
     """
-    non_empty = [f for f in frames if not f.empty]
+    non_empty = [frame for frame in frames if not frame.empty]
     if not non_empty:
-        # Preserve the frame type (GeoDataFrame vs DataFrame) of the
-        # input even when every chunk is empty — ``_get_resp_data``
-        # returns ``gpd.GeoDataFrame()`` on empty geopd responses, and
-        # returning a plain ``pd.DataFrame()`` here would downgrade
-        # the type in a downstream ``pd.concat([result, geo_page])`` to
-        # a plain DataFrame and strip geometry/CRS.
         return frames[0] if frames else pd.DataFrame()
     if len(non_empty) == 1:
-        # Single-completed-chunk fast path. Return a copy so callers
-        # who treat ``ChunkedCall.partial_frame`` as a fresh result
-        # (the property docstring says "live; recomputed per access")
-        # don't accidentally mutate ``_chunks[0][0]`` in place.
         return non_empty[0].copy()
+
     combined = pd.concat(non_empty, ignore_index=True)
-    if dedup and "id" in combined.columns:
-        has_id = combined["id"].notna()
-        if has_id.all():
-            combined = combined.drop_duplicates(subset="id", ignore_index=True)
-        elif has_id.any():
-            # Mixed: dedupe only the id-bearing rows; preserve id-less
-            # rows verbatim (their order relative to id-bearing rows
-            # may shift, which is acceptable — dedup can't be id-keyed
-            # for rows without an id).
-            id_rows = combined[has_id].drop_duplicates(subset="id")
-            no_id_rows = combined[~has_id]
-            combined = pd.concat([id_rows, no_id_rows], ignore_index=True)
+    if "id" not in combined.columns:
+        return combined
+
+    has_id = combined["id"].notna()
+    if has_id.all():
+        return combined.drop_duplicates(subset="id", ignore_index=True)
+    if has_id.any():
+        id_rows = combined[has_id].drop_duplicates(subset="id")
+        no_id_rows = combined[~has_id]
+        return pd.concat([id_rows, no_id_rows], ignore_index=True)
     return combined
 
 

dataretrieval/ogc/planning.py

@@ -9,9 +9,11 @@
 drive it.
 
 Result recombination — reassembling the per-chunk frames and responses
-back into one result (:func:`_combine_chunk_frames`,
-:func:`_combine_chunk_responses`, etc.) — lives in the sibling
-:mod:`dataretrieval.ogc.combining` module, which callers import directly.
+back into one result
+(:func:`~dataretrieval.ogc.combining._combine_chunk_frames`,
+:func:`~dataretrieval.ogc.combining._combine_chunk_responses`, etc.) — lives in
+the sibling :mod:`dataretrieval.ogc.combining` module, which callers import
+directly.
 """
 
 from __future__ import annotations
@@ -549,28 +551,6 @@ def total(self) -> int:
         """
         return math.prod((len(self.chunks[ax.arg_key]) for ax in self.axes), start=1)
 
-    @property
-    def has_filter_axis(self) -> bool:
-        """Whether the plan splits along the cql-text ``filter`` axis.
-
-        Filter-axis chunks can overlap (a feature matching multiple
-        OR-clauses appears in each clause's chunk), so deduplication is
-        required when combining their frames.  List-axis chunks (the
-        common case — multi-value list parameters like
-        ``monitoring_location_id``) never overlap, so the combine step
-        can skip the ``drop_duplicates`` call.
-
-        Returns
-        -------
-        bool
-            ``True`` when the plan has a filter axis with >1 chunk;
-            ``False`` otherwise.
-        """
-        return any(
-            ax.joiner != _LIST_SEP and len(self.chunks[ax.arg_key]) > 1
-            for ax in self.axes
-        )
-
     def iter_sub_args(self) -> Iterator[dict[str, Any]]:
         """
         Yield substituted args for each sub-request, in deterministic

dataretrieval/ogc/retry.py

@@ -19,7 +19,7 @@
 import httpx
 import pandas as pd
 
-from dataretrieval.exceptions import RateLimited, ServiceUnavailable, TransientError
+from dataretrieval.exceptions import RateLimited, TransientError
 from dataretrieval.ogc import progress as _progress
 from dataretrieval.ogc.interruptions import (
     ChunkInterrupted,
@@ -193,30 +193,19 @@ def backoff(self, attempt: int, retry_after: float | None) -> float:
 def _classify_transient(
     exc: BaseException,
 ) -> tuple[type[ChunkInterrupted], float | None] | None:
-    """
-    Classify a SINGLE exception as a known transient (resumable) failure.
-
-    Does NOT walk the ``__cause__`` chain — inspects only the exception
-    passed in. Factored out of :func:`_classify_chunk_error` (which walks
-    the ``__cause__`` chain and calls this on each link) so the transient
-    taxonomy lives in one place. :func:`_retryable` deliberately does *not*
-    delegate here — it uses a narrower set (``httpx.InvalidURL`` and a bare
-    ``httpx.HTTPError`` are resumable but not worth an automatic retry).
-
-    Parameters
-    ----------
-    exc : BaseException
-        A single exception to classify.
+    """Classify one exception as a transient, resumable failure.
 
-    Returns
-    -------
-    tuple[type[ChunkInterrupted], float or None] or None
-        ``(interrupted_class, retry_after)`` for a recognized transient
-        failure; ``None`` otherwise.
+    This function owns the shared exception taxonomy; it deliberately does not
+    walk ``__cause__``. :func:`_classify_chunk_error` walks wrapped pagination
+    failures, while :func:`_retryable` applies the narrower automatic-retry
+    policy to this classification.
     """
     if isinstance(exc, RateLimited):
         return QuotaExhausted, exc.retry_after
-    if isinstance(exc, ServiceUnavailable):
+    if isinstance(exc, TransientError):
+        # Every non-rate-limit typed transient is a service interruption. This
+        # fallback keeps future TransientError subclasses resumable after their
+        # inline retries are exhausted.
         return ServiceInterrupted, exc.retry_after
     if isinstance(exc, (httpx.HTTPError, httpx.InvalidURL)):
         return ServiceInterrupted, None
@@ -272,37 +261,20 @@ def _classify_chunk_error(
 
 
 def _retryable(exc: BaseException) -> tuple[bool, float | None]:
-    """
-    Decide whether ``exc`` is a transient worth an automatic retry.
-
-    Only the *top-level* exception is inspected — unlike
-    :func:`_classify_chunk_error`, which walks the ``__cause__`` chain.
-    The distinction matters because ``_paginate`` raises an
-    initial-request transient (429 / 5xx / :class:`httpx.TransportError`)
-    *raw*, but wraps a mid-pagination failure as a base ``DataRetrievalError``.
-    So a raw transient means a sub-request that made no progress and is cheap to
-    re-issue, whereas a mid-pagination failure is left to escalate to a
-    resumable :class:`ChunkInterrupted` rather than re-walked from page 1
-    (which would re-spend the quota just exhausted). ``httpx.InvalidURL``
-    is never retried — a too-long cursor won't fix on a retry.
+    """Decide whether a top-level transient is worth an automatic retry.
 
-    Returns
-    -------
-    tuple[bool, float or None]
-        ``(retryable, retry_after)`` — the server ``Retry-After`` hint
-        (seconds) when the transient carried one, else ``None``.
+    Wrapped mid-pagination failures are not retried from page one; they instead
+    escalate to a resumable :class:`ChunkInterrupted`. ``httpx.InvalidURL`` and
+    non-transport ``httpx.HTTPError`` instances are resumable but deterministic,
+    so they are classified without being retried.
     """
-    # Only initial-page raw transients (TransientError, TransportError) are
-    # retryable — not a wrapped mid-pagination DataRetrievalError. This is a
-    # deliberately narrower taxonomy than _classify_transient's: RateLimited
-    # and ServiceUnavailable are caught via their TransientError base, but
-    # httpx.InvalidURL and bare httpx.HTTPError are intentionally excluded
-    # (a too-long cursor or non-transport HTTP error won't fix on retry), so
-    # this stays a separate check rather than delegating to _classify_transient.
-    if isinstance(exc, TransientError):
-        return True, exc.retry_after
-    if isinstance(exc, httpx.TransportError):
-        return True, None
+    classification = _classify_transient(exc)
+    if classification is None:
+        return False, None
+
+    _, retry_after = classification
+    if isinstance(exc, (TransientError, httpx.TransportError)):
+        return True, retry_after
     return False, None
 
 

dataretrieval/ogc/shaping.py

@@ -101,17 +101,12 @@ def _get_resp_data(
 
     Notes
     -----
-    The non-geopandas branch builds the frame directly from each
-    feature's ``properties`` dict, plus the top-level ``id`` and
-    ``geometry.coordinates`` columns — the ``id`` column is always
-    added (so the downstream rename to the service-specific output id
-    works even on an all-None id), while the ``geometry`` column is
-    added only when at least one feature carries geometry. This skips
-    the GeoJSON envelope entirely, so
-    newly-added Feature-level fields (e.g. ``geometry.type`` after
-    USGS migrated to full GeoJSON geometry objects) can't leak into
-    the result frame; no reactive drop-list needs maintenance every
-    time the upstream schema grows.
+    The non-geopandas branch normalizes each feature's ``properties``
+    object, flattening nested dictionaries with an underscore separator, then
+    adds the top-level ``id`` and ``geometry.coordinates`` columns. The ``id``
+    column is always added so the downstream service-specific rename works even
+    when all IDs are missing; ``geometry`` is added only when coordinates are
+    present. Feature-level envelope fields are deliberately excluded.
     """
     if body is None:
         body = resp.json()
@@ -128,17 +123,11 @@ def _get_resp_data(
         return _empty_feature_frame(geopd)
 
     if not geopd:
-        # Build the frame directly from the flat properties dicts rather than
-        # routing through ``pd.json_normalize``.  The OGC API returns flat
-        # (non-nested) property objects, so ``json_normalize``'s recursive
-        # flattening is unnecessary overhead (~2.5× slower than plain
-        # ``pd.DataFrame`` for typical page sizes).  The ``id`` key is merged
-        # into each row dict up-front so the column is always present (may be
-        # all-None) — ``_arrange_cols`` downstream relies on it for the rename
-        # to the service-specific output_id (daily_id, channel_measurements_id,
-        # …).
-        rows = [{**(f.get("properties") or {}), "id": f.get("id")} for f in features]
-        df = pd.DataFrame(rows)
+        properties = [feature.get("properties") or {} for feature in features]
+        df = pd.json_normalize(properties, sep="_")
+        # Always materialize the feature-level ID (possibly all-None) so
+        # ``_arrange_cols`` can perform the documented service-specific rename.
+        df["id"] = [feature.get("id") for feature in features]
         _attach_coordinates(df, features)
         return df
 

dataretrieval/wateruse.py

@@ -398,7 +398,7 @@ def _next_page_url(response: httpx.Response) -> str | None:
     url = response.links.get("next", {}).get("url")
     if not url:
         return None
-    return url.replace("https://water.usgs.gov", "https://api.water.usgs.gov", 1)
+    return str(url).replace("https://water.usgs.gov", "https://api.water.usgs.gov", 1)
 
 
 def _nwdc_error_detail(response: httpx.Response) -> str | None: