Skip to content

Commit a2c5627

Browse files
committed
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.
1 parent 6baa4ab commit a2c5627

10 files changed

Lines changed: 708 additions & 28 deletions

File tree

README.md

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,58 @@ df, metadata = waterdata.get_continuous(
105105
print(f"Retrieved {len(df)} continuous gage height measurements")
106106
```
107107

108+
#### Speeding up large downloads with `parallel_chunks`
109+
110+
By default the getters split a multi-value request only as far as the server's
111+
~8 KB URL limit forces — the fewest sub-requests. For a **large, paginated**
112+
pull that is needlessly conservative: every sub-request pages through its own
113+
results, so dividing the query into more, smaller sub-requests lets those pages
114+
be fetched **in parallel**. `parallel_chunks(n)` opts a single call into that
115+
finer split, fanning it out into `n` sub-requests. It pays off only when the
116+
result is large enough to span many pages *and* the query has a multi-value
117+
argument to divide (such as a list of monitoring locations); on a small query —
118+
or one with nothing to split — it just adds requests, so it is a deliberate,
119+
scoped `with` block, never the default.
120+
121+
```python
122+
from dataretrieval import waterdata
123+
124+
# All stream gages in Ohio, then 20 years of their daily discharge — large
125+
# enough to span many pages, so it profits from a finer split.
126+
sites, _ = waterdata.get_monitoring_locations(state="Ohio", site_type_code="ST")
127+
128+
with waterdata.parallel_chunks(32): # fan out into 32 sub-requests
129+
df, md = waterdata.get_daily(
130+
monitoring_location_id=sites["monitoring_location_id"].tolist(),
131+
parameter_code="00060", # discharge
132+
time="2004-01-01/2023-12-31",
133+
)
134+
```
135+
136+
`n` is the number of sub-requests to fan the call out into. It is capped by how
137+
many values there are to split, and each sub-request costs a request against
138+
your hourly [rate limit](https://api.waterdata.usgs.gov/signup/); since how many
139+
run *at once* is capped separately by `API_USGS_CONCURRENT` (default 32), the
140+
useful range is roughly `2` up to that value.
141+
142+
Benchmark — a fixed 271-site subset of Ohio stream gages
143+
(`get_daily`, `parameter_code="00060"`), with a small fixed page size
144+
(`limit=250`) so every run fetches roughly the same number of pages (isolating
145+
the effect of parallelism). Each `n` was run against its own cold 1-year time
146+
window so no run is served from the server's data-window cache:
147+
148+
| `n` | parallelism | pages | wall-clock | speedup |
149+
| ---- | ----------- | ----- | ----------------------- | ------- |
150+
| off | 1 | ~30 | 9.5 s / 9.1 s (2 runs) ||
151+
| `8` | 8 | ~32 | 2.2 s / 1.9 s | ~4.5× |
152+
| `32` | 32 | 54 | 1.2 s | ~|
153+
154+
The gain comes from overlapping each sub-request's per-page latency and
155+
server-side work, so the exact multiplier scales with how many pages the pull
156+
spans — a larger pull (more pages) has more parallelism to exploit. The extra
157+
sub-requests each cost quota, so reserve a large `n` for pulls you know are
158+
large.
159+
108160
Visit the
109161
[API Reference](https://doi-usgs.github.io/dataretrieval-python/reference/waterdata.html)
110162
for more information and examples on available services and input parameters.

dataretrieval/__init__.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,11 @@
4444
URLTooLong,
4545
)
4646

47+
# Parallel-chunks control (a context manager). Defined with the chunker in
48+
# ``dataretrieval.ogc.chunking``; surfaced here for a stable public path
49+
# ``from dataretrieval import parallel_chunks``.
50+
from dataretrieval.ogc.chunking import parallel_chunks
51+
4752
# Resumable chunk-interruption exceptions. They are defined in
4853
# ``dataretrieval.ogc.interruptions`` rather than ``dataretrieval.exceptions``
4954
# because they carry pandas/httpx state and a resumable ``ChunkedCall`` handle,
@@ -93,5 +98,7 @@
9398
"ChunkInterrupted",
9499
"QuotaExhausted",
95100
"ServiceInterrupted",
101+
# parallel-chunks control (defined in ogc.chunking)
102+
"parallel_chunks",
96103
"__version__",
97104
]

dataretrieval/ogc/chunking.py

Lines changed: 131 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,12 @@
99
cartesian product of chunks. Requests that already fit get a trivial
1010
single-step plan — ``ChunkedCall`` has one code path either way.
1111
12+
Parallel chunks: the planner is conservative by default — it splits only as far as
13+
the byte limit forces. A caller who knows their result is large can opt into a
14+
finer split via the ``parallel_chunks(n)`` context manager, which fans the query
15+
out into ``n`` parallel sub-requests; ``n`` drives :meth:`ChunkPlan._refine`. See
16+
``parallel_chunks`` for the why and the when.
17+
1218
This module owns the *execution* half — the event loop and bounded
1319
concurrency that drive a plan to completion (``ChunkedCall``) plus the
1420
public ``multi_value_chunked`` decorator. The neighboring concerns live in
@@ -69,14 +75,15 @@
6975
import functools
7076
import os
7177
from collections.abc import Callable, Iterator
78+
from contextlib import contextmanager
7279
from contextvars import copy_context
7380
from typing import Any, cast
7481

7582
import httpx
7683
import pandas as pd
7784
from anyio.from_thread import start_blocking_portal
7885

79-
from dataretrieval.utils import HTTPX_DEFAULTS, Ambient
86+
from dataretrieval.utils import HTTPX_DEFAULTS, Ambient, _require_positive_int
8087

8188
from . import progress as _progress
8289
from .interruptions import (
@@ -172,6 +179,118 @@ def get_active_client() -> httpx.AsyncClient | None:
172179
return _chunked_client.get()
173180

174181

182+
# Parallel-chunks dial: opt-in to fan a query out *more finely* than the byte
183+
# limit alone requires. Scoped to a ``with parallel_chunks(...):`` block (a
184+
# ContextVar), deliberately NOT an env var (see :func:`parallel_chunks` for
185+
# why). The ambient holds ``n`` — the requested cap on the plan's total
186+
# sub-request count; ``1`` (the default, outside any block) means "off — chunk
187+
# only as much as the byte limit needs, no extra fan-out".
188+
_parallel_chunks: Ambient[int] = Ambient("ogc_parallel_chunks", 1)
189+
190+
191+
@contextmanager
192+
def parallel_chunks(n: int) -> Iterator[None]:
193+
"""
194+
Fan the OGC getters' multi-value requests out into ``n`` parallel sub-requests.
195+
196+
By default the Water Data / NGWMN getters chunk a request only as much as
197+
the server's ~8 KB URL-byte limit forces — the fewest sub-requests that
198+
fit. That is the safe default, but it can be *needlessly* conservative:
199+
because every sub-request paginates, splitting a large result further costs
200+
little or no extra quota *as long as each sub-request still spans many
201+
pages* — rows-per-chunk far exceeding the page size (ten states pulled as
202+
one request then page nearly as many times as ten per-state requests
203+
would). When a split leaves each sub-request only a page or two, its partial
204+
final page is extra, so finer chunks do add some requests. This context
205+
manager lets a caller who *knows* their pull is large ask for that finer
206+
split — trading roughly the same pages for more, smaller sub-requests, which
207+
gives smoother progress, more even concurrency, and a smaller unit of
208+
retry/resume.
209+
210+
Because the library can't tell in advance whether a query is large (ten
211+
states over a short window might fit in a single page, where extra chunks
212+
would only burn quota), this is a *deliberate* per-call knob rather than an
213+
automatic behavior or a process-wide environment variable — scoping it to a
214+
``with`` block keeps an aggressive setting from leaking into unrelated calls
215+
and accidentally spending quota. Outside any block the getters use the
216+
conservative default. Only the OGC getters (Water Data, NGWMN) read this;
217+
wrapping a legacy NWIS call in the block is a harmless no-op.
218+
219+
Parameters
220+
----------
221+
n : int
222+
The number of sub-requests to fan the whole call out into — a positive
223+
integer such as ``2``, ``8``, or ``32``. It caps the plan's *total*
224+
sub-request count (the cartesian product across every multi-value
225+
argument combined, not per argument), so several multi-value arguments
226+
cannot multiply past it. The cap is a ceiling, never exceeded: the
227+
actual count is bounded below by what the ~8 KB URL limit already
228+
forces and above by ``n``, so an ``n`` larger than the input allows
229+
simply yields one sub-request per value, and with several multi-value
230+
arguments the total may land somewhat below ``n`` because splits are
231+
whole (the plan can't always divide evenly onto ``n``); ``n=1`` asks
232+
for no extra fan-out.
233+
234+
Each sub-request fetches at least one page, so it costs at least one
235+
request against your hourly rate limit — a larger ``n`` spends more
236+
quota. And because how many sub-requests run *at once* is capped
237+
separately by ``API_USGS_CONCURRENT`` (default 32), an ``n`` beyond that
238+
adds quota without adding parallelism; the useful range is roughly ``2``
239+
up to ``API_USGS_CONCURRENT``.
240+
241+
Yields
242+
------
243+
None
244+
245+
Raises
246+
------
247+
ValueError
248+
If ``n`` is not a positive integer — raised on ``with`` entry, before
249+
any request is issued, so a bad value fails loudly rather than silently
250+
doing nothing.
251+
252+
Notes
253+
-----
254+
Fanning out carries the same consequences as the byte-limit chunking the
255+
getters already do for oversized requests; opting in just brings them to a
256+
request that would otherwise be a single call:
257+
258+
- ``max_rows``: each sub-request paginates up to ``max_rows`` rows
259+
independently, then the combined result is sorted and truncated to
260+
``max_rows``. So a call with ``max_rows`` set returns a *different*
261+
(though still valid and deterministically sorted) row set inside a
262+
``parallel_chunks`` block than without one — the cap is drawn from the
263+
union of the sub-requests, not a single stream. Don't pair a tight
264+
``max_rows`` preview with ``parallel_chunks`` if you need exactly the
265+
rows the un-fanned call would return.
266+
- Resumability: a single request either fully succeeds or fully fails,
267+
but a fanned-out call can fail partway (e.g. a mid-call rate-limit) and
268+
raise a resumable :class:`~dataretrieval.exceptions.ChunkInterrupted`
269+
(or ``QuotaExhausted``) carrying the completed sub-requests, which you
270+
finish with ``exc.call.resume()``.
271+
- Cross-sub-request de-duplication keys on the feature ``id``; features
272+
with no ``id`` can't be deduped, so overlapping filter clauses split
273+
across chunks may yield duplicate rows.
274+
275+
Examples
276+
--------
277+
>>> from dataretrieval import waterdata
278+
>>> with waterdata.parallel_chunks(32):
279+
... df, md = waterdata.get_daily(
280+
... monitoring_location_id=many_sites, parameter_code="00060"
281+
... ) # doctest: +SKIP
282+
283+
See Also
284+
--------
285+
ChunkPlan._refine : the planning-side effect of ``n``.
286+
"""
287+
# Fail loudly on a bad ``n`` at ``with`` entry, before any request. Shared
288+
# rules with ``max_rows`` via the helper (accepts numpy ints, rejects bool).
289+
_require_positive_int(n, "parallel_chunks(n)", examples="2, 8, 32")
290+
with _parallel_chunks(n):
291+
yield
292+
293+
175294
class ChunkedCall:
176295
"""
177296
Stateful handle for a chunked call.
@@ -591,8 +710,9 @@ def multi_value_chunked(
591710
``async def fetch(args) -> (df, response)``, and drives it to
592711
completion via :meth:`ChunkedCall.resume`. The plan splits multi-value
593712
list params and the cql-text filter so each sub-request URL fits the
594-
byte limit; an already-fitting request is a one-step plan. See the
595-
module docstring for the concurrency model.
713+
byte limit; an already-fitting request is a one-step plan, unless an
714+
active :func:`parallel_chunks` block asks the plan to fan out more
715+
finely. See the module docstring for the concurrency model.
596716
597717
Parameters
598718
----------
@@ -636,7 +756,14 @@ def wrapper(
636756
finalize: _Finalize = _passthrough_result,
637757
) -> tuple[pd.DataFrame, Any]:
638758
limit = _OGC_URL_BYTE_LIMIT if url_limit is None else url_limit
639-
plan = ChunkPlan(args, build_request, limit)
759+
# Read the parallel_chunks dial ``n`` from the ambient set by
760+
# ``parallel_chunks`` (1 = off outside any such block; otherwise the
761+
# requested total sub-request cap). It only affects *planning*, done
762+
# here up front, so a later resume — which re-issues the
763+
# already-planned sub-requests — needs no snapshot.
764+
plan = ChunkPlan(
765+
args, build_request, limit, max_chunks=_parallel_chunks.get()
766+
)
640767
retry_policy = RetryPolicy.from_env()
641768
# The concurrency cap is resolved inside ``resume()`` from
642769
# ``API_USGS_CONCURRENT``; ``1`` is a sequential gather,

dataretrieval/ogc/engine.py

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
import functools
2828
import json
2929
import logging
30-
import numbers
3130
import re
3231
from collections.abc import (
3332
AsyncIterator,
@@ -59,6 +58,7 @@
5958
_default_headers,
6059
_get,
6160
_network_error,
61+
_require_positive_int,
6262
)
6363

6464
# Set up logger for this module
@@ -876,18 +876,12 @@ def get_ogc_data(
876876
- Handles optional arguments such as `convert_type`.
877877
- Applies column cleanup and reordering based on service and properties.
878878
"""
879-
# Enforce a genuine positive integer: a float (even ``10.0``) or ``bool``
880-
# would pass a bare ``< 1`` check and then crash deep in
879+
# Enforce a genuine positive integer up front: a float (even ``10.0``) or
880+
# ``bool`` would pass a bare ``< 1`` check and then crash deep in
881881
# ``pd.DataFrame.head`` with an opaque ``TypeError`` after HTTP I/O has
882-
# already fired. ``numbers.Integral`` (not ``int``) so numpy integers —
883-
# e.g. ``max_rows`` derived from a numpy/pandas computation — are accepted;
884-
# ``bool`` is an ``Integral`` subtype, so exclude it explicitly.
885-
if max_rows is not None and (
886-
not isinstance(max_rows, numbers.Integral)
887-
or isinstance(max_rows, bool)
888-
or max_rows < 1
889-
):
890-
raise ValueError(f"max_rows must be a positive integer (got {max_rows!r}).")
882+
# already fired. Shared with ``parallel_chunks(n)`` via the helper.
883+
if max_rows is not None:
884+
_require_positive_int(max_rows, "max_rows")
891885

892886
if dialect is None:
893887
dialect = _DEFAULT_DIALECT

0 commit comments

Comments
 (0)