diff --git a/README.md b/README.md index d904e6d6..9a9a7966 100644 --- a/README.md +++ b/README.md @@ -105,6 +105,58 @@ df, metadata = waterdata.get_continuous( print(f"Retrieved {len(df)} continuous gage height measurements") ``` +#### Speeding up large downloads with `parallel_chunks` + +By default the getters split a multi-value request only as far as the server's +~8 KB URL limit forces — the fewest sub-requests. For a **large, paginated** +pull that is needlessly conservative: every sub-request pages through its own +results, so dividing the query into more, smaller sub-requests lets those pages +be fetched **in parallel**. `parallel_chunks(n)` opts a single call into that +finer split, fanning it out into `n` sub-requests. It pays off only when the +result is large enough to span many pages *and* the query has a multi-value +argument to divide (such as a list of monitoring locations); on a small query — +or one with nothing to split — it just adds requests, so it is a deliberate, +scoped `with` block, never the default. + +```python +from dataretrieval import waterdata + +# All stream gages in Ohio, then 20 years of their daily discharge — large +# enough to span many pages, so it profits from a finer split. +sites, _ = waterdata.get_monitoring_locations(state="Ohio", site_type_code="ST") + +with waterdata.parallel_chunks(32): # fan out into 32 sub-requests + df, md = waterdata.get_daily( + monitoring_location_id=sites["monitoring_location_id"].tolist(), + parameter_code="00060", # discharge + time="2004-01-01/2023-12-31", + ) +``` + +`n` is the number of sub-requests to fan the call out into. It is capped by how +many values there are to split, and each sub-request costs a request against +your hourly [rate limit](https://api.waterdata.usgs.gov/signup/); since how many +run *at once* is capped separately by `API_USGS_CONCURRENT` (default 32), the +useful range is roughly `2` up to that value. + +Benchmark — a fixed 271-site subset of Ohio stream gages +(`get_daily`, `parameter_code="00060"`), with a small fixed page size +(`limit=250`) so every run fetches roughly the same number of pages (isolating +the effect of parallelism). Each `n` was run against its own cold 1-year time +window so no run is served from the server's data-window cache: + +| `n` | parallelism | pages | wall-clock | speedup | +| ---- | ----------- | ----- | ----------------------- | ------- | +| off | 1 | ~30 | 9.5 s / 9.1 s (2 runs) | 1× | +| `8` | 8 | ~32 | 2.2 s / 1.9 s | ~4.5× | +| `32` | 32 | 54 | 1.2 s | ~8× | + +The gain comes from overlapping each sub-request's per-page latency and +server-side work, so the exact multiplier scales with how many pages the pull +spans — a larger pull (more pages) has more parallelism to exploit. The extra +sub-requests each cost quota, so reserve a large `n` for pulls you know are +large. + Visit the [API Reference](https://doi-usgs.github.io/dataretrieval-python/reference/waterdata.html) for more information and examples on available services and input parameters. diff --git a/dataretrieval/__init__.py b/dataretrieval/__init__.py index c9df1c45..469fe0f5 100644 --- a/dataretrieval/__init__.py +++ b/dataretrieval/__init__.py @@ -44,6 +44,11 @@ URLTooLong, ) +# Parallel-chunks control (a context manager). Defined with the chunker in +# ``dataretrieval.ogc.chunking``; surfaced here for a stable public path +# ``from dataretrieval import parallel_chunks``. +from dataretrieval.ogc.chunking import parallel_chunks + # Resumable chunk-interruption exceptions. They are defined in # ``dataretrieval.ogc.interruptions`` rather than ``dataretrieval.exceptions`` # because they carry pandas/httpx state and a resumable ``ChunkedCall`` handle, @@ -93,5 +98,7 @@ "ChunkInterrupted", "QuotaExhausted", "ServiceInterrupted", + # parallel-chunks control (defined in ogc.chunking) + "parallel_chunks", "__version__", ] diff --git a/dataretrieval/ogc/chunking.py b/dataretrieval/ogc/chunking.py index 00d2766c..d0dae01c 100644 --- a/dataretrieval/ogc/chunking.py +++ b/dataretrieval/ogc/chunking.py @@ -9,6 +9,12 @@ cartesian product of chunks. Requests that already fit get a trivial single-step plan — ``ChunkedCall`` has one code path either way. +Parallel chunks: the planner is conservative by default — it splits only as far as +the byte limit forces. A caller who knows their result is large can opt into a +finer split via the ``parallel_chunks(n)`` context manager, which fans the query +out into ``n`` parallel sub-requests; ``n`` drives :meth:`ChunkPlan._refine`. See +``parallel_chunks`` for the why and the when. + This module owns the *execution* half — the event loop and bounded concurrency that drive a plan to completion (``ChunkedCall``) plus the public ``multi_value_chunked`` decorator. The neighboring concerns live in @@ -69,6 +75,7 @@ import functools import os from collections.abc import Callable, Iterator +from contextlib import contextmanager from contextvars import copy_context from typing import Any, cast @@ -76,7 +83,7 @@ import pandas as pd from anyio.from_thread import start_blocking_portal -from dataretrieval.utils import HTTPX_DEFAULTS, Ambient +from dataretrieval.utils import HTTPX_DEFAULTS, Ambient, _require_positive_int from . import progress as _progress from .interruptions import ( @@ -172,6 +179,118 @@ def get_active_client() -> httpx.AsyncClient | None: return _chunked_client.get() +# Parallel-chunks dial: opt-in to fan a query out *more finely* than the byte +# limit alone requires. Scoped to a ``with parallel_chunks(...):`` block (a +# ContextVar), deliberately NOT an env var (see :func:`parallel_chunks` for +# why). The ambient holds ``n`` — the requested cap on the plan's total +# sub-request count; ``1`` (the default, outside any block) means "off — chunk +# only as much as the byte limit needs, no extra fan-out". +_parallel_chunks: Ambient[int] = Ambient("ogc_parallel_chunks", 1) + + +@contextmanager +def parallel_chunks(n: int) -> Iterator[None]: + """ + Fan the OGC getters' multi-value requests out into ``n`` parallel sub-requests. + + By default the Water Data / NGWMN getters chunk a request only as much as + the server's ~8 KB URL-byte limit forces — the fewest sub-requests that + fit. That is the safe default, but it can be *needlessly* conservative: + because every sub-request paginates, splitting a large result further costs + little or no extra quota *as long as each sub-request still spans many + pages* — rows-per-chunk far exceeding the page size (ten states pulled as + one request then page nearly as many times as ten per-state requests + would). When a split leaves each sub-request only a page or two, its partial + final page is extra, so finer chunks do add some requests. This context + manager lets a caller who *knows* their pull is large ask for that finer + split — trading roughly the same pages for more, smaller sub-requests, which + gives smoother progress, more even concurrency, and a smaller unit of + retry/resume. + + Because the library can't tell in advance whether a query is large (ten + states over a short window might fit in a single page, where extra chunks + would only burn quota), this is a *deliberate* per-call knob rather than an + automatic behavior or a process-wide environment variable — scoping it to a + ``with`` block keeps an aggressive setting from leaking into unrelated calls + and accidentally spending quota. Outside any block the getters use the + conservative default. Only the OGC getters (Water Data, NGWMN) read this; + wrapping a legacy NWIS call in the block is a harmless no-op. + + Parameters + ---------- + n : int + The number of sub-requests to fan the whole call out into — a positive + integer such as ``2``, ``8``, or ``32``. It caps the plan's *total* + sub-request count (the cartesian product across every multi-value + argument combined, not per argument), so several multi-value arguments + cannot multiply past it. The cap is a ceiling, never exceeded: the + actual count is bounded below by what the ~8 KB URL limit already + forces and above by ``n``, so an ``n`` larger than the input allows + simply yields one sub-request per value, and with several multi-value + arguments the total may land somewhat below ``n`` because splits are + whole (the plan can't always divide evenly onto ``n``); ``n=1`` asks + for no extra fan-out. + + Each sub-request fetches at least one page, so it costs at least one + request against your hourly rate limit — a larger ``n`` spends more + quota. And because how many sub-requests run *at once* is capped + separately by ``API_USGS_CONCURRENT`` (default 32), an ``n`` beyond that + adds quota without adding parallelism; the useful range is roughly ``2`` + up to ``API_USGS_CONCURRENT``. + + Yields + ------ + None + + Raises + ------ + ValueError + If ``n`` is not a positive integer — raised on ``with`` entry, before + any request is issued, so a bad value fails loudly rather than silently + doing nothing. + + Notes + ----- + Fanning out carries the same consequences as the byte-limit chunking the + getters already do for oversized requests; opting in just brings them to a + request that would otherwise be a single call: + + - ``max_rows``: each sub-request paginates up to ``max_rows`` rows + independently, then the combined result is sorted and truncated to + ``max_rows``. So a call with ``max_rows`` set returns a *different* + (though still valid and deterministically sorted) row set inside a + ``parallel_chunks`` block than without one — the cap is drawn from the + union of the sub-requests, not a single stream. Don't pair a tight + ``max_rows`` preview with ``parallel_chunks`` if you need exactly the + rows the un-fanned call would return. + - Resumability: a single request either fully succeeds or fully fails, + but a fanned-out call can fail partway (e.g. a mid-call rate-limit) and + raise a resumable :class:`~dataretrieval.exceptions.ChunkInterrupted` + (or ``QuotaExhausted``) carrying the completed sub-requests, which you + finish with ``exc.call.resume()``. + - Cross-sub-request de-duplication keys on the feature ``id``; features + with no ``id`` can't be deduped, so overlapping filter clauses split + across chunks may yield duplicate rows. + + Examples + -------- + >>> from dataretrieval import waterdata + >>> with waterdata.parallel_chunks(32): + ... df, md = waterdata.get_daily( + ... monitoring_location_id=many_sites, parameter_code="00060" + ... ) # doctest: +SKIP + + See Also + -------- + ChunkPlan._refine : the planning-side effect of ``n``. + """ + # Fail loudly on a bad ``n`` at ``with`` entry, before any request. Shared + # rules with ``max_rows`` via the helper (accepts numpy ints, rejects bool). + _require_positive_int(n, "parallel_chunks(n)", examples="2, 8, 32") + with _parallel_chunks(n): + yield + + class ChunkedCall: """ Stateful handle for a chunked call. @@ -591,8 +710,9 @@ def multi_value_chunked( ``async def fetch(args) -> (df, response)``, and drives it to completion via :meth:`ChunkedCall.resume`. The plan splits multi-value list params and the cql-text filter so each sub-request URL fits the - byte limit; an already-fitting request is a one-step plan. See the - module docstring for the concurrency model. + byte limit; an already-fitting request is a one-step plan, unless an + active :func:`parallel_chunks` block asks the plan to fan out more + finely. See the module docstring for the concurrency model. Parameters ---------- @@ -636,7 +756,14 @@ def wrapper( finalize: _Finalize = _passthrough_result, ) -> tuple[pd.DataFrame, Any]: limit = _OGC_URL_BYTE_LIMIT if url_limit is None else url_limit - plan = ChunkPlan(args, build_request, limit) + # Read the parallel_chunks dial ``n`` from the ambient set by + # ``parallel_chunks`` (1 = off outside any such block; otherwise the + # requested total sub-request cap). It only affects *planning*, done + # here up front, so a later resume — which re-issues the + # already-planned sub-requests — needs no snapshot. + plan = ChunkPlan( + args, build_request, limit, max_chunks=_parallel_chunks.get() + ) retry_policy = RetryPolicy.from_env() # The concurrency cap is resolved inside ``resume()`` from # ``API_USGS_CONCURRENT``; ``1`` is a sequential gather, diff --git a/dataretrieval/ogc/engine.py b/dataretrieval/ogc/engine.py index a9a7b9e1..668fa7b7 100644 --- a/dataretrieval/ogc/engine.py +++ b/dataretrieval/ogc/engine.py @@ -27,7 +27,6 @@ import functools import json import logging -import numbers import re from collections.abc import ( AsyncIterator, @@ -59,6 +58,7 @@ _default_headers, _get, _network_error, + _require_positive_int, ) # Set up logger for this module @@ -876,18 +876,12 @@ def get_ogc_data( - Handles optional arguments such as `convert_type`. - Applies column cleanup and reordering based on service and properties. """ - # Enforce a genuine positive integer: a float (even ``10.0``) or ``bool`` - # would pass a bare ``< 1`` check and then crash deep in + # Enforce a genuine positive integer up front: a float (even ``10.0``) or + # ``bool`` would pass a bare ``< 1`` check and then crash deep in # ``pd.DataFrame.head`` with an opaque ``TypeError`` after HTTP I/O has - # already fired. ``numbers.Integral`` (not ``int``) so numpy integers — - # e.g. ``max_rows`` derived from a numpy/pandas computation — are accepted; - # ``bool`` is an ``Integral`` subtype, so exclude it explicitly. - if max_rows is not None and ( - not isinstance(max_rows, numbers.Integral) - or isinstance(max_rows, bool) - or max_rows < 1 - ): - raise ValueError(f"max_rows must be a positive integer (got {max_rows!r}).") + # already fired. Shared with ``parallel_chunks(n)`` via the helper. + if max_rows is not None: + _require_positive_int(max_rows, "max_rows") if dialect is None: dialect = _DEFAULT_DIALECT diff --git a/dataretrieval/ogc/planning.py b/dataretrieval/ogc/planning.py index 68c337ae..e80d0643 100644 --- a/dataretrieval/ogc/planning.py +++ b/dataretrieval/ogc/planning.py @@ -288,6 +288,21 @@ def _extract_axes(args: dict[str, Any]) -> list[_Axis]: return axes +def _split_at(chunks: list[list[str]], idx: int) -> None: + """Replace ``chunks[idx]`` in place with its two contiguous halves. + + The single primitive both planning passes use to fan an axis out. It + preserves the partition invariants every consumer relies on: *coverage* + (each atom survives, exactly once) and *contiguous, deterministic order* + (resume and :meth:`ChunkPlan.iter_sub_args` depend on it). Kept in one + place so those invariants can't drift between :meth:`ChunkPlan._plan` + (byte-driven) and :meth:`ChunkPlan._refine` (fan-out-driven). + """ + chunk = chunks[idx] + mid = len(chunk) // 2 + chunks[idx : idx + 1] = [chunk[:mid], chunk[mid:]] + + class ChunkPlan: """ Strategy for issuing one user-level request as a sequence of @@ -312,7 +327,22 @@ class ChunkPlan: Factory that turns a kwargs dict into a sized httpx request, e.g. ``_construct_api_requests``. url_limit : int - Byte budget for the request (URL + body). + Byte budget for the request (URL + body) — a hard ceiling every + sub-request must fit. + max_chunks : int, optional + Hard cap on the plan's total sub-request count (default ``1`` = off). + ``1`` chunks only as much as ``url_limit`` requires — the most + conservative plan, fewest sub-requests — so a fitting request is a + passthrough. A cap of ``2`` or more fans the plan out to up to + ``max_chunks`` sub-requests overall (the cartesian product across axes, + never fewer than the byte budget already forces) — capped as a whole, + not per axis, so several multi-value axes can't multiply past the cap. + The plan never exceeds the cap and may land below it when no whole + split lands on it exactly. ``max_chunks`` is a sub-request count, so a + value below ``1`` (``0`` or negative) is a caller error and raises + ``ValueError``. Set from the + :func:`~dataretrieval.ogc.chunking.parallel_chunks` ``n``; see + :meth:`_refine`. Attributes ---------- @@ -337,6 +367,8 @@ class ChunkPlan: Unchunkable If the request needs chunking but even the singleton plan doesn't fit ``url_limit``. + ValueError + If ``max_chunks`` is less than 1 (0 or negative). """ def __init__( @@ -344,7 +376,19 @@ def __init__( args: dict[str, Any], build_request: Callable[..., httpx.Request], url_limit: int, + max_chunks: int = 1, ) -> None: + if max_chunks < 1: + # ``max_chunks`` is a sub-request *count*: the minimum is ``1`` + # (the ambient default outside any ``parallel_chunks`` block), + # which means "off — no extra fan-out". ``0`` or negative is a + # meaningless count and can only be a caller bug, so fail loudly + # rather than silently no-op. The public ``parallel_chunks(n)`` + # already rejects ``n < 1``; this guards direct construction. + raise ValueError( + f"max_chunks must be >= 1 (1 disables fan-out); got {max_chunks!r}." + ) + self.args = args self.axes: list[_Axis] = [] self.chunks: dict[str, list[list[str]]] = {} @@ -352,10 +396,10 @@ def __init__( axes = _extract_axes(args) if not axes: - # No chunkable axis: nothing to split. If the single request fits, - # run it verbatim (the common passthrough). ``_safe_request_bytes`` - # treats an un-constructable URL (httpx.InvalidURL, > 64 KB) as over - # budget. + # No chunkable axis: nothing to split, and ``parallel_chunks`` has + # nothing to act on either. If the single request fits, run it + # verbatim (the common passthrough). ``_safe_request_bytes`` treats + # an un-constructable URL (httpx.InvalidURL, > 64 KB) as over budget. if _safe_request_bytes(build_request, args, url_limit) <= url_limit: return # Over budget. A filter the chunker doesn't manage — cql-json — is @@ -388,14 +432,30 @@ def __init__( except httpx.InvalidURL: initial_request = None + fits = False if initial_request is not None: self.canonical_url = str(initial_request.url) - if _request_bytes(initial_request) <= url_limit: - return + fits = _request_bytes(initial_request) <= url_limit + + # A request that already fits and hasn't opted into finer chunking is + # the common passthrough: leave ``axes``/``chunks`` empty so + # ``total == 1`` and ``iter_sub_args`` yields the original args + # verbatim. ``max_chunks == 1`` (off / no extra fan-out) means + # "don't split", so it takes this path; only ``max_chunks >= 2`` asks + # for extra fan-out and sets the axes up to be refined below. + if fits and max_chunks <= 1: + return self.axes = axes self.chunks = {axis.arg_key: [list(axis.atoms)] for axis in axes} - self._plan(build_request, url_limit) + if not fits: + # Hard pass: greedy-halve until every worst-case sub-request fits + # the byte budget (may raise ``Unchunkable``). + self._plan(build_request, url_limit) + # Soft pass: optionally split further than the byte budget requires. + # Purely additive — never re-raises, and the byte budget stays + # satisfied; a no-op at ``max_chunks == 1``. + self._refine(max_chunks) if self.canonical_url is None: # Original URL was un-constructable (httpx.InvalidURL); fall @@ -447,10 +507,69 @@ def _plan( f"sub-request). Reduce input sizes, shorten or simplify " f"the filter, or split the call manually." ) - axis_chunks = self.chunks[biggest_axis.arg_key] - chunk = axis_chunks[biggest_idx] - mid = len(chunk) // 2 - axis_chunks[biggest_idx : biggest_idx + 1] = [chunk[:mid], chunk[mid:]] + _split_at(self.chunks[biggest_axis.arg_key], biggest_idx) + + def _refine(self, max_chunks: int) -> None: + """ + Fan the plan out more finely than the byte budget alone requires — + the ``parallel_chunks`` dial (see + :func:`~dataretrieval.ogc.chunking.parallel_chunks` for why a caller + would want this, and :class:`ChunkPlan`'s ``max_chunks`` parameter for + the cap's contract: total-not-per-axis, a hard ceiling that may land + below the cap). + + Implementation. Each split multiplies the plan by ``(k+1)/k`` for the + chosen axis (adding ``total // k`` sub-requests, not one), so a split + is taken only when it keeps :attr:`total` within the cap; when no + in-budget split remains the plan stops *below* the cap rather than + overshooting (two even axes can reach 4 but not 5, so a cap of 5 yields + 4). Each split picks the single largest splittable chunk among the + in-budget axes (ties broken by axis-extraction order, then lowest + index), so growth is distributed round-robin rather than one axis + saturating before another is touched. Purely additive — only ever + *splits* existing chunks, so the byte pass's work and the ``url_limit`` + invariant are both preserved, and it never raises. A no-op at + ``max_chunks == 1``. + + Parameters + ---------- + max_chunks : int + The ``parallel_chunks(n)`` value; see :class:`ChunkPlan`'s + ``max_chunks`` parameter for the full contract. + """ + if max_chunks <= 1: + return + while True: + total = self.total + if total >= max_chunks: + return + # Largest splittable chunk among the axes whose split still fits the + # cap. Splitting any chunk of an axis with ``k`` chunks turns that + # ``k`` into ``k+1``, so it adds ``total // k`` sub-requests (the + # product of the other axes) regardless of which chunk — hence the + # budget test is per axis, not per chunk. Skipping an over-budget + # axis makes ``max_chunks`` a true ceiling. The ranking key is atom + # count (``len``), not URL bytes like ``_plan`` — this pass balances + # work across sub-requests rather than fitting a byte budget. A + # chunk of size 1 can't be split further. Stable input order breaks + # ties by axis order, then lowest index within an axis. + candidate: tuple[_Axis, int] | None = None + candidate_size = -1 + for axis in self.axes: + axis_chunks = self.chunks[axis.arg_key] + if total + total // len(axis_chunks) > max_chunks: + continue # any split of this axis would overshoot the cap + for idx, chunk in enumerate(axis_chunks): + if len(chunk) <= 1: + continue + if len(chunk) > candidate_size: + candidate, candidate_size = (axis, idx), len(chunk) + if candidate is None: + # Every axis is saturated at one atom per chunk or would + # overshoot the cap; stop below it rather than exceed it. + return + axis, idx = candidate + _split_at(self.chunks[axis.arg_key], idx) def _worst_case_args(self) -> dict[str, Any]: """ diff --git a/dataretrieval/utils.py b/dataretrieval/utils.py index 8d28ee86..acdf7822 100644 --- a/dataretrieval/utils.py +++ b/dataretrieval/utils.py @@ -4,6 +4,7 @@ from __future__ import annotations +import numbers import os import warnings from collections.abc import Callable, Iterable, Iterator @@ -66,6 +67,36 @@ def __call__(self, value: _T) -> Iterator[None]: self._var.reset(token) +def _require_positive_int( + value: int, name: str, *, examples: str | None = None +) -> None: + """Validate that ``value`` is a positive integer, else raise ``ValueError``. + + Accepts any :class:`numbers.Integral` (so a numpy/pandas integer passes, + not only ``int``) but rejects ``bool`` — an ``Integral`` subtype that is + nonsensical as a count. A non-integer (float, str, ``None``) or a value + ``< 1`` raises before any I/O, rather than crashing later (e.g. deep in + ``pd.DataFrame.head``). Shared by the user-facing count knobs ``max_rows`` + and ``parallel_chunks(n)`` so their boundary validation can't drift. + (``ChunkPlan.max_chunks`` is an internal, already-``int`` precondition with + its own domain-specific message, so it keeps a lighter ``< 1`` guard rather + than routing through here.) + + Parameters + ---------- + value : int + The value to check. Typed ``int`` for callers, but validated at + runtime because the real value may be anything the user passed. + name : str + Parameter name, used as the subject of the error message. + examples : str, optional + Illustrative values appended to the message (e.g. ``"2, 8, 32"``). + """ + if not isinstance(value, numbers.Integral) or isinstance(value, bool) or value < 1: + eg = f", e.g. {examples}" if examples else "" + raise ValueError(f"{name} must be a positive integer{eg} (got {value!r}).") + + def _default_headers() -> dict[str, str]: """Build the default HTTP headers for a USGS web-API request. diff --git a/dataretrieval/waterdata/__init__.py b/dataretrieval/waterdata/__init__.py index 99b6e178..eb231469 100644 --- a/dataretrieval/waterdata/__init__.py +++ b/dataretrieval/waterdata/__init__.py @@ -9,6 +9,7 @@ from __future__ import annotations +from dataretrieval.ogc.chunking import parallel_chunks from dataretrieval.ogc.filters import FILTER_LANG # Public API exports @@ -50,6 +51,7 @@ "PROFILE_LOOKUP", "SERVICES", "WATERDATA_SERVICES", + "parallel_chunks", "get_channel", "get_codes", "get_combined_metadata", diff --git a/docs/source/userguide/errors.rst b/docs/source/userguide/errors.rst index e2dc3ef1..28da515f 100644 --- a/docs/source/userguide/errors.rst +++ b/docs/source/userguide/errors.rst @@ -96,6 +96,45 @@ condition clears -- only the unfinished sub-requests are re-issued. except ChunkInterrupted as again: exc = again +Chunk a large request more finely +================================= + +By default the getters split an over-large request only as much as the +server's ~8 KB URL limit forces -- the fewest sub-requests. Because each +sub-request paginates, splitting a large result further costs little or no +extra quota *as long as each sub-request still spans many pages* (ten states +pulled as one request then page nearly as many times as ten per-state requests +would; a split that leaves each sub-request only a page or two adds its partial +final page). So if you *know* your pull is large you can ask for a finer split +with ``parallel_chunks(n)`` -- trading roughly the same pages for more, smaller +sub-requests, which gives smoother progress, more even concurrency, and a +smaller unit of retry/resume. It is a scoped ``with`` +block, so an aggressive setting can't leak into unrelated calls and +accidentally spend quota: + +.. code-block:: python + + from dataretrieval import waterdata + + with waterdata.parallel_chunks(32): + df, md = waterdata.get_daily( + monitoring_location_id=many_sites, parameter_code="00060" + ) + +``n`` is a positive integer (e.g. ``2``, ``8``, ``32``) -- the number of +sub-requests to fan the call out into; a non-integer or non-positive value +raises ``ValueError`` at the ``with``. It caps the *total* sub-request count +across every multi-value argument combined (not per argument), bounded below by +what the byte limit already forces and above by how many values there are to +split, so several multi-value arguments can't multiply past it and ``n=1`` asks +for no extra fan-out. Each sub-request costs a request against your hourly rate +limit, and because how many run *at once* is capped separately by +``API_USGS_CONCURRENT`` (default 32) an ``n`` beyond that adds quota without +adding parallelism -- the useful range is roughly ``2`` up to +``API_USGS_CONCURRENT``. There is no "off" level: simply don't enter the block +unless you already expect a large, multi-page result -- on a query that would +have fit in a single page, extra chunks only burn quota. + The full taxonomy ================= diff --git a/tests/utils_test.py b/tests/utils_test.py index b3201419..30950294 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -208,6 +208,20 @@ def test_chunk_interruptions_exported_at_top_level(self): dataretrieval.ChunkInterrupted, dataretrieval.DataRetrievalError ) + def test_parallel_chunks_exported_at_top_level_and_waterdata(self): + """The ``parallel_chunks`` context manager is reachable both from the top + level (``from dataretrieval import parallel_chunks``) and from the + user-facing ``dataretrieval.waterdata`` namespace, and both resolve to + the single object defined in ``dataretrieval.ogc.chunking``.""" + import dataretrieval + from dataretrieval import waterdata + from dataretrieval.ogc import chunking + + assert dataretrieval.parallel_chunks is chunking.parallel_chunks + assert waterdata.parallel_chunks is chunking.parallel_chunks + assert "parallel_chunks" in dataretrieval.__all__ + assert "parallel_chunks" in waterdata.__all__ + class Test_BaseMetadata: """Tests of BaseMetadata""" diff --git a/tests/waterdata_chunking_test.py b/tests/waterdata_chunking_test.py index 5a946946..600f05d8 100644 --- a/tests/waterdata_chunking_test.py +++ b/tests/waterdata_chunking_test.py @@ -42,8 +42,10 @@ from dataretrieval.ogc.chunking import ( ChunkedCall, _chunked_client, + _parallel_chunks, get_active_client, multi_value_chunked, + parallel_chunks, ) from dataretrieval.ogc.interruptions import ( ChunkInterrupted, @@ -2105,3 +2107,296 @@ async def fetch(args): assert "finalized" in df.columns assert md[0] == "METADATA" assert calls["finalize"] >= 1 + + +# --------------------------------------------------------------------------- +# Parallel chunks: the opt-in dial ``parallel_chunks(n)`` to fan a query out +# MORE finely than the byte limit alone requires (``ChunkPlan._refine`` + the +# ``parallel_chunks`` context manager). ``_fake_build``'s base is 200 bytes, so +# a handful of short atoms sits far under ``url_limit=8000`` — the byte pass +# passes it through untouched, and any splitting below is the ``n`` cap alone. +# ``ChunkPlan`` takes the integer cap (``max_chunks``) directly; +# ``parallel_chunks(n)`` publishes ``n`` onto it. The cap bounds the plan's +# *total* sub-request count (the cartesian product across axes), not each axis +# independently — see ``test_cap_caps_the_total_across_axes``. +# --------------------------------------------------------------------------- + + +def test_default_preserves_passthrough(): + """The default ``max_chunks`` (1 = off) must not perturb the existing + plan: a multi-value request that fits the byte limit is still the trivial + passthrough (no axes, ``total == 1``), byte-for-byte the pre-feature + behavior.""" + args = {"monitoring_location_id": ["A", "B", "C", "D"]} + plan = ChunkPlan(args, _fake_build, url_limit=8000) # default max_chunks=1 + assert plan.axes == [] + assert plan.total == 1 + assert list(plan.iter_sub_args()) == [args] + + +def test_unit_cap_preserves_passthrough(): + """``max_chunks=1`` means "no extra fan-out", so a fitting multi-value + request stays the trivial passthrough (no axes, ``total == 1``, + ``iter_sub_args`` yields the original args verbatim) — identical to the + default (off), not a materialized one-chunk-per-axis plan.""" + args = {"monitoring_location_id": ["A", "B", "C", "D"]} + plan = ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=1) + assert plan.axes == [] + assert plan.total == 1 + assert list(plan.iter_sub_args()) == [args] + + +@pytest.mark.parametrize("bad", [0, -1]) +def test_invalid_cap_raises(bad): + """``max_chunks`` is a sub-request count, so a value below 1 (``0`` or + negative) is a caller bug, not a silent no-op: it raises ``ValueError`` at + construction. (The public ``parallel_chunks(n)`` already rejects ``n < 1``; + this pins the same guard on direct construction.)""" + args = {"monitoring_location_id": ["A", "B", "C", "D"]} + with pytest.raises(ValueError, match="max_chunks must be >= 1"): + ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=bad) + + +@pytest.mark.parametrize( + ("max_chunks", "expected_pieces"), + [(1, 1), (2, 2), (8, 8), (16, 10), (32, 10)], +) +def test_cap_ramps_then_saturates(max_chunks, expected_pieces): + """A single 10-atom axis that fits the byte limit splits into + ``min(10, cap)`` pieces: 1 (off), 2, 8, then saturating at 10 (one atom per + chunk) once the cap overshoots the atom count. Monotonic and bounded, and + whenever it splits the partition is a cover — every atom exactly once. (The + cap-1 passthrough has no axis to cover; see the passthrough test.)""" + atoms = [f"S{i:02d}" for i in range(10)] + plan = ChunkPlan( + {"monitoring_location_id": atoms}, + _fake_build, + url_limit=8000, + max_chunks=max_chunks, + ) + assert plan.total == expected_pieces + if expected_pieces > 1: + flattened = [ + a for chunk in plan.chunks["monitoring_location_id"] for a in chunk + ] + assert sorted(flattened) == sorted(atoms) + + +def test_cap_bounds_fan_out_for_a_long_axis(): + """The cap holds fan-out to ``n``: at ``n=32`` a 100-atom axis fans into + ``n`` pieces — NOT 100 singletons — so ``parallel_chunks(32)`` on a huge + list can't detonate into hundreds of sub-requests. Every atom is still + covered exactly once.""" + high = 32 + atoms = [f"X{i:03d}" for i in range(100)] + plan = ChunkPlan( + {"monitoring_location_id": atoms}, + _fake_build, + url_limit=8000, + max_chunks=high, + ) + assert plan.total == high + flattened = [a for chunk in plan.chunks["monitoring_location_id"] for a in chunk] + assert sorted(flattened) == sorted(atoms) + + +def test_cap_below_byte_split_does_not_reduce_fan_out(): + """The cap is purely additive — it can only split further, never coarsen. + A request the byte budget already fans into K>2 chunks is untouched by a + cap of 2 (below K), so the byte-driven plan is preserved.""" + # Heavy axis of four 30-char atoms; a limit tight enough that the byte pass + # must drive every atom into its own sub-request (4 pieces > the cap of 2). + args = {"monitoring_location_id": ["X" * 30, "Y" * 30, "Z" * 30, "W" * 30]} + baseline = ChunkPlan(args, _fake_build, url_limit=250, max_chunks=1) + assert baseline.total > 2 # byte pass alone already fanned out past 2 + refined = ChunkPlan(args, _fake_build, url_limit=250, max_chunks=2) + # cap 2 < baseline pieces → refine is a no-op here. + assert refined.total == baseline.total + + +def test_cap_never_exceeds_the_byte_budget(): + """Refining on top of an over-budget request keeps the hard invariant: + every sub-request still fits ``url_limit`` (splitting only ever shrinks + a chunk), and the fan-out is at least what the byte pass required.""" + args = {"monitoring_location_id": ["X" * 30, "Y" * 30, "Z" * 30, "W" * 30]} + limit = 310 + byte_only = ChunkPlan(args, _fake_build, url_limit=limit, max_chunks=1) + plan = ChunkPlan(args, _fake_build, url_limit=limit, max_chunks=32) + assert plan.total >= byte_only.total + for sub in plan.iter_sub_args(): + assert _safe_request_bytes(_fake_build, sub, limit) <= limit + + +def test_cap_refines_the_filter_axis(): + """The dial treats the cql-text ``filter`` axis like any other: an + under-budget filter of N top-level OR-clauses is split along that axis + into ``min(N, cap)`` pieces.""" + clauses = [f"p='{i}'" for i in range(8)] + args = {"filter": " OR ".join(clauses)} + plan = ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=4) + assert len(plan.chunks["filter"]) == 4 # min(8, 4) + assert plan.total == 4 + + +def test_cap_caps_the_total_across_axes(): + """With more than one multi-value axis the cap bounds the *total* + sub-request count (the cartesian product), not each axis independently — + the blast-radius guardrail the dial exists for. Two 6-atom axes at a cap + of 4 top out at 4 sub-requests total, not 4x4=16; growth is distributed + round-robin across axes rather than one axis alone climbing to the cap.""" + args = { + "monitoring_location_id": [f"L{i}" for i in range(6)], + "parameter_code": [f"{i:05d}" for i in range(6)], + } + plan = ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=4) + assert plan.total == 4 + # Every atom on every axis is still covered exactly once. + for key, atoms in ( + ("monitoring_location_id", args["monitoring_location_id"]), + ("parameter_code", args["parameter_code"]), + ): + flattened = [a for chunk in plan.chunks[key] for a in chunk] + assert sorted(flattened) == sorted(atoms) + + +def test_cap_bounds_fan_out_across_many_axes(): + """The guardrail holds regardless of axis count: three multi-value axes at + a cap of 30 fan out to *at most* 30 sub-requests total — never the + ``30 ** 3`` a per-axis cap would allow, and never *over* the cap either. + 30 is deliberately not evenly reachable by these axes: a single split + multiplies the plan by more than one, so the naive ``while total < cap`` + the first refine used stepped past 30 (to 32). The cap is a hard ceiling — + the property neither the single-axis-only cap nor that naive loop + guaranteed.""" + cap = 30 + # Three chunkable axes (two list axes + the filter OR-axis), each with 10 + # atoms — under the old per-axis cap this would have been cap**3. + args = { + "monitoring_location_id": [f"L{i}" for i in range(10)], + "parameter_code": [f"{i:05d}" for i in range(10)], + "filter": " OR ".join(f"p='{i}'" for i in range(10)), + } + plan = ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=cap) + assert 1 < plan.total <= cap # fanned out, but never past the ceiling + + +@pytest.mark.parametrize( + "atoms_per_axis, cap", + [ + (4, 5), # pre-fix loop overshot 5 -> 6 + (8, 10), # pre-fix loop overshot 10 -> 12 + (10, 7), # pre-fix loop overshot 7 -> 8 + ], +) +def test_cap_is_a_hard_ceiling_never_overshoots(atoms_per_axis, cap): + """The cap is a hard ceiling, not a soft target. With two multi-value axes + a single split multiplies the plan by ``(k+1)/k`` for the split axis — + adding the product of the *other* axes, not one — so a naive + ``while total < cap`` loop steps *past* the cap. These are exactly the + (atoms, cap) combos that loop overshot (5->6, 10->12, 7->8). The plan must + fan out and cover every atom once, but never exceed the cap, landing below + it when no whole split lands on it exactly (two even axes reach 4, not 5).""" + args = { + "monitoring_location_id": [f"L{i:03d}" for i in range(atoms_per_axis)], + "parameter_code": [f"{i:05d}" for i in range(atoms_per_axis)], + } + plan = ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=cap) + assert 1 < plan.total <= cap # fanned out, but never past the ceiling + # Every atom on every axis is still covered exactly once. + for key, atoms in args.items(): + flattened = [a for chunk in plan.chunks[key] for a in chunk] + assert sorted(flattened) == sorted(atoms) + + +def test_cap_does_not_mask_unchunkable(): + """A request with nothing to split that still busts the byte limit must + raise ``Unchunkable`` regardless of the cap — the soft pass has no axis to + act on and must not swallow the hard failure.""" + args = {"monitoring_location_id": "one-huge-scalar"} + with pytest.raises(Unchunkable): + ChunkPlan(args, _fake_build, url_limit=10, max_chunks=32) + + +def test_parallel_chunks_publishes_n_on_the_ambient(): + """The context manager publishes ``n`` on the ambient for the block and + restores the previous value on exit — including proper nesting.""" + assert _parallel_chunks.get() == 1 # default (off, = no extra fan-out) + with parallel_chunks(32): + assert _parallel_chunks.get() == 32 + with parallel_chunks(2): + assert _parallel_chunks.get() == 2 + assert _parallel_chunks.get() == 32 # outer restored + assert _parallel_chunks.get() == 1 # default (off) outside any block + + +@pytest.mark.parametrize( + "bad", + [ + 0, # not positive + -1, # negative + 1.5, # a float, not an int + "8", # a string, even a numeric one + "high", # the old level names are gone + None, # None not accepted + True, # bool is an int subclass but nonsensical here + ["8"], # a list + ], +) +def test_parallel_chunks_rejects_non_positive_int(bad): + """``n`` must be a positive integer; every other shape — zero, negative, a + float, a string (including a numeric one and the old level names), ``None``, + a ``bool``, a list — raises ``ValueError`` at ``with`` entry, before any + request, and leaves the ambient untouched.""" + with pytest.raises(ValueError, match="must be a positive integer"): + with parallel_chunks(bad): + pass + assert _parallel_chunks.get() == 1 # default (off) — unchanged by a rejected call + + +def test_parallel_chunks_drives_end_to_end_fan_out(): + """End-to-end: the same fitting request passes through as a single call by + default, but fans into ``n`` sub-requests inside a ``parallel_chunks(n)`` + block — and the combined result still recovers every atom exactly once.""" + sites = [f"S{i:02d}" for i in range(8)] + + calls: list[tuple[str, ...]] = [] + + @multi_value_chunked(build_request=_fake_build, url_limit=8000) + async def fetch(args): + chunk = tuple(args["monitoring_location_id"]) + calls.append(chunk) + return pd.DataFrame({"site": list(chunk)}), _ok_response() + + # Default: comfortably under the byte limit → one passthrough call. + df_plain, _ = fetch({"monitoring_location_id": sites}) + assert len(calls) == 1 + assert sorted(df_plain["site"]) == sorted(sites) + + calls.clear() + with parallel_chunks(8): + df_fine, _ = fetch({"monitoring_location_id": sites}) + # 8 atoms at n=8 → 8 singleton sub-requests. + assert len(calls) == 8 + assert all(len(chunk) == 1 for chunk in calls) + # Union across chunks recovers the original set, once each. + assert sorted(a for chunk in calls for a in chunk) == sorted(sites) + assert sorted(df_fine["site"]) == sorted(sites) + + +@pytest.mark.parametrize("n", [1, 2, 3, 8]) +def test_parallel_chunks_supports_arbitrary_n(n): + """An arbitrary ``n`` (not only 2/8/32) fans an under-limit request into + exactly ``n`` sub-requests, together covering every site once — including + ``n=1``, the explicit no-op that stays a single passthrough call.""" + sites = [f"S{i:02d}" for i in range(8)] + calls: list[int] = [] + + @multi_value_chunked(build_request=_fake_build, url_limit=8000) + async def fetch(args): + calls.append(len(args["monitoring_location_id"])) + return pd.DataFrame(), _ok_response() + + with parallel_chunks(n): + fetch({"monitoring_location_id": sites}) + assert len(calls) == n + assert sum(calls) == 8