feat(waterdata): add parallel_chunks to control OGC chunk fan-out

[thodson-usgs]

Summary

Adds waterdata.parallel_chunks(n) — a context manager to control how finely the OGC waterdata (and NGWMN) getters split multi-value requests into chunked sub-requests: it fans a call out into n parallel sub-requests.

Today the chunker splits a request only as much as the server's ~8 KB URL-byte limit forces — the fewest sub-requests. 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: ten states pulled as one request then page nearly as many times as ten per-state requests would. In that situation finer chunks buy smoother progress, more even concurrency, and a smaller unit of retry/resume. (When a split instead leaves each sub-request only a page or two, its partial final page is extra — so on a smaller pull finer chunks do add some requests; reserve a large n for pulls you know are large.)

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), so this is a deliberate, scoped knob the user sets with their own judgment — not automatic, and not a process-wide env var (which would be a quota footgun). Scoping it to a with block keeps an aggressive setting from leaking into unrelated calls.

from dataretrieval import waterdata

# Default: chunk only as much as the URL limit needs.
df, md = waterdata.get_daily(monitoring_location_id=many_sites)

# Opt into a finer split for a pull you know is large:
with waterdata.parallel_chunks(32):
    df, md = waterdata.get_daily(
        monitoring_location_id=many_sites, parameter_code="00060"
    )

Measured speedup (271 Ohio discharge sites, get_daily, cold cache, each n on its own time window): a large paginated pull ran ~6× faster at the production page size and up to ~12× faster when per-request latency dominates — the fan-out parallelizes what is otherwise sequential pagination. See the new README section for the table.

The dial

parallel_chunks(n) takes a positive integer — the number of sub-requests to fan the whole call out into. 2, 8, and 32 are typical values; a non-integer, non-positive value, or a bool raises ValueError at the with.

with waterdata.parallel_chunks(32):      # fan out into 32 sub-requests
    df, md = waterdata.get_daily(
        monitoring_location_id=many_sites, parameter_code="00060"
    )

n caps the plan's total sub-request count — the cartesian product across every multi-value argument combined, not each argument independently — so several multi-value arguments can't multiply past it. The actual count is bounded below by what the ~8 KB URL limit already forces and above by how many values there are to split (so an n larger than the input allows just yields one sub-request per value); n=1 asks for no extra fan-out. There is no "off" — not entering the block is off.

Cost and the useful range. 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 how many sub-requests run at once is capped separately by API_USGS_CONCURRENT (default 32), so an n beyond that adds quota without adding parallelism. The useful range is therefore roughly 2 up to API_USGS_CONCURRENT. Fan-out volume and simultaneity stay deliberately independent controls.

Exported as waterdata.parallel_chunks and, for parity with ChunkInterrupted, at the top level as dataretrieval.parallel_chunks.

Implementation

• ChunkPlan._refine(max_chunks) — a soft pass that runs after the existing hard byte pass (_plan). It splits the largest splittable chunk across every axis (round-robin) until the plan's total sub-request count reaches the cap, only ever splitting further (via the shared _split_at primitive), so the url_limit invariant always holds and it never raises. A no-op at cap 0, so the default path is byte-for-byte unchanged (passthrough preserved).
• n is read from an Ambient (contextvar) set by the context manager, at plan-construction time inside multi_value_chunked's wrapper — so a later resume() (which re-issues already-planned sub-requests) needs no extra snapshot.
• parallel_chunks(n) validates n inline (a positive int, rejecting bool/float/str/None) and publishes it on the ambient; ChunkPlan reads that plain int as max_chunks. There is no level→cap lookup table — the value the user passes is the cap.

Tests & checks

• Parallel-chunks unit + end-to-end tests in tests/waterdata_chunking_test.py, plus an export-surface test; covers the cap→pieces ramp/saturation (with cover-partition checks), support for an arbitrary n (not just 2/8/32), the total-cap product bound across multiple multi-value axes, the guardrail on long axes, byte-budget preservation, filter-axis + multi-axis behavior, n-validation (rejecting 0/negative/float/str/None/bool), n=1 as an explicit no-op, context-manager scoping/nesting, and the passthrough-unchanged default.
• ruff check, ruff format --check, and mypy --strict all clean.
• NEWS.md, a userguide section, and a README usage section + benchmark table updated.

Note

The dial has been through a few shapes: an off/1–5/max scale, a per-axis cap derived from the concurrency width, then a fixed "low"/"medium"/"high" enum — and is now a plain integer n, the number of sub-requests to fan out into. The integer is more expressive than three arbitrary tiers, gives precise control, and mirrors the int-valued API_USGS_CONCURRENT; 2/8/32 survive only as documented examples. The knob was also renamed chunk_granularity → parallel_chunks: "parallel_chunks" says what it does — fan a query into more, parallel sub-requests — without colliding with API_USGS_CONCURRENT, the separate in-flight cap. Still a draft.

🤖 Generated with Claude Code