From d18b0c22ffd975be00442597bc71fd7d3a3c1c63 Mon Sep 17 00:00:00 2001 From: thodson-usgs Date: Wed, 12 Aug 2026 20:35:04 -0500 Subject: [PATCH 1/5] fix(ogc): keep the frame type when a result is empty Every OGC getter returned a ``GeoDataFrame`` when its filter matched at least one row and a plain ``DataFrame`` when it matched none. ``_empty_feature_frame`` exists to guarantee an empty page keeps the concrete class -- its docstring calls itself "the single home for this empty-page contract", because a later ``pd.concat([empty, geo_page])`` would otherwise strip geometry and CRS. But ``_deal_with_empty`` runs last, on every getter, 160 lines away in the same file, and built a fresh ``pd.DataFrame(columns=...)`` -- discarding the guarantee it was handed. So ``.geometry``, ``.to_crs()``, ``.plot()`` and ``pd.concat`` across a loop of per-site calls broke only on the empty result: the case least likely to be covered by a caller's tests, and invisible from ``df.columns`` because the ``geometry`` column is still listed. Reindexing the frame we were given applies the schema's columns while keeping the class -- and the ``geometry`` dtype with it. The plain-pandas path is unchanged: ``pd.DataFrame().reindex(columns=[...])`` is still a ``pd.DataFrame``. The regression test uses a ``DataFrame`` subclass rather than real geopandas, so the invariant is checked even where geopandas is absent (it is not in the ``[test]`` extra). Verified to fail against the old line and pass against the new one. 788 passed, mypy --strict clean, all hooks pass. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01BTaSm7HmVb94RSJiKW4WAS --- dataretrieval/ogc/shaping.py | 21 ++++++++++++- tests/waterdata_utils_test.py | 59 +++++++++++++++++++++++++++++++++++ 2 files changed, 79 insertions(+), 1 deletion(-) diff --git a/dataretrieval/ogc/shaping.py b/dataretrieval/ogc/shaping.py index d97d58b5..64b760e8 100644 --- a/dataretrieval/ogc/shaping.py +++ b/dataretrieval/ogc/shaping.py @@ -207,10 +207,29 @@ def _deal_with_empty( endpoint=collection, req_type="schema", base_url=base_url ) properties = list(schema.get("properties", {}).keys()) - return pd.DataFrame(columns=properties) + return _empty_result(return_list, properties) return return_list +def _empty_result(template: pd.DataFrame, columns: list[str]) -> pd.DataFrame: + """Empty result frame shaped like a non-empty one. + + Matches the non-empty path on the three things a caller can observe: the + concrete class, ``object`` dtypes (an all-NaN reindex would give + ``float64``, breaking ``.str``), and -- under geopandas -- an *active* + geometry column with the CRS, without which ``.geometry`` and + ``.to_crs()`` raise on a frame that looks geospatial. ``geometry`` is + appended when absent because :func:`_arrange_cols` does the same. + """ + empty = pd.DataFrame({name: pd.Series(dtype=object) for name in columns}) + if not (GEOPANDAS and isinstance(template, gpd.GeoDataFrame)): + return empty + if "geometry" not in empty.columns: + empty["geometry"] = pd.Series(dtype=object) + empty["geometry"] = gpd.GeoSeries(empty["geometry"], dtype="geometry") + return gpd.GeoDataFrame(empty, geometry="geometry", crs=_CRS) + + def _arrange_cols( df: pd.DataFrame, properties: list[str] | None, diff --git a/tests/waterdata_utils_test.py b/tests/waterdata_utils_test.py index 9081b701..8644d2ed 100644 --- a/tests/waterdata_utils_test.py +++ b/tests/waterdata_utils_test.py @@ -33,6 +33,7 @@ from dataretrieval.ogc.schema import _check_ogc_requests from dataretrieval.ogc.shaping import ( _arrange_cols, + _deal_with_empty, _get_resp_data, _to_snake_case, ) @@ -605,6 +606,64 @@ class _Sentinel: assert isinstance(result, _Sentinel) +@pytest.mark.skipif(not _shaping_module.GEOPANDAS, reason="requires geopandas") +def test_empty_result_matches_a_non_empty_one(): + """An empty result must be usable exactly like a non-empty one. + + A getter used to return a ``GeoDataFrame`` when its filter matched a row + and a plain ``DataFrame`` when it matched none. Preserving only the class + is not enough: the frame also needs ``object`` dtypes (an all-NaN reindex + gives ``float64``, breaking ``.str``) and an *active* geometry column with + the CRS, or ``.geometry`` / ``.to_crs()`` raise on something that looks + geospatial. + """ + import geopandas as gpd + + out = _deal_with_empty( + gpd.GeoDataFrame(), + ["monitoring_location_id", "geometry"], + "daily", + base_url="x", + ) + + assert isinstance(out, gpd.GeoDataFrame) and out.empty + assert out["monitoring_location_id"].dtype == object + assert out.geometry is not None + assert out.crs == "EPSG:4326" + # Usable the way a non-empty result is. + out["monitoring_location_id"].str.startswith("USGS") + real = gpd.GeoDataFrame( + { + "monitoring_location_id": ["USGS-1"], + "geometry": gpd.points_from_xy([1], [2]), + }, + crs="EPSG:4326", + ) + combined = pd.concat([out, real], ignore_index=True) + assert isinstance(combined, gpd.GeoDataFrame) and combined.crs == "EPSG:4326" + + +@pytest.mark.skipif(not _shaping_module.GEOPANDAS, reason="requires geopandas") +def test_empty_result_carries_geometry_even_when_properties_omit_it(): + """``_arrange_cols`` re-appends ``geometry`` on the non-empty path, so an + explicit ``properties=`` must not leave the empty result plain.""" + import geopandas as gpd + + out = _deal_with_empty( + gpd.GeoDataFrame(), ["monitoring_location_id"], "daily", base_url="x" + ) + assert isinstance(out, gpd.GeoDataFrame) + assert out.crs == "EPSG:4326" + + +def test_empty_result_without_geopandas_stays_plain(): + """A plain input must not be promoted, and keeps ``object`` dtypes.""" + out = _deal_with_empty(pd.DataFrame(), ["a", "b"], "daily", base_url="x") + assert type(out) is pd.DataFrame + assert list(out.columns) == ["a", "b"] + assert all(dtype.kind == "O" for dtype in out.dtypes) + + def test_get_resp_data_empty_preserves_geopd_type(): """Same as the stats-side preservation: ``_get_resp_data``'s ``numberReturned == 0`` short-circuit must return a From ee9518a72eb9d563e7c179131c56b286e6b07ba2 Mon Sep 17 00:00:00 2001 From: thodson-usgs Date: Thu, 13 Aug 2026 03:53:31 -0500 Subject: [PATCH 2/5] fix(ogc): preserve empty geospatial frames through concat --- dataretrieval/ogc/shaping.py | 39 ++++++++++++++-------- tests/waterdata_utils_test.py | 62 ++++++++++++++++------------------- 2 files changed, 53 insertions(+), 48 deletions(-) diff --git a/dataretrieval/ogc/shaping.py b/dataretrieval/ogc/shaping.py index 64b760e8..e379d906 100644 --- a/dataretrieval/ogc/shaping.py +++ b/dataretrieval/ogc/shaping.py @@ -52,7 +52,7 @@ def _empty_feature_frame(geopd: bool) -> pd.DataFrame: single home for this empty-page contract, shared by the feature-frame builders that flatten GeoJSON pages. """ - return gpd.GeoDataFrame() if geopd else pd.DataFrame() + return gpd.GeoDataFrame(geometry=[], crs=_CRS) if geopd else pd.DataFrame() def _attach_coordinates(df: pd.DataFrame, features: list[dict[str, Any]]) -> None: @@ -214,20 +214,31 @@ def _deal_with_empty( def _empty_result(template: pd.DataFrame, columns: list[str]) -> pd.DataFrame: """Empty result frame shaped like a non-empty one. - Matches the non-empty path on the three things a caller can observe: the - concrete class, ``object`` dtypes (an all-NaN reindex would give - ``float64``, breaking ``.str``), and -- under geopandas -- an *active* - geometry column with the CRS, without which ``.geometry`` and - ``.to_crs()`` raise on a frame that looks geospatial. ``geometry`` is - appended when absent because :func:`_arrange_cols` does the same. + Reindex the frame produced by the response path so its concrete class is + preserved. New, empty pandas columns are cast to ``object`` (matching + ``DataFrame(columns=...)`` and keeping string accessors usable). Under + geopandas, also mirror :func:`_arrange_cols` by retaining ``geometry`` and + making it active with the source CRS, or the API's CRS when the empty page + did not carry one. """ - empty = pd.DataFrame({name: pd.Series(dtype=object) for name in columns}) - if not (GEOPANDAS and isinstance(template, gpd.GeoDataFrame)): - return empty - if "geometry" not in empty.columns: - empty["geometry"] = pd.Series(dtype=object) - empty["geometry"] = gpd.GeoSeries(empty["geometry"], dtype="geometry") - return gpd.GeoDataFrame(empty, geometry="geometry", crs=_CRS) + is_geo = GEOPANDAS and isinstance(template, gpd.GeoDataFrame) + result_columns = list(columns) + if is_geo and "geometry" not in result_columns: + result_columns.append("geometry") + + empty = template.reindex(columns=result_columns) + if not is_geo: + return empty.astype(object) + + for name in result_columns: + if name != "geometry": + empty[name] = empty[name].astype(object) + empty = empty.set_geometry("geometry") + try: + source_crs = template.crs + except AttributeError: + source_crs = None + return empty.set_crs(source_crs or _CRS) def _arrange_cols( diff --git a/tests/waterdata_utils_test.py b/tests/waterdata_utils_test.py index 8644d2ed..0123f3ae 100644 --- a/tests/waterdata_utils_test.py +++ b/tests/waterdata_utils_test.py @@ -606,32 +606,47 @@ class _Sentinel: assert isinstance(result, _Sentinel) +def test_empty_result_preserves_frame_subclass_without_geopandas(): + """Type preservation must not depend on geopandas being installed.""" + + class Frame(pd.DataFrame): + @property + def _constructor(self): + return Frame + + out = _deal_with_empty(Frame(), ["a", "b"], "daily", base_url="x") + + assert type(out) is Frame + assert list(out.columns) == ["a", "b"] + assert all(dtype.kind == "O" for dtype in out.dtypes) + assert ( + type(_deal_with_empty(pd.DataFrame(), ["a"], "daily", base_url="x")) + is pd.DataFrame + ) + + @pytest.mark.skipif(not _shaping_module.GEOPANDAS, reason="requires geopandas") -def test_empty_result_matches_a_non_empty_one(): - """An empty result must be usable exactly like a non-empty one. - - A getter used to return a ``GeoDataFrame`` when its filter matched a row - and a plain ``DataFrame`` when it matched none. Preserving only the class - is not enough: the frame also needs ``object`` dtypes (an all-NaN reindex - gives ``float64``, breaking ``.str``) and an *active* geometry column with - the CRS, or ``.geometry`` / ``.to_crs()`` raise on something that looks - geospatial. - """ +def test_empty_result_matches_a_non_empty_geodataframe(): + """An empty geospatial result remains usable like a non-empty one.""" import geopandas as gpd + response = _resp_ok([]) + page = _get_resp_data(response, geopd=True) + template = pd.concat([page], ignore_index=True) out = _deal_with_empty( - gpd.GeoDataFrame(), - ["monitoring_location_id", "geometry"], + template, + ["monitoring_location_id"], "daily", base_url="x", ) + assert isinstance(template, gpd.GeoDataFrame) assert isinstance(out, gpd.GeoDataFrame) and out.empty assert out["monitoring_location_id"].dtype == object assert out.geometry is not None assert out.crs == "EPSG:4326" - # Usable the way a non-empty result is. out["monitoring_location_id"].str.startswith("USGS") + out.to_crs("EPSG:3857") real = gpd.GeoDataFrame( { "monitoring_location_id": ["USGS-1"], @@ -643,27 +658,6 @@ def test_empty_result_matches_a_non_empty_one(): assert isinstance(combined, gpd.GeoDataFrame) and combined.crs == "EPSG:4326" -@pytest.mark.skipif(not _shaping_module.GEOPANDAS, reason="requires geopandas") -def test_empty_result_carries_geometry_even_when_properties_omit_it(): - """``_arrange_cols`` re-appends ``geometry`` on the non-empty path, so an - explicit ``properties=`` must not leave the empty result plain.""" - import geopandas as gpd - - out = _deal_with_empty( - gpd.GeoDataFrame(), ["monitoring_location_id"], "daily", base_url="x" - ) - assert isinstance(out, gpd.GeoDataFrame) - assert out.crs == "EPSG:4326" - - -def test_empty_result_without_geopandas_stays_plain(): - """A plain input must not be promoted, and keeps ``object`` dtypes.""" - out = _deal_with_empty(pd.DataFrame(), ["a", "b"], "daily", base_url="x") - assert type(out) is pd.DataFrame - assert list(out.columns) == ["a", "b"] - assert all(dtype.kind == "O" for dtype in out.dtypes) - - def test_get_resp_data_empty_preserves_geopd_type(): """Same as the stats-side preservation: ``_get_resp_data``'s ``numberReturned == 0`` short-circuit must return a From d6ea9e14a25fc80b5055fa6299df000e1dbd6448 Mon Sep 17 00:00:00 2001 From: thodson-usgs Date: Thu, 13 Aug 2026 04:39:37 -0500 Subject: [PATCH 3/5] fix(ogc): choose frame shape per request --- dataretrieval/ngwmn.py | 1 + dataretrieval/ogc/engine.py | 59 ++++++++-- dataretrieval/ogc/shaping.py | 167 ++++++++++++--------------- dataretrieval/waterdata/reference.py | 6 +- dataretrieval/waterdata/utils.py | 6 + tests/ngwmn_test.py | 46 ++++++++ tests/waterdata_filters_test.py | 2 +- tests/waterdata_test.py | 23 ++++ tests/waterdata_utils_test.py | 63 +++++++--- 9 files changed, 255 insertions(+), 118 deletions(-) diff --git a/dataretrieval/ngwmn.py b/dataretrieval/ngwmn.py index 154899fa..cc0cf9bf 100644 --- a/dataretrieval/ngwmn.py +++ b/dataretrieval/ngwmn.py @@ -105,6 +105,7 @@ def _get(service: str, local_vars: dict[str, Any]) -> tuple[pd.DataFrame, BaseMe service, output_id=_NGWMN_OUTPUT_ID, base_url=NGWMN_OGC_API_URL, + spatial=service == "sites", dialect=NGWMN_DIALECT, ) diff --git a/dataretrieval/ogc/engine.py b/dataretrieval/ogc/engine.py index 5c5236d8..be63e69c 100644 --- a/dataretrieval/ogc/engine.py +++ b/dataretrieval/ogc/engine.py @@ -145,7 +145,10 @@ async def _paginate( def _ogc_parse_response( - resp: httpx.Response, *, geopd: bool + resp: httpx.Response, + *, + geopd: bool, + include_geometry: bool = True, ) -> tuple[pd.DataFrame, str | None]: """Parse one OGC API page: extract the DataFrame and the next-page URL. @@ -156,7 +159,12 @@ def _ogc_parse_response( """ body = resp.json() return ( - _get_resp_data(resp, geopd=geopd, body=body), + _get_resp_data( + resp, + geopd=geopd, + include_geometry=include_geometry, + body=body, + ), _next_req_url(resp, body=body) or None, ) @@ -166,6 +174,7 @@ async def _walk_pages( req: httpx.Request, client: httpx.AsyncClient | None = None, *, + include_geometry: bool = True, row_cap: int | None = None, ) -> tuple[pd.DataFrame, httpx.Response]: """ @@ -179,7 +188,10 @@ async def _walk_pages( Parameters ---------- geopd : bool - Whether geopandas is installed (drives geometry handling). + Whether this request's pages use an active GeoDataFrame. + include_geometry : bool, optional + Whether plain-DataFrame pages retain raw coordinate lists. False for + nonspatial collections and requests using ``skip_geometry=True``. req : httpx.Request The initial HTTP request to send. client : httpx.AsyncClient, optional @@ -216,7 +228,11 @@ async def follow_up(cursor: str, sess: httpx.AsyncClient) -> httpx.Response: return await _paginate( req, - parse_response=functools.partial(_ogc_parse_response, geopd=geopd), + parse_response=functools.partial( + _ogc_parse_response, + geopd=geopd, + include_geometry=include_geometry, + ), follow_up=follow_up, client=client, row_cap=row_cap, @@ -229,6 +245,7 @@ def get_ogc_data( output_id: str, *, base_url: str, + spatial: bool, max_rows: int | None = None, extra_id_cols: frozenset[str] | set[str] = frozenset(), dialect: OgcDialect | None = None, @@ -266,6 +283,10 @@ def get_ogc_data( rejected at send time, surfacing as a NetworkError about an unknown service. Requiring it moves that mistake to the call site, where mypy catches it. + spatial : bool + Whether this collection's result contract includes feature geometry. + The adapter supplies this semantic fact; ``skip_geometry`` and + geopandas availability then select the concrete frame representation. extra_id_cols : set or frozenset, optional Synthetic id columns to push to the end of a result frame (see :func:`_arrange_cols`). Defaults to an empty set. @@ -314,6 +335,12 @@ def get_ogc_data( convert_type = args.pop("convert_type", False) args = {k: v for k, v in args.items() if v is not None} + # Choose one semantic frame shape for the whole request. Every page and + # the all-empty finalizer receive these same values, so frame type never + # depends on whether a particular page happens to carry geometry. + include_geometry = spatial and not bool(args.get("skip_geometry", False)) + geopd = GEOPANDAS and include_geometry + # Post-processing is injected into the chunker rather than applied here, # so it runs on *every* exit: the normal return AND a later # ``exc.call.resume()`` after a ChunkInterrupted (which never re-enters @@ -328,6 +355,8 @@ def get_ogc_data( output_id=output_id, convert_type=convert_type, collection=collection, + geopd=geopd, + include_geometry=include_geometry, max_rows=max_rows, extra_id_cols=extra_id_cols, dialect=dialect, @@ -354,7 +383,12 @@ def get_ogc_data( # ``(df, BaseMetadata)`` shape rather than a raw response pair. return FanOut( [req], - functools.partial(_walk_pages, GEOPANDAS, row_cap=max_rows), + functools.partial( + _walk_pages, + geopd, + include_geometry=include_geometry, + row_cap=max_rows, + ), RetryPolicy.from_env(), finalize, canonical_url=str(req.url), @@ -371,7 +405,11 @@ def get_ogc_data( _construct_api_requests, base_url=base_url, dialect=dialect ) fetch = functools.partial( - _fetch_once, build_request=build_request, row_cap=max_rows + _fetch_once, + build_request=build_request, + geopd=geopd, + include_geometry=include_geometry, + row_cap=max_rows, ) run = chunking.multi_value_chunked(build_request=build_request)(fetch) # No progress block here: the executor that emits the events owns the line @@ -383,6 +421,8 @@ async def _fetch_once( args: dict[str, Any], *, build_request: Callable[..., httpx.Request], + geopd: bool, + include_geometry: bool, row_cap: int | None = None, ) -> tuple[pd.DataFrame, httpx.Response]: """Send one prepared-args OGC request asynchronously; return (frame, response). @@ -400,4 +440,9 @@ async def _fetch_once( synchronously. The return shape is ``(frame, response)``. """ req = build_request(**args) - return await _walk_pages(geopd=GEOPANDAS, req=req, row_cap=row_cap) + return await _walk_pages( + geopd=geopd, + req=req, + include_geometry=include_geometry, + row_cap=row_cap, + ) diff --git a/dataretrieval/ogc/shaping.py b/dataretrieval/ogc/shaping.py index e379d906..760e6f41 100644 --- a/dataretrieval/ogc/shaping.py +++ b/dataretrieval/ogc/shaping.py @@ -43,16 +43,33 @@ ) -def _empty_feature_frame(geopd: bool) -> pd.DataFrame: - """Empty result frame for a page that carries no features. - - Returns a ``GeoDataFrame`` when geopandas is available so a downstream - ``pd.concat([empty_page, geo_page])`` doesn't downgrade a geopandas - user's result to a plain ``DataFrame`` (stripping geometry/CRS). The - single home for this empty-page contract, shared by the feature-frame - builders that flatten GeoJSON pages. +def _empty_feature_frame( + geopd: bool, + columns: list[str] | None = None, + *, + include_geometry: bool = True, +) -> pd.DataFrame: + """Construct an empty feature frame for one request's geometry mode. + + ``geopd`` and ``include_geometry`` are selected once before pagination and + reused for every page and the final all-empty result. This keeps empty and + non-empty pages in the same frame family, so ordinary ``pd.concat`` is + sufficient. ``columns`` is supplied only when finalization has fetched the + collection schema; page-level empties intentionally remain schema-light. """ - return gpd.GeoDataFrame(geometry=[], crs=_CRS) if geopd else pd.DataFrame() + result_columns = list(columns or []) + if include_geometry: + if "geometry" not in result_columns: + result_columns.append("geometry") + else: + result_columns = [name for name in result_columns if name != "geometry"] + + data = {name: pd.Series(dtype=object) for name in result_columns} + if not geopd: + return pd.DataFrame(data, columns=result_columns) + + data["geometry"] = gpd.GeoSeries([], crs=_CRS) + return gpd.GeoDataFrame(data, columns=result_columns, geometry="geometry", crs=_CRS) def _attach_coordinates(df: pd.DataFrame, features: list[dict[str, Any]]) -> None: @@ -86,6 +103,7 @@ def _get_resp_data( resp: httpx.Response, geopd: bool, *, + include_geometry: bool = True, body: dict[str, Any] | None = None, ) -> pd.DataFrame: """ @@ -97,8 +115,11 @@ def _get_resp_data( The HTTP response object expected to contain a JSON body with a "features" key. geopd : bool - Indicates whether geopandas is installed and should be used to - handle geometries. + Whether this request's result contract uses an active GeoDataFrame. + Selected once before pagination rather than inferred from page content. + include_geometry : bool, optional + Whether plain-DataFrame results retain raw coordinate lists. False for + nonspatial collections and ``skip_geometry=True`` requests. body : dict, optional Pre-parsed JSON body for ``resp``. When provided, skips the ``resp.json()`` call — useful when the caller has already @@ -132,11 +153,11 @@ def _get_resp_data( # actually carry features. An absent/empty ``features`` is also the real # schema-drift shape (a 200 with no features) — treat it as empty rather # than crash with a ``KeyError`` downstream, which ``_paginate`` would - # mistake for a transient transport error. ``_empty_feature_frame`` - # preserves the GeoDataFrame type on the short-circuit (see its docstring). + # mistake for a transient transport error. The request-level mode gives + # this page the same frame family as every non-empty page in the walk. features = body.get("features") or [] if not features: - return _empty_feature_frame(geopd) + return _empty_feature_frame(geopd, include_geometry=include_geometry) if not geopd: properties = [feature.get("properties") or {} for feature in features] @@ -144,25 +165,18 @@ def _get_resp_data( # Always materialize the feature-level ID (possibly all-None) so # ``_arrange_cols`` can perform the documented collection-specific rename. df["id"] = [feature.get("id") for feature in features] - _attach_coordinates(df, features) + if include_geometry: + _attach_coordinates(df, features) return df - # Organize json into geodataframe and make sure id column comes along. - # NGWMN observation collections (water levels, lithology, …) return - # features with no ``geometry`` key at all; ``_geo_feature_frame`` absorbs - # that, and the all-null check below then yields a plain DataFrame. + # A spatial request remains geospatial even when this particular page + # carries only null/missing geometries; changing frame family based on page + # contents makes pagination concat order-dependent. df = _geo_feature_frame(features) - # Mirror the non-geopandas branch's defensive ``f.get("id")`` so a feature - # missing a top-level ``id`` yields None rather than a KeyError. + # Mirror the plain branch's defensive ``f.get("id")`` so a feature missing + # a top-level id yields None rather than a KeyError. df["id"] = [f.get("id") for f in features] - df = df[["id"] + [col for col in df.columns if col != "id"]] - - # If no geometry present, then return pandas dataframe. A geodataframe - # is not needed. - if df["geometry"].isnull().all(): - df = pd.DataFrame(df.drop(columns="geometry")) - - return df + return df[["id"] + [col for col in df.columns if col != "id"]] def _deal_with_empty( @@ -171,74 +185,32 @@ def _deal_with_empty( collection: str, *, base_url: str, + geopd: bool, + include_geometry: bool, ) -> pd.DataFrame: + """Apply the collection schema when an entire result is empty. + + Page-level empties already carry the request's frame family. The complete + column list is available only from explicit ``properties`` or the + collection schema, so construct the fully shaped empty once here rather + than fetching schema during pagination. """ - Handles empty DataFrame results by returning a DataFrame with appropriate columns. + if not return_list.empty: + return return_list - If `return_list` is empty, determines the column names to use: - - If `properties` is not provided or contains only NaN values, - retrieves schema properties from the specified collection. - - Otherwise, uses the provided `properties` list as column names. + if not properties or all(pd.isna(properties)): + from dataretrieval.ogc.schema import _check_ogc_requests - Parameters - ---------- - return_list : pd.DataFrame - The DataFrame to check for emptiness. - properties : Optional[List[str]] - List of property names to use as columns, or None. - collection : str - The collection endpoint to query for schema properties if needed. - base_url : str - OGC API base URL to use for that schema query — the API the empty - result came from, not any particular collection. + schema, _ = _check_ogc_requests( + endpoint=collection, req_type="schema", base_url=base_url + ) + properties = list(schema.get("properties", {}).keys()) - Returns - ------- - pd.DataFrame - The original DataFrame if not empty, otherwise an empty - DataFrame with the appropriate columns. - """ - if return_list.empty: - if not properties or all(pd.isna(properties)): - # Schema lookup performs HTTP only for an empty result. - from dataretrieval.ogc.schema import _check_ogc_requests - - schema, _ = _check_ogc_requests( - endpoint=collection, req_type="schema", base_url=base_url - ) - properties = list(schema.get("properties", {}).keys()) - return _empty_result(return_list, properties) - return return_list - - -def _empty_result(template: pd.DataFrame, columns: list[str]) -> pd.DataFrame: - """Empty result frame shaped like a non-empty one. - - Reindex the frame produced by the response path so its concrete class is - preserved. New, empty pandas columns are cast to ``object`` (matching - ``DataFrame(columns=...)`` and keeping string accessors usable). Under - geopandas, also mirror :func:`_arrange_cols` by retaining ``geometry`` and - making it active with the source CRS, or the API's CRS when the empty page - did not carry one. - """ - is_geo = GEOPANDAS and isinstance(template, gpd.GeoDataFrame) - result_columns = list(columns) - if is_geo and "geometry" not in result_columns: - result_columns.append("geometry") - - empty = template.reindex(columns=result_columns) - if not is_geo: - return empty.astype(object) - - for name in result_columns: - if name != "geometry": - empty[name] = empty[name].astype(object) - empty = empty.set_geometry("geometry") - try: - source_crs = template.crs - except AttributeError: - source_crs = None - return empty.set_crs(source_crs or _CRS) + return _empty_feature_frame( + geopd, + properties, + include_geometry=include_geometry, + ) def _arrange_cols( @@ -391,6 +363,8 @@ def _finalize_ogc( output_id: str, convert_type: bool, collection: str, + geopd: bool, + include_geometry: bool, max_rows: int | None = None, extra_id_cols: frozenset[str] | set[str] = frozenset(), dialect: OgcDialect | None = None, @@ -419,7 +393,14 @@ def _finalize_ogc( """ if dialect is None: dialect = DEFAULT_DIALECT - frame = _deal_with_empty(frame, properties, collection, base_url=base_url) + frame = _deal_with_empty( + frame, + properties, + collection, + base_url=base_url, + geopd=geopd, + include_geometry=include_geometry, + ) # Normalize to PEP-8 snake_case column names *first*, so the dialect's # ``time_cols``/``numerical_cols``/``sort_cols`` (all snake_case) match # regardless of whether the API returns snake_case (Water Data, where diff --git a/dataretrieval/waterdata/reference.py b/dataretrieval/waterdata/reference.py index 14a1884b..1d988640 100644 --- a/dataretrieval/waterdata/reference.py +++ b/dataretrieval/waterdata/reference.py @@ -112,7 +112,11 @@ def get_reference_table( if limit is not None: query_args["limit"] = limit return get_ogc_data( - args=query_args, output_id=output_id, collection=collection, max_rows=max_rows + args=query_args, + output_id=output_id, + collection=collection, + max_rows=max_rows, + spatial=False, ) diff --git a/dataretrieval/waterdata/utils.py b/dataretrieval/waterdata/utils.py index d2cc9182..3a0ec5c3 100644 --- a/dataretrieval/waterdata/utils.py +++ b/dataretrieval/waterdata/utils.py @@ -178,6 +178,8 @@ def get_ogc_data( output_id: str | None = None, max_rows: int | None = None, cql_body: str | None = None, + *, + spatial: bool = True, ) -> tuple[pd.DataFrame, BaseMetadata]: """Water-Data wrapper over :func:`~dataretrieval.ogc.get_ogc_data`. @@ -204,6 +206,9 @@ def get_ogc_data( cql_body : str, optional A verbatim CQL2 JSON body to POST instead of building the query from ``args`` (see the facade's ``cql_body``). Used by :func:`get_cql`. + spatial : bool, optional + Whether the collection carries feature geometry. Water Data's typed + feature collections do; reference tables pass ``False``. Returns ------- @@ -221,6 +226,7 @@ def get_ogc_data( output_id, max_rows=max_rows, base_url=OGC_API_URL, + spatial=spatial, extra_id_cols=_EXTRA_ID_COLS, dialect=WATERDATA_DIALECT, cql_body=cql_body, diff --git a/tests/ngwmn_test.py b/tests/ngwmn_test.py index 5921aa1e..78dd0fac 100644 --- a/tests/ngwmn_test.py +++ b/tests/ngwmn_test.py @@ -307,6 +307,28 @@ def test_get_sites_skip_geometry(httpx_mock): assert _queries(httpx_mock, "sites")[0]["skipGeometry"] == ["true"] +def test_get_sites_empty_skip_geometry_is_plain(httpx_mock): + """The request's skip-geometry mode also governs an all-empty result.""" + httpx_mock.add_response( + method="GET", + url=_schema_re("sites"), + json={ + "properties": { + "geometry": {}, + "monitoring_location_id": {}, + "state_name": {}, + } + }, + ) + _mock(httpx_mock, "sites", _collection([])) + + df, _ = ngwmn.get_sites(monitoring_location_id="USGS-00000000", skip_geometry=True) + + assert type(df) is DataFrame + assert "geometry" not in df.columns + assert "monitoring_location_id" in df.columns + + def test_get_sites_state_accepts_name_postal_or_fips(httpx_mock): """The single ``state`` parameter accepts a full name, postal code, or FIPS code, and all three are normalized to the full ``state_name`` that the @@ -340,6 +362,28 @@ def test_get_providers(httpx_mock): assert "geometry" not in df.columns +def test_get_providers_empty_stays_plain(httpx_mock): + """A nonspatial collection stays tabular when its query has no matches.""" + httpx_mock.add_response( + method="GET", + url=_schema_re("providers"), + json={ + "properties": { + "agency_code": {}, + "organization_type": {}, + "state": {}, + } + }, + ) + _mock(httpx_mock, "providers", _collection([])) + + df, _ = ngwmn.get_providers(agency_code="NONE") + + assert type(df) is DataFrame + assert "geometry" not in df.columns + assert {"agency_code", "organization_type", "state"}.issubset(df.columns) + + def test_get_providers_state_accepts_name_postal_or_fips(httpx_mock): """``get_providers`` normalizes any state encoding to the uppercase postal code that the ``providers`` collection queries on -- the other half of the @@ -515,8 +559,10 @@ def test_empty_result_returns_typed_empty_frame(httpx_mock): df, _ = ngwmn.get_water_level(monitoring_location_id=_SITE) + assert type(df) is DataFrame assert df.empty assert "monitoring_location_id" in df.columns + assert "geometry" not in df.columns # --- live upstream monitor --------------------------------------------------- diff --git a/tests/waterdata_filters_test.py b/tests/waterdata_filters_test.py index 9ca439c6..ff6c6f25 100644 --- a/tests/waterdata_filters_test.py +++ b/tests/waterdata_filters_test.py @@ -164,7 +164,7 @@ def test_long_filter_fans_out_into_multiple_requests(): expr = _filter_chunking_clauses() sent_filters: list[str] = [] - async def fake_walk_pages(*, geopd, req, row_cap=None): + async def fake_walk_pages(*, geopd, req, include_geometry=True, row_cap=None): idx = len(sent_filters) sent_filters.append(_query_params(req).get("filter", [None])[0]) return pd.DataFrame({"id": [f"chunk-{idx}"], "value": [idx]}), _fake_response() diff --git a/tests/waterdata_test.py b/tests/waterdata_test.py index 7d8fe80f..e3436925 100644 --- a/tests/waterdata_test.py +++ b/tests/waterdata_test.py @@ -762,6 +762,29 @@ def test_get_daily_no_geometry(httpx_mock): assert _sent(httpx_mock, "daily")[0]["skipGeometry"] == ["true"] +def test_get_daily_empty_no_geometry(httpx_mock): + """An empty skip-geometry request has the same plain shape as a hit.""" + httpx_mock.add_response( + method="GET", + url=_schema_url("daily"), + json={ + "properties": { + "geometry": {}, + "id": {}, + "monitoring_location_id": {}, + "value": {}, + } + }, + ) + _mock_items(httpx_mock, "daily", body={"features": [], "links": []}) + + df, _ = get_daily(monitoring_location_id="USGS-NOT-FOUND", skip_geometry=True) + + assert type(df) is DataFrame + assert "geometry" not in df.columns + assert {"monitoring_location_id", "value"}.issubset(df.columns) + + def test_get_continuous(httpx_mock): """Continuous observations are timestamped (not date-only), so ``time`` comes back as a UTC-aware datetime column.""" diff --git a/tests/waterdata_utils_test.py b/tests/waterdata_utils_test.py index 0123f3ae..d76fb7d1 100644 --- a/tests/waterdata_utils_test.py +++ b/tests/waterdata_utils_test.py @@ -54,6 +54,8 @@ extra_id_cols=_EXTRA_ID_COLS, dialect=WATERDATA_DIALECT, base_url=OGC_API_URL, + geopd=False, + include_geometry=True, ) _LOGGER_NAME = _utils_module.__name__ @@ -606,24 +608,30 @@ class _Sentinel: assert isinstance(result, _Sentinel) -def test_empty_result_preserves_frame_subclass_without_geopandas(): - """Type preservation must not depend on geopandas being installed.""" - - class Frame(pd.DataFrame): - @property - def _constructor(self): - return Frame - - out = _deal_with_empty(Frame(), ["a", "b"], "daily", base_url="x") - - assert type(out) is Frame - assert list(out.columns) == ["a", "b"] - assert all(dtype.kind == "O" for dtype in out.dtypes) - assert ( - type(_deal_with_empty(pd.DataFrame(), ["a"], "daily", base_url="x")) - is pd.DataFrame +def test_empty_result_uses_request_geometry_contract(): + """All-empty shaping follows the request mode, not the template class.""" + with_geometry = _deal_with_empty( + pd.DataFrame(), + ["a", "b"], + "daily", + base_url="x", + geopd=False, + include_geometry=True, + ) + without_geometry = _deal_with_empty( + pd.DataFrame(), + ["a", "geometry", "b"], + "daily", + base_url="x", + geopd=False, + include_geometry=False, ) + assert type(with_geometry) is pd.DataFrame + assert list(with_geometry.columns) == ["a", "b", "geometry"] + assert all(dtype.kind == "O" for dtype in with_geometry.dtypes) + assert list(without_geometry.columns) == ["a", "b"] + @pytest.mark.skipif(not _shaping_module.GEOPANDAS, reason="requires geopandas") def test_empty_result_matches_a_non_empty_geodataframe(): @@ -638,6 +646,8 @@ def test_empty_result_matches_a_non_empty_geodataframe(): ["monitoring_location_id"], "daily", base_url="x", + geopd=True, + include_geometry=True, ) assert isinstance(template, gpd.GeoDataFrame) @@ -680,6 +690,27 @@ class _Sentinel: assert isinstance(result, _Sentinel) +@pytest.mark.skipif(not _shaping_module.GEOPANDAS, reason="requires geopandas") +def test_get_resp_data_keeps_one_geospatial_type_when_geometry_is_missing(): + """A spatial request cannot change frame family based on page contents.""" + import geopandas as gpd + + empty = _get_resp_data(_resp_ok([]), geopd=True) + missing = _get_resp_data( + _resp_ok([{"id": "1", "properties": {"value": 1}}]), + geopd=True, + ) + + assert isinstance(empty, gpd.GeoDataFrame) + assert isinstance(missing, gpd.GeoDataFrame) + assert missing.geometry.isna().all() + for pages in ([empty, missing], [missing, empty]): + combined = pd.concat(pages, ignore_index=True) + assert isinstance(combined, gpd.GeoDataFrame) + assert combined.geometry.name == "geometry" + assert combined.crs == "EPSG:4326" + + def test_get_resp_data_attaches_wgs84_crs(): """A geometry-bearing Water Data page should come back tagged as EPSG:4326 (the CRS the coordinates are published in), so callers can From b2a94ff59d8317d5bf3844acd4aa4f6cea7161bb Mon Sep 17 00:00:00 2001 From: thodson-usgs Date: Thu, 13 Aug 2026 08:58:11 -0500 Subject: [PATCH 4/5] refactor(ogc): aggregate raw feature pages --- benchmarks/ogc_raw_feature_pagination.py | 241 ++++++++++++++++++ dataretrieval/ogc/engine.py | 106 ++++---- dataretrieval/ogc/shaping.py | 84 ++----- dataretrieval/transport/pagination.py | 80 ++++-- dataretrieval/waterdata/stats.py | 14 +- tests/ogc_live_test.py | 298 +++++++++++++++++++++++ tests/transport_test.py | 45 ++++ tests/waterdata_filters_test.py | 7 +- tests/waterdata_test.py | 83 +++++++ tests/waterdata_utils_test.py | 104 ++++---- 10 files changed, 857 insertions(+), 205 deletions(-) create mode 100644 benchmarks/ogc_raw_feature_pagination.py create mode 100644 tests/ogc_live_test.py diff --git a/benchmarks/ogc_raw_feature_pagination.py b/benchmarks/ogc_raw_feature_pagination.py new file mode 100644 index 00000000..0829ecb3 --- /dev/null +++ b/benchmarks/ogc_raw_feature_pagination.py @@ -0,0 +1,241 @@ +"""Benchmark raw OGC feature aggregation against per-page frame conversion. + +Run from the repository root:: + + python benchmarks/ogc_raw_feature_pagination.py + +The default workload shapes 10,000 features in ten pages for both spatial and +nonspatial collections. ``page_frames`` reproduces PR #373's per-page +conversion boundary with the current equivalent feature shaper; +``raw_features`` uses the experiment's completed-chunk boundary. The report +includes conversion count, median wall time, and median isolated-process peak +RSS. Guardrails are informational, not assertions, because timing and memory +measurements vary by machine. +""" + +from __future__ import annotations + +import argparse +import gc +import json +import platform +import resource +import statistics +import subprocess +import sys +import time +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import pandas as pd + +from dataretrieval.ogc.engine import _combine_feature_pages +from dataretrieval.ogc.shaping import _feature_frame + +Feature = dict[str, Any] +Pages = list[list[Feature]] +Strategy = Callable[[Pages, bool], tuple[pd.DataFrame, int]] + + +@dataclass(frozen=True) +class Measurement: + """Median measurements for one strategy and geometry mode.""" + + mode: str + strategy: str + conversions: int + wall_seconds: float + peak_mib: float + + +def _pages(row_count: int, page_size: int, *, spatial: bool) -> Pages: + """Build deterministic GeoJSON pages outside the measured region.""" + features: list[Feature] = [] + for index in range(row_count): + feature: Feature = { + "type": "Feature", + "id": f"feature-{index}", + "properties": { + "value": index / 10, + "name": f"site-{index % 100}", + "quality": {"approved": index % 2 == 0}, + }, + } + if spatial: + feature["geometry"] = { + "type": "Point", + "coordinates": [-125.0 + index / row_count, 25.0 + index / row_count], + } + features.append(feature) + return [ + features[start : start + page_size] for start in range(0, row_count, page_size) + ] + + +def _page_frames(pages: Pages, spatial: bool) -> tuple[pd.DataFrame, int]: + """Model the former implementation: shape every page, then concatenate.""" + frames = [ + _feature_frame(page, geopd=spatial, include_geometry=spatial) for page in pages + ] + return pd.concat(frames, ignore_index=True), len(frames) + + +def _raw_features(pages: Pages, spatial: bool) -> tuple[pd.DataFrame, int]: + """Run the current implementation: flatten pages, then shape once.""" + features = _combine_feature_pages(pages, row_cap=None) + return ( + _feature_frame(features, geopd=spatial, include_geometry=spatial), + 1, + ) + + +def _peak_rss_mib() -> float: + """Return this process's maximum resident set size in MiB.""" + peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + # macOS reports bytes; Linux and the other supported CI platforms report KiB. + return peak / (1024 * 1024) if sys.platform == "darwin" else peak / 1024 + + +def _run_worker( + name: str, + *, + spatial: bool, + row_count: int, + page_size: int, +) -> dict[str, float | int]: + """Run one isolated measurement and return its JSON payload.""" + command = [ + sys.executable, + str(Path(__file__).resolve()), + "--worker", + name, + "--rows", + str(row_count), + "--page-size", + str(page_size), + ] + if spatial: + command.append("--spatial") + output = subprocess.check_output(command, text=True) + return json.loads(output) + + +def _measure( + name: str, + *, + spatial: bool, + repeats: int, + expected_rows: int, + page_size: int, +) -> Measurement: + """Measure one strategy repeatedly in fresh processes.""" + results = [ + _run_worker( + name, + spatial=spatial, + row_count=expected_rows, + page_size=page_size, + ) + for _ in range(repeats) + ] + conversions = {int(item["conversions"]) for item in results} + row_counts = {int(item["rows"]) for item in results} + if len(conversions) != 1 or row_counts != {expected_rows}: + raise RuntimeError(f"{name} produced inconsistent output") + + return Measurement( + mode="spatial" if spatial else "nonspatial", + strategy=name, + conversions=conversions.pop(), + wall_seconds=statistics.median(float(item["wall_seconds"]) for item in results), + peak_mib=statistics.median(float(item["peak_mib"]) for item in results), + ) + + +def _ratio(current: float, baseline: float) -> float: + return current / baseline if baseline else float("inf") + + +def _print_report(measurements: list[Measurement]) -> None: + """Print a Markdown-friendly result table and informational guardrails.""" + print(f"Python {platform.python_version()} on {platform.platform()}") + print("\n| mode | strategy | conversions | wall (s) | peak RSS MiB |") + print("| --- | --- | ---: | ---: | ---: |") + for item in measurements: + print( + f"| {item.mode} | {item.strategy} | {item.conversions} | " + f"{item.wall_seconds:.4f} | {item.peak_mib:.2f} |" + ) + + print("\nInformational guardrails (raw / page-frame):") + for mode in ("spatial", "nonspatial"): + by_name = {item.strategy: item for item in measurements if item.mode == mode} + baseline = by_name["page_frames"] + current = by_name["raw_features"] + wall_ratio = _ratio(current.wall_seconds, baseline.wall_seconds) + peak_ratio = _ratio(current.peak_mib, baseline.peak_mib) + wall_status = "PASS" if wall_ratio <= 1.10 else "REVIEW" + peak_status = "PASS" if peak_ratio <= 1.25 else "REVIEW" + print( + f"- {mode}: wall {wall_ratio:.2f}x ({wall_status}, <=1.10x); " + f"peak {peak_ratio:.2f}x ({peak_status}, <=1.25x)" + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", type=int, default=10_000) + parser.add_argument("--page-size", type=int, default=1_000) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument( + "--worker", + choices=("page_frames", "raw_features"), + help=argparse.SUPPRESS, + ) + parser.add_argument("--spatial", action="store_true", help=argparse.SUPPRESS) + args = parser.parse_args() + if args.rows <= 0 or args.page_size <= 0 or args.repeats <= 0: + parser.error("--rows, --page-size, and --repeats must be positive") + + strategies: dict[str, Strategy] = { + "page_frames": _page_frames, + "raw_features": _raw_features, + } + if args.worker is not None: + pages = _pages(args.rows, args.page_size, spatial=args.spatial) + gc.collect() + started = time.perf_counter() + frame, conversions = strategies[args.worker](pages, args.spatial) + wall_seconds = time.perf_counter() - started + print( + json.dumps( + { + "rows": len(frame), + "conversions": conversions, + "wall_seconds": wall_seconds, + "peak_mib": _peak_rss_mib(), + } + ) + ) + return + + measurements: list[Measurement] = [] + for spatial in (True, False): + for name in strategies: + measurements.append( + _measure( + name, + spatial=spatial, + repeats=args.repeats, + expected_rows=args.rows, + page_size=args.page_size, + ) + ) + + _print_report(measurements) + + +if __name__ == "__main__": + main() diff --git a/dataretrieval/ogc/engine.py b/dataretrieval/ogc/engine.py index be63e69c..2a6b1132 100644 --- a/dataretrieval/ogc/engine.py +++ b/dataretrieval/ogc/engine.py @@ -51,7 +51,7 @@ _switch_arg_id, _switch_properties_id, ) -from dataretrieval.ogc.shaping import GEOPANDAS, _finalize_ogc, _get_resp_data +from dataretrieval.ogc.shaping import GEOPANDAS, _feature_frame, _finalize_ogc from dataretrieval.transport.fanout import FanOut, active_client from dataretrieval.transport.links import resolve_next_url from dataretrieval.transport.pagination import paginate @@ -100,7 +100,7 @@ def _next_req_url( # rather than ``numberReturned``: the main Water Data API reports # ``numberReturned`` but the NGWMN OGC API omits it, so trusting it would # refuse to follow a ``next`` link on a page that actually carries - # features (mirrors the same guard in :func:`_get_resp_data`). + # features (mirrors the same guard in :func:`_ogc_parse_response`). if not (body.get("features") or []): return None for link in body.get("links", []): @@ -146,27 +146,36 @@ async def _paginate( def _ogc_parse_response( resp: httpx.Response, - *, - geopd: bool, - include_geometry: bool = True, -) -> tuple[pd.DataFrame, str | None]: - """Parse one OGC API page: extract the DataFrame and the next-page URL. - - The parse strategy :func:`_walk_pages` hands to - :func:`_paginate`. Coerces falsy cursors (empty href, etc.) to - ``None`` so the paginate loop's ``while cursor is not None`` - terminates instead of spinning on a meaningless value. +) -> tuple[list[dict[str, Any]], str | None]: + """Parse one OGC page into raw features and its next-page cursor. + + Missing or null ``features`` is the service's empty-page shape. Other + container errors are rejected here so a malformed page is reported with + pagination context rather than failing later inside pandas or geopandas. + Feature members remain schema-tolerant: ``id``, ``properties``, and + ``geometry`` are optional and interpreted only during final shaping. """ body = resp.json() - return ( - _get_resp_data( - resp, - geopd=geopd, - include_geometry=include_geometry, - body=body, - ), - _next_req_url(resp, body=body) or None, - ) + raw_features = body.get("features") + if raw_features is None: + features: list[dict[str, Any]] = [] + elif not isinstance(raw_features, list): + raise ValueError("OGC response 'features' must be a list") + else: + features = [] + for index, feature in enumerate(raw_features): + if not isinstance(feature, dict): + raise ValueError(f"OGC feature at index {index} must be a mapping") + features.append(feature) + return features, _next_req_url(resp, body=body) or None + + +def _combine_feature_pages( + pages: list[list[dict[str, Any]]], row_cap: int | None +) -> list[dict[str, Any]]: + """Flatten feature pages in arrival order and apply the download cap.""" + features = [feature for page in pages for feature in page] + return features if row_cap is None else features[:row_cap] async def _walk_pages( @@ -177,36 +186,35 @@ async def _walk_pages( include_geometry: bool = True, row_cap: int | None = None, ) -> tuple[pd.DataFrame, httpx.Response]: - """ - Iterate paginated OGC API responses and aggregate them into one DataFrame. + """Aggregate raw OGC feature pages, then shape one completed chunk. - Thin wrapper that hands off to :func:`_paginate` with - OGC-specific strategies: pages are parsed via :func:`_get_resp_data` - (through :func:`_ogc_parse_response`) and the next-page cursor is the - URL from the response's ``links`` array (per :func:`_next_req_url`). + Thin wrapper around :func:`paginate` with OGC-specific strategies: each + page parser returns its raw ``features`` list and next-link cursor, the + paginator combines and caps those lists, and :func:`_feature_frame` + converts the completed chunk once. Parameters ---------- geopd : bool - Whether this request's pages use an active GeoDataFrame. + Whether the completed chunk uses an active GeoDataFrame. include_geometry : bool, optional - Whether plain-DataFrame pages retain raw coordinate lists. False for - nonspatial collections and requests using ``skip_geometry=True``. + Whether a plain completed-chunk DataFrame retains raw coordinate lists. + False for nonspatial collections and ``skip_geometry=True`` requests. req : httpx.Request The initial HTTP request to send. client : httpx.AsyncClient, optional - Caller-borrowed client; ``None`` defers client management to - :func:`_paginate`. + Caller-borrowed client; ``None`` lets :func:`paginate` manage one. row_cap : int, optional - Stop following pages once this many rows have accumulated and - truncate to exactly this many. ``None`` (default) walks every page. - An early-stop download bound only — the combined-result cap is - applied in :func:`~dataretrieval.ogc.shaping._finalize_ogc`. + Stop following pages once this many raw features have accumulated and + truncate to exactly this many before frame conversion. ``None`` + (default) walks every page. An early-stop download bound only — the + combined-result cap is applied in + :func:`~dataretrieval.ogc.shaping._finalize_ogc`. Returns ------- pd.DataFrame - A DataFrame containing the aggregated results from all pages. + The completed chunk frame built from the capped feature sequence. httpx.Response Aggregated response — initial-request URL (for query identity), final page's headers (so downstream sees current rate-limit @@ -226,16 +234,22 @@ async def _walk_pages( async def follow_up(cursor: str, sess: httpx.AsyncClient) -> httpx.Response: return await sess.request(method, cursor, headers=headers, content=content) - return await _paginate( + features, response = await paginate( req, - parse_response=functools.partial( - _ogc_parse_response, + parse_response=_ogc_parse_response, + follow_up=follow_up, + client=client if client is not None else active_client(), + raise_for_status=_raise_for_non_200, + row_cap=row_cap, + combine_pages=_combine_feature_pages, + ) + return ( + _feature_frame( + features, geopd=geopd, include_geometry=include_geometry, ), - follow_up=follow_up, - client=client, - row_cap=row_cap, + response, ) @@ -335,9 +349,9 @@ def get_ogc_data( convert_type = args.pop("convert_type", False) args = {k: v for k, v in args.items() if v is not None} - # Choose one semantic frame shape for the whole request. Every page and - # the all-empty finalizer receive these same values, so frame type never - # depends on whether a particular page happens to carry geometry. + # Choose one semantic frame shape for the whole request. Every completed + # chunk and the all-empty finalizer receive these same values, so frame type + # never depends on whether a particular page happens to carry geometry. include_geometry = spatial and not bool(args.get("skip_geometry", False)) geopd = GEOPANDAS and include_geometry diff --git a/dataretrieval/ogc/shaping.py b/dataretrieval/ogc/shaping.py index 760e6f41..2ad3a2af 100644 --- a/dataretrieval/ogc/shaping.py +++ b/dataretrieval/ogc/shaping.py @@ -52,10 +52,11 @@ def _empty_feature_frame( """Construct an empty feature frame for one request's geometry mode. ``geopd`` and ``include_geometry`` are selected once before pagination and - reused for every page and the final all-empty result. This keeps empty and - non-empty pages in the same frame family, so ordinary ``pd.concat`` is - sufficient. ``columns`` is supplied only when finalization has fetched the - collection schema; page-level empties intentionally remain schema-light. + reused when a completed chunk has no features and when the final combined + result is empty. This keeps empty and non-empty chunks in the same frame + family. ``columns`` is supplied only when finalization has fetched the + collection schema; completed-chunk empties intentionally remain + schema-light until then. """ result_columns = list(columns or []) if include_geometry: @@ -99,83 +100,38 @@ def _geo_feature_frame(features: list[dict[str, Any]]) -> pd.DataFrame: ) -def _get_resp_data( - resp: httpx.Response, +def _feature_frame( + features: list[dict[str, Any]], geopd: bool, *, include_geometry: bool = True, - body: dict[str, Any] | None = None, ) -> pd.DataFrame: - """ - Extracts and normalizes data from an HTTP response containing GeoJSON features. + """Convert one completed chunk's GeoJSON features into a frame. - Parameters - ---------- - resp : httpx.Response - The HTTP response object expected to contain a JSON body - with a "features" key. - geopd : bool - Whether this request's result contract uses an active GeoDataFrame. - Selected once before pagination rather than inferred from page content. - include_geometry : bool, optional - Whether plain-DataFrame results retain raw coordinate lists. False for - nonspatial collections and ``skip_geometry=True`` requests. - body : dict, optional - Pre-parsed JSON body for ``resp``. When provided, skips the - ``resp.json()`` call — useful when the caller has already - decoded the body for its own use (avoids a second parse pass). + The page walk retains feature dictionaries and calls this once after all + pages for the chunk have been combined and capped. ``geopd`` and + ``include_geometry`` are request-level decisions, so frame family never + depends on which page supplied a feature or whether its geometry is null. - Returns - ------- - gpd.GeoDataFrame or pd.DataFrame - A ``GeoDataFrame`` when ``geopd`` is True; otherwise a plain - ``DataFrame`` carrying the feature properties plus an ``id`` - column (always present, possibly all-None) and a ``geometry`` - column (coordinates list) when at least one feature includes - geometry. Returns an empty ``DataFrame`` when no features are - returned. - - Notes - ----- - The non-geopandas branch normalizes each feature's ``properties`` object, + The plain-DataFrame branch normalizes each feature's ``properties`` object, flattening nested dictionaries with an underscore separator, then adds the - top-level ``id`` and a ``geometry`` column containing the coordinates. The - ``id`` column is always added so the downstream collection-specific rename - works even when all IDs are missing; ``geometry`` is added only when - coordinates are present. Feature-level envelope fields are deliberately - excluded. + top-level ``id`` and optional coordinate-list ``geometry`` column. The + ``id`` column is always materialized so final shaping can rename it to the + collection-specific identifier. """ - if body is None: - body = resp.json() - # Key the empty-result short-circuit off ``features`` rather than - # ``numberReturned``: the main Water Data API reports ``numberReturned``, - # but the NGWMN OGC API omits it, so trusting it would discard pages that - # actually carry features. An absent/empty ``features`` is also the real - # schema-drift shape (a 200 with no features) — treat it as empty rather - # than crash with a ``KeyError`` downstream, which ``_paginate`` would - # mistake for a transient transport error. The request-level mode gives - # this page the same frame family as every non-empty page in the walk. - features = body.get("features") or [] if not features: return _empty_feature_frame(geopd, include_geometry=include_geometry) if not geopd: 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 collection-specific rename. df["id"] = [feature.get("id") for feature in features] if include_geometry: _attach_coordinates(df, features) return df - # A spatial request remains geospatial even when this particular page - # carries only null/missing geometries; changing frame family based on page - # contents makes pagination concat order-dependent. df = _geo_feature_frame(features) - # Mirror the plain branch's defensive ``f.get("id")`` so a feature missing - # a top-level id yields None rather than a KeyError. - df["id"] = [f.get("id") for f in features] + df["id"] = [feature.get("id") for feature in features] return df[["id"] + [col for col in df.columns if col != "id"]] @@ -190,8 +146,8 @@ def _deal_with_empty( ) -> pd.DataFrame: """Apply the collection schema when an entire result is empty. - Page-level empties already carry the request's frame family. The complete - column list is available only from explicit ``properties`` or the + A completed empty chunk already carries the request's frame family. The + complete column list is available only from explicit ``properties`` or the collection schema, so construct the fully shaped empty once here rather than fetching schema during pagination. """ @@ -387,7 +343,7 @@ def _finalize_ogc( ``max_rows`` is applied here (after dedup/sort, on the *combined* frame) rather than only per-chunk, so a chunked call's total is bounded to exactly ``max_rows`` and a resumed call honors the cap too. The - per-page ``row_cap`` bound in the engine is only an early-stop download + chunk page-walk ``row_cap`` in the engine is only an early-stop download bound. ``base_url`` is required and captured with the finalizer so resumed calls query the same API's schema when their combined result is empty. """ diff --git a/dataretrieval/transport/pagination.py b/dataretrieval/transport/pagination.py index 523463c6..0023cd89 100644 --- a/dataretrieval/transport/pagination.py +++ b/dataretrieval/transport/pagination.py @@ -9,7 +9,7 @@ import logging from collections.abc import AsyncIterator, Awaitable, Callable from contextlib import asynccontextmanager -from typing import Any, TypeVar +from typing import Any, TypeVar, overload import httpx import pandas as pd @@ -67,6 +67,19 @@ def paginated_failure_message(pages_collected: int, cause: BaseException) -> str ) +def _combine_frame_pages( + pages: list[pd.DataFrame], row_cap: int | None +) -> pd.DataFrame: + """Combine DataFrame pages using the paginator's established behavior.""" + result = pd.concat(pages, ignore_index=True) + return result if row_cap is None else result.head(row_cap) + + +_Page = TypeVar("_Page") +_Result = TypeVar("_Result") + + +@overload async def paginate( initial_req: httpx.Request, *, @@ -75,25 +88,52 @@ async def paginate( raise_for_status: Callable[[httpx.Response], None], client: httpx.AsyncClient | None = None, row_cap: int | None = None, -) -> tuple[pd.DataFrame, httpx.Response]: + combine_pages: None = None, +) -> tuple[pd.DataFrame, httpx.Response]: ... + + +@overload +async def paginate( + initial_req: httpx.Request, + *, + parse_response: Callable[[httpx.Response], tuple[_Page, _Cursor | None]], + follow_up: Callable[[_Cursor, httpx.AsyncClient], Awaitable[httpx.Response]], + raise_for_status: Callable[[httpx.Response], None], + combine_pages: Callable[[list[_Page], int | None], _Result], + client: httpx.AsyncClient | None = None, + row_cap: int | None = None, +) -> tuple[_Result, httpx.Response]: ... + + +async def paginate( + initial_req: httpx.Request, + *, + parse_response: Callable[[httpx.Response], tuple[Any, Any | None]], + follow_up: Callable[[Any, httpx.AsyncClient], Awaitable[httpx.Response]], + raise_for_status: Callable[[httpx.Response], None], + client: httpx.AsyncClient | None = None, + row_cap: int | None = None, + combine_pages: Callable[[list[Any], int | None], Any] | None = None, +) -> tuple[Any, httpx.Response]: """Fetch and combine pages until the injected parser returns no cursor. - The service adapter supplies response parsing, cursor following, and status - mapping. This loop owns client lifecycle, repeated-cursor protection, - optional row capping, progress updates, failure wrapping, and response - metadata aggregation. + The service adapter supplies response parsing, cursor following, status + mapping, and optionally how its natural page payloads combine. DataFrame + pages retain the established concatenation default. This loop owns client + lifecycle, repeated-cursor protection, optional row capping, progress + updates, failure wrapping, and response metadata aggregation. """ logger.debug("Requesting: %s", initial_req.url) reporter = _progress.current() - def report_page(page: httpx.Response, frame: pd.DataFrame) -> None: + def report_page(page: httpx.Response, payload: Any) -> None: note_progress() # a walk still delivering pages is not stalled if reporter is not None: reporter.set_rate_remaining( page.headers.get(_QUOTA_HEADER), limit=page.headers.get("x-ratelimit-limit"), ) - reporter.add_page(rows=len(frame)) + reporter.add_page(rows=len(payload)) async with _client_for(client) as session: response = await session.send(initial_req) @@ -102,15 +142,15 @@ def report_page(page: httpx.Response, frame: pd.DataFrame) -> None: total_elapsed = _safe_elapsed(response) try: - frame, cursor = parse_response(response) + payload, cursor = parse_response(response) except Exception as exc: # noqa: BLE001 logger.warning("Initial response parse failed.") raise DataRetrievalError(paginated_failure_message(0, exc)) from exc - frames = [frame] - nrows = len(frame) + pages = [payload] + nrows = len(payload) seen: set[Any] = set() - report_page(response, frame) + report_page(response, payload) while ( cursor is not None @@ -121,17 +161,17 @@ def report_page(page: httpx.Response, frame: pd.DataFrame) -> None: try: response = await follow_up(cursor, session) raise_for_status(response) - frame, cursor = parse_response(response) - frames.append(frame) - nrows += len(frame) + payload, cursor = parse_response(response) + pages.append(payload) + nrows += len(payload) total_elapsed += _safe_elapsed(response) - report_page(response, frame) + report_page(response, payload) except Exception as exc: # noqa: BLE001 logger.warning( "Request failed at cursor %r. Data download interrupted.", cursor ) raise DataRetrievalError( - paginated_failure_message(len(frames), exc) + paginated_failure_message(len(pages), exc) ) from exc final_response = _merge_response( @@ -139,10 +179,8 @@ def report_page(page: httpx.Response, frame: pd.DataFrame) -> None: headers_from=response, elapsed=total_elapsed, ) - result = pd.concat(frames, ignore_index=True) - if row_cap is not None: - result = result.head(row_cap) - return result, final_response + combine = _combine_frame_pages if combine_pages is None else combine_pages + return combine(pages, row_cap), final_response def run_paginated( diff --git a/dataretrieval/waterdata/stats.py b/dataretrieval/waterdata/stats.py index 9239f324..f3b177fa 100644 --- a/dataretrieval/waterdata/stats.py +++ b/dataretrieval/waterdata/stats.py @@ -58,13 +58,13 @@ def _handle_nesting( Notes ----- The non-geopandas branch uses the same schema-aware extraction as - :func:`engine._get_resp_data`: it builds the per-feature outer frame - directly from each feature's ``properties`` (minus the nested - ``data`` field, which is unrolled separately below via the - ``record_path`` json_normalize), then adds ``geometry`` only when - present. Unlike :func:`engine._get_resp_data`, no top-level ``id`` - column is added — stats features don't carry one, so this matches the - geopandas branch. Skipping the GeoJSON envelope keeps newly-added + :func:`~dataretrieval.ogc.shaping._feature_frame`: it builds the + per-feature outer frame directly from each feature's ``properties`` (minus + the nested ``data`` field, which is unrolled separately below via the + ``record_path`` json_normalize), then adds ``geometry`` only when present. + Unlike :func:`~dataretrieval.ogc.shaping._feature_frame`, no top-level + ``id`` column is added — stats features don't carry one, so this matches + the geopandas branch. Skipping the GeoJSON envelope keeps newly-added fields like ``geometry.type`` from leaking into the result. """ if body is None: diff --git a/tests/ogc_live_test.py b/tests/ogc_live_test.py new file mode 100644 index 00000000..86ebfa72 --- /dev/null +++ b/tests/ogc_live_test.py @@ -0,0 +1,298 @@ +"""Opt-in live conformance checks for OGC and shared pagination adapters. + +These tests deliberately use small, fixed queries. They are deselected by the +project's default ``-m 'not live'`` configuration and run only in the scheduled +live workflow or explicitly with ``pytest -m live``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import httpx +import pandas as pd +import pytest +from geopandas import GeoDataFrame + +from dataretrieval import ngwmn, parallel_chunks, waterdata, wateruse + +if TYPE_CHECKING: + from dataretrieval.utils import BaseMetadata + +pytestmark = pytest.mark.live + +_WATERDATA_SITES = ["USGS-01646500", "USGS-05427718"] +_NGWMN_SITES = ["USGS-272838082142201", "USGS-404159100494601"] + + +def _assert_spatial(frame: pd.DataFrame) -> None: + assert isinstance(frame, GeoDataFrame) + assert frame.geometry.name == "geometry" + assert frame.crs is not None + assert frame.crs.to_epsg() == 4326 + + +def _assert_metadata(metadata: BaseMetadata) -> None: + assert metadata.url + assert metadata.query_time.total_seconds() >= 0 + + +def _assert_serial_parallel_equal( + serial: pd.DataFrame, + parallel: pd.DataFrame, + *, + id_column: str, +) -> None: + """Compare fan-out values without conflating equivalent null sentinels.""" + assert type(serial) is type(parallel) + assert list(serial.columns) == list(parallel.columns) + + left = serial.sort_values(id_column).reset_index(drop=True) + right = parallel.sort_values(id_column).reset_index(drop=True) + if isinstance(left, GeoDataFrame): + assert left.crs == right.crs + assert left.geometry.equals(right.geometry) + left = left.drop(columns=left.geometry.name) + right = right.drop(columns=right.geometry.name) + + left = left.astype(object).where(left.notna(), None) + right = right.astype(object).where(right.notna(), None) + pd.testing.assert_frame_equal(left, right, check_dtype=False) + + +def _record_live_requests(monkeypatch) -> list[httpx.Request]: + """Capture real outgoing requests while preserving the live transport.""" + requests: list[httpx.Request] = [] + original_send = httpx.AsyncClient.send + + async def recording_send(self, request, *args, **kwargs): + requests.append(request) + return await original_send(self, request, *args, **kwargs) + + monkeypatch.setattr(httpx.AsyncClient, "send", recording_send) + return requests + + +def test_waterdata_spatial_and_skip_geometry_live(): + """The same spatial collection honors both requested frame families.""" + spatial, spatial_md = waterdata.get_monitoring_locations( + monitoring_location_id=_WATERDATA_SITES[0] + ) + tabular, tabular_md = waterdata.get_monitoring_locations( + monitoring_location_id=_WATERDATA_SITES[0], skip_geometry=True + ) + + assert len(spatial) == len(tabular) == 1 + _assert_spatial(spatial) + assert type(tabular) is pd.DataFrame + assert "geometry" not in tabular.columns + _assert_metadata(spatial_md) + _assert_metadata(tabular_md) + + +def test_waterdata_nonspatial_reference_live(): + """Reference collections stay plain DataFrames on the raw-feature path.""" + frame, metadata = waterdata.get_reference_table( + "parameter-codes", query={"id": "00060"} + ) + + assert type(frame) is pd.DataFrame + assert frame["parameter_code"].tolist() == ["00060"] + assert "geometry" not in frame.columns + _assert_metadata(metadata) + + +def test_waterdata_empty_spatial_schema_live(): + """An empty spatial result is schema-complete and remains geospatial.""" + frame, metadata = waterdata.get_monitoring_locations( + monitoring_location_id="USGS-999999999999999" + ) + + assert frame.empty + _assert_spatial(frame) + assert "monitoring_location_id" in frame.columns + _assert_metadata(metadata) + + +def test_waterdata_pagination_and_max_rows_live(): + """A one-row page is followed once, then the raw-feature cap stops paging.""" + frame, metadata = waterdata.get_daily( + monitoring_location_id=_WATERDATA_SITES[0], + parameter_code="00060", + time="2024-01-01/2024-01-03", + limit=1, + max_rows=2, + ) + + assert len(frame) == 2 + _assert_spatial(frame) + _assert_metadata(metadata) + + +def test_waterdata_cql_post_parallel_chunks_live(monkeypatch): + """A serial CQL2 POST and two-chunk fan-out return equivalent frames.""" + requests = _record_live_requests(monkeypatch) + serial, _ = waterdata.get_monitoring_locations( + monitoring_location_id=_WATERDATA_SITES + ) + serial_item_requests = [ + request + for request in requests + if "/collections/monitoring-locations/items" in str(request.url) + ] + assert len(serial_item_requests) == 1 + assert serial_item_requests[0].method == "POST" + requests.clear() + + monkeypatch.setenv("API_USGS_CONCURRENT", "2") + with parallel_chunks(2): + parallel, metadata = waterdata.get_monitoring_locations( + monitoring_location_id=_WATERDATA_SITES + ) + + item_requests = [ + request + for request in requests + if "/collections/monitoring-locations/items" in str(request.url) + ] + assert len(item_requests) == 2 + _assert_serial_parallel_equal( + serial, + parallel, + id_column="monitoring_location_id", + ) + _assert_spatial(parallel) + _assert_metadata(metadata) + + +def test_ngwmn_spatial_parallel_chunks_live(monkeypatch): + """NGWMN site fan-out equals the serial completed spatial frame.""" + serial, _ = ngwmn.get_sites(monitoring_location_id=_NGWMN_SITES) + + requests = _record_live_requests(monkeypatch) + monkeypatch.setenv("API_USGS_CONCURRENT", "2") + with parallel_chunks(2): + parallel, metadata = ngwmn.get_sites(monitoring_location_id=_NGWMN_SITES) + + item_requests = [ + request + for request in requests + if "/collections/sites/items" in str(request.url) + ] + assert len(item_requests) == 2 + assert all(request.method == "GET" for request in item_requests) + _assert_serial_parallel_equal( + serial, + parallel, + id_column="monitoring_location_id", + ) + _assert_spatial(parallel) + _assert_metadata(metadata) + + +def test_ngwmn_paginates_providers_live(): + """A partial second providers page exercises NGWMN next links.""" + frame, metadata = ngwmn.get_providers(state="WI", limit=35) + + assert len(frame) > 35 + assert type(frame) is pd.DataFrame + assert "geometry" not in frame.columns + _assert_metadata(metadata) + + +def test_ngwmn_providers_are_nonspatial_live(): + """An exact NGWMN provider query stays a plain DataFrame.""" + frame, metadata = ngwmn.get_providers(agency_code="USGS", limit=100) + + assert not frame.empty + assert type(frame) is pd.DataFrame + assert "geometry" not in frame.columns + _assert_metadata(metadata) + + +def test_ngwmn_geometry_free_observations_live(): + """Observation features without geometry remain a plain DataFrame.""" + frame, metadata = ngwmn.get_water_level( + monitoring_location_id=_NGWMN_SITES[0], + limit=10000, + ) + + assert not frame.empty + assert type(frame) is pd.DataFrame + assert "geometry" not in frame.columns + _assert_metadata(metadata) + + +def test_ngwmn_empty_spatial_live(): + """An empty NGWMN site result preserves its spatial frame contract.""" + frame, metadata = ngwmn.get_sites(monitoring_location_id="USGS-999999999999999") + + assert frame.empty + _assert_spatial(frame) + assert "monitoring_location_id" in frame.columns + _assert_metadata(metadata) + + +def test_statistics_shared_paginator_live(): + """Statistics retains the shared paginator's default DataFrame path.""" + frame, metadata = waterdata.get_stats_por( + monitoring_location_id=_WATERDATA_SITES[0], + parameter_code="00060", + computation_type="arithmetic_mean", + normal_type="MOY", + page_size=1, + ) + + assert not frame.empty + assert isinstance(frame, pd.DataFrame) + _assert_metadata(metadata) + + +def test_ratings_shared_paginator_live(): + """Ratings pages STAC search frames, then retrieves one small asset.""" + ratings = waterdata.get_ratings( + monitoring_location_id="USGS-01104475", + file_type="exsa", + limit=1, + ) + + assert ratings + assert all(type(frame) is pd.DataFrame for frame in ratings.values()) + assert all(not frame.empty for frame in ratings.values()) + + +def test_wateruse_shared_paginator_live(): + """NWDC Water Use retains the shared paginator's DataFrame combiner.""" + frame, metadata = wateruse.get_wateruse( + model="wu-public-supply-wd", + variable="pswdtot", + huc="010900020502", + time_resolution="monthly", + start_date="2020-01", + end_date="2020-12", + limit=1, + ) + + assert not frame.empty + assert type(frame) is pd.DataFrame + assert "huc12_id" in frame.columns + _assert_metadata(metadata) + + +def test_wateruse_list_fan_out_live(monkeypatch): + """NWDC Water Use combines two list-shaped location chunks.""" + monkeypatch.setenv("API_USGS_CONCURRENT", "2") + frame, metadata = wateruse.get_wateruse( + model="wu-public-supply-wd", + variable="pswdtot", + state=["RI", "DE"], + time_resolution="monthly", + start_date="2020-01", + end_date="2020-01", + limit=600, + ) + + assert not frame.empty + assert type(frame) is pd.DataFrame + assert "huc12_id" in frame.columns + _assert_metadata(metadata) diff --git a/tests/transport_test.py b/tests/transport_test.py index dcd4b5c0..7ab933a0 100644 --- a/tests/transport_test.py +++ b/tests/transport_test.py @@ -95,6 +95,51 @@ async def follow(cursor: str, session: httpx.AsyncClient) -> httpx.Response: assert client.get.await_count == 1 +def test_paginate_combines_adapter_page_payloads() -> None: + """An adapter may combine its natural sized page type without transport + pretending that payload is a DataFrame. + """ + first = _response(url="https://example.test/page/1") + second = _response(url="https://example.test/page/2") + client = mock.AsyncMock(spec=httpx.AsyncClient) + client.send.return_value = first + client.get.return_value = second + + def parse(response: httpx.Response) -> tuple[list[int], str | None]: + if response is first: + return [1, 2], "next" + return [3, 4], None + + async def follow(cursor: str, session: httpx.AsyncClient) -> httpx.Response: + assert cursor == "next" + return await session.get(str(second.url)) + + def combine_pages(pages: list[list[int]], row_cap: int | None) -> list[int]: + combined = list(itertools.chain.from_iterable(pages)) + return combined if row_cap is None else combined[:row_cap] + + reporter = mock.Mock() + with mock.patch( + "dataretrieval.transport.pagination._progress.current", + return_value=reporter, + ): + values, _ = asyncio.run( + paginate( + httpx.Request("GET", first.url), + parse_response=parse, + follow_up=follow, + raise_for_status=_raise_for_status, + client=client, + row_cap=3, + combine_pages=combine_pages, + ) + ) + + assert values == [1, 2, 3] + assert reporter.add_page.call_args_list == [mock.call(rows=2), mock.call(rows=2)] + assert client.get.await_count == 1 + + def test_retry_sync_retries_transient_then_succeeds(monkeypatch) -> None: attempts = 0 slept: list[float] = [] diff --git a/tests/waterdata_filters_test.py b/tests/waterdata_filters_test.py index ff6c6f25..3be97d51 100644 --- a/tests/waterdata_filters_test.py +++ b/tests/waterdata_filters_test.py @@ -231,10 +231,9 @@ async def fake_walk_pages(*_args, **_kwargs): def test_empty_chunks_do_not_downgrade_geodataframe(): """A mix of empty and non-empty chunk responses must not - downgrade a GeoDataFrame-typed result to a plain DataFrame. - ``_get_resp_data`` returns ``pd.DataFrame()`` on empty responses, - which would otherwise strip geometry/CRS from the concatenated - output.""" + downgrade a GeoDataFrame-typed result to a plain DataFrame. Fan-out must + ignore an empty plain chunk frame when real geospatial chunks are present, + preserving geometry and CRS in the combined output.""" pytest.importorskip("geopandas") import geopandas as gpd from shapely.geometry import Point diff --git a/tests/waterdata_test.py b/tests/waterdata_test.py index e3436925..2c67137c 100644 --- a/tests/waterdata_test.py +++ b/tests/waterdata_test.py @@ -13,6 +13,7 @@ import pytest from pandas import DataFrame +from dataretrieval.exceptions import DataRetrievalError from dataretrieval.ogc.requests import ( _check_monitoring_location_id, _normalize_str_iterable, @@ -661,6 +662,45 @@ def test_samples_service_profile_routes_to_its_endpoint( # --- daily / continuous ------------------------------------------------------ +@pytest.mark.parametrize( + ("features", "message"), + [ + ({"not": "a list"}, "'features' must be a list"), + (["not-a-feature"], "feature at index 0 must be a mapping"), + ], +) +def test_get_daily_rejects_malformed_feature_pages(httpx_mock, features, message): + """Malformed page structure fails with page context, not a pandas error.""" + _mock_items( + httpx_mock, + "daily", + body={"features": features, "links": []}, + ) + + with pytest.raises(DataRetrievalError, match=message): + get_daily(monitoring_location_id="USGS-05427718") + + +def test_get_daily_reports_malformed_later_page(httpx_mock): + """Structural failure on page two reports the completed page count.""" + next_url = f"{_OGC_BASE}/collections/daily/items?cursor=bad-page" + first = _fixture("daily") + first["features"] = first["features"][:1] + first["links"] = [{"rel": "next", "href": next_url}] + _mock_items(httpx_mock, "daily", body=first) + _mock_items( + httpx_mock, + "daily", + body={"features": ["not-a-feature"], "links": []}, + ) + + with pytest.raises( + DataRetrievalError, + match=r"after collecting 1 page\(s\).*feature at index 0 must be a mapping", + ): + get_daily(monitoring_location_id="USGS-05427718") + + def test_get_daily(httpx_mock): """A daily query returns tidy rows with the collection id renamed to ``daily_id`` and moved last, dates as ``date`` objects, values numeric.""" @@ -785,6 +825,49 @@ def test_get_daily_empty_no_geometry(httpx_mock): assert {"monitoring_location_id", "value"}.issubset(df.columns) +@pytest.mark.parametrize("body", [{}, {"features": None, "links": []}]) +def test_get_daily_empty_spatial_envelopes(httpx_mock, body): + """Missing and null feature lists retain the spatial collection contract.""" + geopandas = pytest.importorskip("geopandas") + httpx_mock.add_response( + method="GET", + url=_schema_url("daily"), + json={ + "properties": { + "id": {}, + "monitoring_location_id": {}, + "value": {}, + } + }, + ) + _mock_items(httpx_mock, "daily", body=body) + + df, _ = get_daily(monitoring_location_id="USGS-NOT-FOUND") + + assert isinstance(df, geopandas.GeoDataFrame) + assert df.empty + assert df.crs == "EPSG:4326" + assert {"monitoring_location_id", "value", "geometry"}.issubset(df.columns) + + +def test_get_daily_empty_terminal_page_preserves_spatial_result(httpx_mock): + """A terminal empty page cannot downgrade the completed spatial chunk.""" + geopandas = pytest.importorskip("geopandas") + next_url = f"{_OGC_BASE}/collections/daily/items?cursor=terminal-empty" + first = _fixture("daily") + first["features"] = first["features"][:1] + first["links"] = [{"rel": "next", "href": next_url}] + _mock_items(httpx_mock, "daily", body=first) + _mock_items(httpx_mock, "daily", body={"features": [], "links": []}) + + df, _ = get_daily(monitoring_location_id="USGS-05427718") + + assert isinstance(df, geopandas.GeoDataFrame) + assert len(df) == 1 + assert df.crs == "EPSG:4326" + assert df.geometry.name == "geometry" + + def test_get_continuous(httpx_mock): """Continuous observations are timestamped (not date-only), so ``time`` comes back as a UTC-aware datetime column.""" diff --git a/tests/waterdata_utils_test.py b/tests/waterdata_utils_test.py index d76fb7d1..0a506b8a 100644 --- a/tests/waterdata_utils_test.py +++ b/tests/waterdata_utils_test.py @@ -34,7 +34,7 @@ from dataretrieval.ogc.shaping import ( _arrange_cols, _deal_with_empty, - _get_resp_data, + _feature_frame, _to_snake_case, ) from dataretrieval.ogc.shaping import _finalize_ogc as _ogc_finalize @@ -360,16 +360,9 @@ def test_walk_pages_wraps_initial_page_parse_error(): assert isinstance(excinfo.value.__cause__, json.JSONDecodeError) -def test_get_resp_data_handles_missing_features_key(): - """Regression: a 200 with ``numberReturned > 0`` but no - ``features`` key (real schema-drift shape) used to crash - ``_get_resp_data`` with ``KeyError`` — wrapped downstream by - ``_paginate`` as a generic transport error. ``_handle_nesting`` - was already hardened against this; ``_get_resp_data`` now mirrors - that defensiveness and returns an empty frame instead.""" - resp = mock.Mock() - resp.json.return_value = {"numberReturned": 1, "links": []} - df = _get_resp_data(resp, geopd=False) +def test_feature_frame_handles_empty_feature_list(): + """An empty completed chunk produces a schema-light frame for finalization.""" + df = _feature_frame([], geopd=False) assert df.empty assert isinstance(df, pd.DataFrame) @@ -377,7 +370,7 @@ def test_get_resp_data_handles_missing_features_key(): def test_next_req_url_follows_link_without_number_returned(): """The NGWMN OGC API omits ``numberReturned`` from its page envelope, so ``_next_req_url`` keys the ``next`` link off ``features`` (mirroring - ``_get_resp_data``) rather than that count -- otherwise a page that carries + ``_feature_frame``) rather than that count -- otherwise a page that carries features but no count stops pagination after page 1 and silently truncates every multi-page result. A page that carries features still follows its ``next`` link even when ``numberReturned`` is absent.""" @@ -638,8 +631,7 @@ def test_empty_result_matches_a_non_empty_geodataframe(): """An empty geospatial result remains usable like a non-empty one.""" import geopandas as gpd - response = _resp_ok([]) - page = _get_resp_data(response, geopd=True) + page = _feature_frame([], geopd=True) template = pd.concat([page], ignore_index=True) out = _deal_with_empty( template, @@ -668,8 +660,8 @@ def test_empty_result_matches_a_non_empty_geodataframe(): assert isinstance(combined, gpd.GeoDataFrame) and combined.crs == "EPSG:4326" -def test_get_resp_data_empty_preserves_geopd_type(): - """Same as the stats-side preservation: ``_get_resp_data``'s +def test_feature_frame_empty_preserves_geopd_type(): + """Same as the stats-side preservation: ``_feature_frame``'s ``numberReturned == 0`` short-circuit must return a ``GeoDataFrame`` (not a plain ``DataFrame``) when geopd is True, so paginating across a sparse intermediate page doesn't downgrade @@ -681,54 +673,47 @@ class _Sentinel: fake_gpd.GeoDataFrame = lambda *a, **kw: _Sentinel() - resp = mock.MagicMock() - resp.json.return_value = {"numberReturned": 0, "features": [], "links": []} - # ``_get_resp_data`` resolves ``gpd`` from the shaping namespace -- patch + # ``_feature_frame`` resolves ``gpd`` from the shaping namespace -- patch # it there, not in ``utils``. with mock.patch.object(_shaping_module, "gpd", fake_gpd, create=True): - result = _get_resp_data(resp, geopd=True) + result = _feature_frame([], geopd=True) assert isinstance(result, _Sentinel) @pytest.mark.skipif(not _shaping_module.GEOPANDAS, reason="requires geopandas") -def test_get_resp_data_keeps_one_geospatial_type_when_geometry_is_missing(): +def test_feature_frame_keeps_one_geospatial_type_when_geometry_is_missing(): """A spatial request cannot change frame family based on page contents.""" import geopandas as gpd - empty = _get_resp_data(_resp_ok([]), geopd=True) - missing = _get_resp_data( - _resp_ok([{"id": "1", "properties": {"value": 1}}]), + empty = _feature_frame([], geopd=True) + missing = _feature_frame( + [{"id": "1", "properties": {"value": 1}}], geopd=True, ) assert isinstance(empty, gpd.GeoDataFrame) assert isinstance(missing, gpd.GeoDataFrame) assert missing.geometry.isna().all() - for pages in ([empty, missing], [missing, empty]): - combined = pd.concat(pages, ignore_index=True) - assert isinstance(combined, gpd.GeoDataFrame) - assert combined.geometry.name == "geometry" - assert combined.crs == "EPSG:4326" + assert missing.geometry.name == "geometry" + assert missing.crs == "EPSG:4326" -def test_get_resp_data_attaches_wgs84_crs(): +def test_feature_frame_attaches_wgs84_crs(): """A geometry-bearing Water Data page should come back tagged as EPSG:4326 (the CRS the coordinates are published in), so callers can run ``to_crs`` / spatial joins without first patching in a CRS by hand. Regression for the ``.crs is None`` reported in issue #342.""" geopandas = pytest.importorskip("geopandas") - resp = _resp_ok( - [ - { - "type": "Feature", - "id": "USGS-01", - "geometry": {"type": "Point", "coordinates": [-76.5, 39.2]}, - "properties": {"monitoring_location_id": "USGS-01"}, - } - ] - ) - df = _get_resp_data(resp, geopd=True) + features = [ + { + "type": "Feature", + "id": "USGS-01", + "geometry": {"type": "Point", "coordinates": [-76.5, 39.2]}, + "properties": {"monitoring_location_id": "USGS-01"}, + } + ] + df = _feature_frame(features, geopd=True) assert isinstance(df, geopandas.GeoDataFrame) assert df.crs == "EPSG:4326" @@ -772,38 +757,31 @@ def test_handle_nesting_tolerates_missing_features_key(): assert df.empty -def test_get_resp_data_always_materializes_id_column(): - """``_get_resp_data`` must always materialize the ``id`` column +def test_feature_frame_always_materializes_id_column(): + """``_feature_frame`` must always materialize the ``id`` column (NaN-filled when no feature carries one) so the downstream ``_arrange_cols`` rename to the collection-specific output_id (``daily_id``, ``channel_measurements_id``, etc.) isn't a silent no-op.""" - resp = mock.MagicMock() - resp.json.return_value = { - "numberReturned": 2, - "features": [ - {"properties": {"val": "a"}}, # no top-level id - {"properties": {"val": "b"}}, # ditto - ], - } - df = _get_resp_data(resp, geopd=False) + features = [ + {"properties": {"val": "a"}}, # no top-level id + {"properties": {"val": "b"}}, # ditto + ] + df = _feature_frame(features, geopd=False) assert "id" in df.columns assert df["id"].isna().all() -def test_get_resp_data_flattens_nested_properties(): +def test_feature_frame_flattens_nested_properties(): """Nested GeoJSON properties keep underscore-separated column names.""" - resp = mock.MagicMock() - resp.json.return_value = { - "features": [ - { - "id": "feature-1", - "properties": {"station": {"code": "A"}, "value": 1}, - } - ] - } + features = [ + { + "id": "feature-1", + "properties": {"station": {"code": "A"}, "value": 1}, + } + ] - df = _get_resp_data(resp, geopd=False) + df = _feature_frame(features, geopd=False) assert df.to_dict(orient="records") == [ {"value": 1, "station_code": "A", "id": "feature-1"} From 896e3770b5e1b73e4dc7c4aa8354266ce8d10cf2 Mon Sep 17 00:00:00 2001 From: thodson-usgs Date: Thu, 13 Aug 2026 09:58:01 -0500 Subject: [PATCH 5/5] refactor(ogc): simplify raw pagination plumbing --- benchmarks/ogc_raw_feature_pagination.py | 2 + dataretrieval/ogc/engine.py | 111 +++++------------------ dataretrieval/transport/pagination.py | 45 +++------ tests/waterdata_filters_test.py | 2 +- tests/waterdata_progress_test.py | 4 +- 5 files changed, 38 insertions(+), 126 deletions(-) diff --git a/benchmarks/ogc_raw_feature_pagination.py b/benchmarks/ogc_raw_feature_pagination.py index 0829ecb3..7564c88e 100644 --- a/benchmarks/ogc_raw_feature_pagination.py +++ b/benchmarks/ogc_raw_feature_pagination.py @@ -31,6 +31,8 @@ import pandas as pd +# Deliberate white-box benchmark: these private imports measure the experiment's +# conversion boundary; they are not supported adapter extension points. from dataretrieval.ogc.engine import _combine_feature_pages from dataretrieval.ogc.shaping import _feature_frame diff --git a/dataretrieval/ogc/engine.py b/dataretrieval/ogc/engine.py index 2a6b1132..4d5ec05b 100644 --- a/dataretrieval/ogc/engine.py +++ b/dataretrieval/ogc/engine.py @@ -26,11 +26,8 @@ import functools import logging -from collections.abc import ( - Awaitable, - Callable, -) -from typing import TYPE_CHECKING, Any, TypeVar +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, cast import httpx import pandas as pd @@ -120,30 +117,6 @@ def _next_req_url( return None -_Cursor = TypeVar("_Cursor") - - -async def _paginate( - initial_req: httpx.Request, - *, - parse_response: Callable[[httpx.Response], tuple[pd.DataFrame, _Cursor | None]], - follow_up: Callable[[_Cursor, httpx.AsyncClient], Awaitable[httpx.Response]], - client: httpx.AsyncClient | None = None, - raise_for_status: Callable[[httpx.Response], None] = _raise_for_non_200, - row_cap: int | None = None, -) -> tuple[pd.DataFrame, httpx.Response]: - """Compatibility wrapper around collection-neutral cursor pagination.""" - session = client if client is not None else active_client() - return await paginate( - initial_req, - parse_response=parse_response, - follow_up=follow_up, - client=session, - raise_for_status=raise_for_status, - row_cap=row_cap, - ) - - def _ogc_parse_response( resp: httpx.Response, ) -> tuple[list[dict[str, Any]], str | None]: @@ -156,18 +129,17 @@ def _ogc_parse_response( ``geometry`` are optional and interpreted only during final shaping. """ body = resp.json() - raw_features = body.get("features") - if raw_features is None: - features: list[dict[str, Any]] = [] - elif not isinstance(raw_features, list): - raise ValueError("OGC response 'features' must be a list") - else: + features = body.get("features") + if features is None: features = [] - for index, feature in enumerate(raw_features): - if not isinstance(feature, dict): - raise ValueError(f"OGC feature at index {index} must be a mapping") - features.append(feature) - return features, _next_req_url(resp, body=body) or None + if not isinstance(features, list): + raise ValueError("OGC response 'features' must be a list") + for index, feature in enumerate(features): + if not isinstance(feature, dict): + raise ValueError(f"OGC feature at index {index} must be a mapping") + return cast("list[dict[str, Any]]", features), _next_req_url( + resp, body=body + ) or None def _combine_feature_pages( @@ -179,53 +151,19 @@ def _combine_feature_pages( async def _walk_pages( - geopd: bool, req: httpx.Request, client: httpx.AsyncClient | None = None, *, + geopd: bool | None = None, include_geometry: bool = True, row_cap: int | None = None, ) -> tuple[pd.DataFrame, httpx.Response]: - """Aggregate raw OGC feature pages, then shape one completed chunk. - - Thin wrapper around :func:`paginate` with OGC-specific strategies: each - page parser returns its raw ``features`` list and next-link cursor, the - paginator combines and caps those lists, and :func:`_feature_frame` - converts the completed chunk once. + """Aggregate raw OGC pages, then shape one completed chunk. - Parameters - ---------- - geopd : bool - Whether the completed chunk uses an active GeoDataFrame. - include_geometry : bool, optional - Whether a plain completed-chunk DataFrame retains raw coordinate lists. - False for nonspatial collections and ``skip_geometry=True`` requests. - req : httpx.Request - The initial HTTP request to send. - client : httpx.AsyncClient, optional - Caller-borrowed client; ``None`` lets :func:`paginate` manage one. - row_cap : int, optional - Stop following pages once this many raw features have accumulated and - truncate to exactly this many before frame conversion. ``None`` - (default) walks every page. An early-stop download bound only — the - combined-result cap is applied in - :func:`~dataretrieval.ogc.shaping._finalize_ogc`. - - Returns - ------- - pd.DataFrame - The completed chunk frame built from the capped feature sequence. - httpx.Response - Aggregated response — initial-request URL (for query identity), - final page's headers (so downstream sees current rate-limit - state), and cumulative ``elapsed`` summed across pages. - - Raises - ------ - DataRetrievalError - See :func:`_paginate`. - httpx.HTTPError - See :func:`_paginate`. + The paginator combines and caps feature lists before the single frame + conversion. ``geopd`` is a test override; production derives it from + geopandas availability and ``include_geometry``. ``row_cap`` only bounds + this chunk's page walk; finalization caps the combined result. """ method = req.method # ``httpx.Request.method`` is already upper-cased. headers = req.headers @@ -246,7 +184,7 @@ async def follow_up(cursor: str, sess: httpx.AsyncClient) -> httpx.Response: return ( _feature_frame( features, - geopd=geopd, + geopd=GEOPANDAS and include_geometry if geopd is None else geopd, include_geometry=include_geometry, ), response, @@ -349,11 +287,8 @@ def get_ogc_data( convert_type = args.pop("convert_type", False) args = {k: v for k, v in args.items() if v is not None} - # Choose one semantic frame shape for the whole request. Every completed - # chunk and the all-empty finalizer receive these same values, so frame type - # never depends on whether a particular page happens to carry geometry. + # Choose one semantic geometry mode for every completed and empty chunk. include_geometry = spatial and not bool(args.get("skip_geometry", False)) - geopd = GEOPANDAS and include_geometry # Post-processing is injected into the chunker rather than applied here, # so it runs on *every* exit: the normal return AND a later @@ -369,7 +304,7 @@ def get_ogc_data( output_id=output_id, convert_type=convert_type, collection=collection, - geopd=geopd, + geopd=GEOPANDAS and include_geometry, include_geometry=include_geometry, max_rows=max_rows, extra_id_cols=extra_id_cols, @@ -399,7 +334,6 @@ def get_ogc_data( [req], functools.partial( _walk_pages, - geopd, include_geometry=include_geometry, row_cap=max_rows, ), @@ -421,7 +355,6 @@ def get_ogc_data( fetch = functools.partial( _fetch_once, build_request=build_request, - geopd=geopd, include_geometry=include_geometry, row_cap=max_rows, ) @@ -435,7 +368,6 @@ async def _fetch_once( args: dict[str, Any], *, build_request: Callable[..., httpx.Request], - geopd: bool, include_geometry: bool, row_cap: int | None = None, ) -> tuple[pd.DataFrame, httpx.Response]: @@ -455,7 +387,6 @@ async def _fetch_once( """ req = build_request(**args) return await _walk_pages( - geopd=geopd, req=req, include_geometry=include_geometry, row_cap=row_cap, diff --git a/dataretrieval/transport/pagination.py b/dataretrieval/transport/pagination.py index 0023cd89..9f51899a 100644 --- a/dataretrieval/transport/pagination.py +++ b/dataretrieval/transport/pagination.py @@ -7,9 +7,9 @@ from __future__ import annotations import logging -from collections.abc import AsyncIterator, Awaitable, Callable +from collections.abc import AsyncIterator, Awaitable, Callable, Sized from contextlib import asynccontextmanager -from typing import Any, TypeVar, overload +from typing import Any, TypeVar, cast import httpx import pandas as pd @@ -75,46 +75,19 @@ def _combine_frame_pages( return result if row_cap is None else result.head(row_cap) -_Page = TypeVar("_Page") -_Result = TypeVar("_Result") +_Page = TypeVar("_Page", bound=Sized) -@overload -async def paginate( - initial_req: httpx.Request, - *, - parse_response: Callable[[httpx.Response], tuple[pd.DataFrame, _Cursor | None]], - follow_up: Callable[[_Cursor, httpx.AsyncClient], Awaitable[httpx.Response]], - raise_for_status: Callable[[httpx.Response], None], - client: httpx.AsyncClient | None = None, - row_cap: int | None = None, - combine_pages: None = None, -) -> tuple[pd.DataFrame, httpx.Response]: ... - - -@overload async def paginate( initial_req: httpx.Request, *, parse_response: Callable[[httpx.Response], tuple[_Page, _Cursor | None]], follow_up: Callable[[_Cursor, httpx.AsyncClient], Awaitable[httpx.Response]], raise_for_status: Callable[[httpx.Response], None], - combine_pages: Callable[[list[_Page], int | None], _Result], client: httpx.AsyncClient | None = None, row_cap: int | None = None, -) -> tuple[_Result, httpx.Response]: ... - - -async def paginate( - initial_req: httpx.Request, - *, - parse_response: Callable[[httpx.Response], tuple[Any, Any | None]], - follow_up: Callable[[Any, httpx.AsyncClient], Awaitable[httpx.Response]], - raise_for_status: Callable[[httpx.Response], None], - client: httpx.AsyncClient | None = None, - row_cap: int | None = None, - combine_pages: Callable[[list[Any], int | None], Any] | None = None, -) -> tuple[Any, httpx.Response]: + combine_pages: Callable[[list[_Page], int | None], _Page] | None = None, +) -> tuple[_Page, httpx.Response]: """Fetch and combine pages until the injected parser returns no cursor. The service adapter supplies response parsing, cursor following, status @@ -126,7 +99,7 @@ async def paginate( logger.debug("Requesting: %s", initial_req.url) reporter = _progress.current() - def report_page(page: httpx.Response, payload: Any) -> None: + def report_page(page: httpx.Response, payload: _Page) -> None: note_progress() # a walk still delivering pages is not stalled if reporter is not None: reporter.set_rate_remaining( @@ -179,7 +152,11 @@ def report_page(page: httpx.Response, payload: Any) -> None: headers_from=response, elapsed=total_elapsed, ) - combine = _combine_frame_pages if combine_pages is None else combine_pages + combine = ( + cast("Callable[[list[_Page], int | None], _Page]", _combine_frame_pages) + if combine_pages is None + else combine_pages + ) return combine(pages, row_cap), final_response diff --git a/tests/waterdata_filters_test.py b/tests/waterdata_filters_test.py index 3be97d51..7cefd62e 100644 --- a/tests/waterdata_filters_test.py +++ b/tests/waterdata_filters_test.py @@ -164,7 +164,7 @@ def test_long_filter_fans_out_into_multiple_requests(): expr = _filter_chunking_clauses() sent_filters: list[str] = [] - async def fake_walk_pages(*, geopd, req, include_geometry=True, row_cap=None): + async def fake_walk_pages(*, req, geopd=None, include_geometry=True, row_cap=None): idx = len(sent_filters) sent_filters.append(_query_params(req).get("filter", [None])[0]) return pd.DataFrame({"id": [f"chunk-{idx}"], "value": [idx]}), _fake_response() diff --git a/tests/waterdata_progress_test.py b/tests/waterdata_progress_test.py index 2f789f43..5b43309c 100644 --- a/tests/waterdata_progress_test.py +++ b/tests/waterdata_progress_test.py @@ -19,13 +19,14 @@ from dataretrieval import progress as _progress from dataretrieval.ogc.chunking import ChunkedCall -from dataretrieval.ogc.engine import _paginate, _walk_pages +from dataretrieval.ogc.engine import _walk_pages from dataretrieval.ogc.planning import ChunkPlan from dataretrieval.progress import ( ProgressReporter, current, progress_context, ) +from dataretrieval.transport.pagination import paginate as _paginate def _run_walk_pages(*, geopd, req, client): @@ -505,6 +506,7 @@ async def run(): req, parse_response=parse_sync, follow_up=follow_up, + raise_for_status=lambda _response: None, client=client, ) return df