feat(ogc): add opt-in parallel chunk fan-out (#341)

* feat(ogc): add opt-in parallel chunk fan-out

Add parallel_chunks(n) so large Water Data and NGWMN queries can split
into more independently paginated sub-requests while preserving the
conservative byte-limit-only default.

Cap optional fan-out across all chunk axes, validate positive integer
inputs, export the helper publicly, and document usage, tradeoffs, and
benchmark results. Add planner, context, and integration coverage.

* docs: simplify parallel_chunks Series example

Author: thodson-usgs

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"],
+        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.

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__",
 ]

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,14 +75,15 @@
 import functools
 import os
 from collections.abc import Callable, Iterator
+from contextlib import contextmanager
 from contextvars import copy_context
 from typing import Any, cast
 
 import httpx
 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,

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