diff --git a/benchmarks/ogc_raw_feature_pagination.py b/benchmarks/ogc_raw_feature_pagination.py new file mode 100644 index 00000000..7564c88e --- /dev/null +++ b/benchmarks/ogc_raw_feature_pagination.py @@ -0,0 +1,243 @@ +"""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 + +# 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 + +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/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..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 @@ -51,7 +48,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 +97,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", []): @@ -120,92 +117,53 @@ 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, *, geopd: bool -) -> 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. + resp: httpx.Response, +) -> 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, body=body), - _next_req_url(resp, body=body) or None, - ) + features = body.get("features") + if features is None: + features = [] + 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( + 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( - 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]: - """ - Iterate paginated OGC API responses and aggregate them into one DataFrame. + """Aggregate raw OGC 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`). - - Parameters - ---------- - geopd : bool - Whether geopandas is installed (drives geometry handling). - req : httpx.Request - The initial HTTP request to send. - client : httpx.AsyncClient, optional - Caller-borrowed client; ``None`` defers client management to - :func:`_paginate`. - 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`. - - Returns - ------- - pd.DataFrame - A DataFrame containing the aggregated results from all pages. - 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 @@ -214,12 +172,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, geopd=geopd), + parse_response=_ogc_parse_response, follow_up=follow_up, - client=client, + 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=GEOPANDAS and include_geometry if geopd is None else geopd, + include_geometry=include_geometry, + ), + response, ) @@ -229,6 +197,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 +235,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 +287,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 geometry mode for every completed and empty chunk. + include_geometry = spatial and not bool(args.get("skip_geometry", False)) + # 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 +304,8 @@ def get_ogc_data( output_id=output_id, convert_type=convert_type, collection=collection, + geopd=GEOPANDAS and include_geometry, + include_geometry=include_geometry, max_rows=max_rows, extra_id_cols=extra_id_cols, dialect=dialect, @@ -354,7 +332,11 @@ 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, + include_geometry=include_geometry, + row_cap=max_rows, + ), RetryPolicy.from_env(), finalize, canonical_url=str(req.url), @@ -371,7 +353,10 @@ 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, + 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 +368,7 @@ async def _fetch_once( args: dict[str, Any], *, build_request: Callable[..., httpx.Request], + 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 +386,8 @@ 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( + req=req, + include_geometry=include_geometry, + row_cap=row_cap, + ) diff --git a/dataretrieval/ogc/shaping.py b/dataretrieval/ogc/shaping.py index d97d58b5..2ad3a2af 100644 --- a/dataretrieval/ogc/shaping.py +++ b/dataretrieval/ogc/shaping.py @@ -43,16 +43,34 @@ ) -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 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. """ - return gpd.GeoDataFrame() 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: @@ -82,87 +100,39 @@ 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, *, - body: dict[str, Any] | None = None, + include_geometry: bool = True, ) -> 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 - Indicates whether geopandas is installed and should be used to - handle geometries. - 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. ``_empty_feature_frame`` - # preserves the GeoDataFrame type on the short-circuit (see its docstring). - 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] 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] - _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. 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. - 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 + df["id"] = [feature.get("id") for feature in features] + return df[["id"] + [col for col in df.columns if col != "id"]] def _deal_with_empty( @@ -171,44 +141,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. + + 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. """ - 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. - - 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()) - schema, _ = _check_ogc_requests( - endpoint=collection, req_type="schema", base_url=base_url - ) - properties = list(schema.get("properties", {}).keys()) - return pd.DataFrame(columns=properties) - return return_list + return _empty_feature_frame( + geopd, + properties, + include_geometry=include_geometry, + ) def _arrange_cols( @@ -361,6 +319,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, @@ -383,13 +343,20 @@ 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. """ 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/transport/pagination.py b/dataretrieval/transport/pagination.py index 523463c6..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 +from typing import Any, TypeVar, cast import httpx import pandas as pd @@ -67,33 +67,46 @@ 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", bound=Sized) + + async def paginate( initial_req: httpx.Request, *, - parse_response: Callable[[httpx.Response], tuple[pd.DataFrame, _Cursor | None]], + 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], client: httpx.AsyncClient | None = None, row_cap: int | None = None, -) -> tuple[pd.DataFrame, 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, 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: _Page) -> 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 +115,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 +134,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 +152,12 @@ 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 = ( + 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 def run_paginated( 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/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/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/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 9ca439c6..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, 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() @@ -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_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 diff --git a/tests/waterdata_test.py b/tests/waterdata_test.py index 7d8fe80f..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.""" @@ -762,6 +802,72 @@ 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) + + +@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 9081b701..0a506b8a 100644 --- a/tests/waterdata_utils_test.py +++ b/tests/waterdata_utils_test.py @@ -33,7 +33,8 @@ from dataretrieval.ogc.schema import _check_ogc_requests from dataretrieval.ogc.shaping import ( _arrange_cols, - _get_resp_data, + _deal_with_empty, + _feature_frame, _to_snake_case, ) from dataretrieval.ogc.shaping import _finalize_ogc as _ogc_finalize @@ -53,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__ @@ -357,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) @@ -374,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.""" @@ -605,8 +601,67 @@ class _Sentinel: assert isinstance(result, _Sentinel) -def test_get_resp_data_empty_preserves_geopd_type(): - """Same as the stats-side preservation: ``_get_resp_data``'s +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(): + """An empty geospatial result remains usable like a non-empty one.""" + import geopandas as gpd + + page = _feature_frame([], geopd=True) + template = pd.concat([page], ignore_index=True) + out = _deal_with_empty( + template, + ["monitoring_location_id"], + "daily", + base_url="x", + geopd=True, + include_geometry=True, + ) + + 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" + out["monitoring_location_id"].str.startswith("USGS") + out.to_crs("EPSG:3857") + 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" + + +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 @@ -618,33 +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) -def test_get_resp_data_attaches_wgs84_crs(): +@pytest.mark.skipif(not _shaping_module.GEOPANDAS, reason="requires geopandas") +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 = _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() + assert missing.geometry.name == "geometry" + assert missing.crs == "EPSG:4326" + + +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" @@ -688,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"}