feat(waterdata): cap chunk_granularity by total sub-requests; document + benchmark

Cap chunk_granularity on the plan's TOTAL sub-request count (2/8/32) rather than per multi-value axis, so several multi-value arguments can't multiply past the ceiling. ChunkPlan._refine now splits the largest splittable chunk across every axis round-robin; behavior is identical for the common single-axis query.

Add a 'Speeding up large downloads' usage section to the README (Water Data API) with a measured cold-cache benchmark: parallelizing a large paginated pull's sub-requests gave ~6x (production page size, cold) up to ~12x (latency-bound) on 271 Ohio discharge sites.

Temper the 'quota-neutral' claim in the docstring, NEWS, and user guide: a finer split is only ~quota-neutral when each sub-request still spans many pages; otherwise each chunk's partial final page adds some requests.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: thodson-usgs

NEWS.md

@@ -1,4 +1,4 @@
-**07/01/2026:** Added `waterdata.chunk_granularity(...)` — a context manager to control how finely the OGC `waterdata` getters split multi-value requests. By default the getters chunk only as much as the server's ~8 KB URL limit forces (the fewest sub-requests). But because every sub-request paginates, splitting a large result further is usually quota-neutral, so a caller who *knows* their pull is large can opt into a finer split for smoother progress, more even concurrency, and a smaller unit of retry/resume: `with waterdata.chunk_granularity("high"): df, md = waterdata.get_daily(monitoring_location_id=many_sites)`. The level is one of `"low"`, `"medium"`, or `"high"` — typed as `waterdata.GranularityLevel` (a `typing.Literal`), so a type checker rejects any other value and an invalid string raises `ValueError` at the `with`. Each level caps how many sub-chunks a multi-value argument is split into — `"low"` / `"medium"` / `"high"` cap at 2 / 8 / 32. That ceiling is a granularity constant, deliberately independent of the fan-out concurrency (`API_USGS_CONCURRENT`): how finely a query splits is orthogonal to how many sub-requests run at once. Capping the aggressive end at 32 bounds the blast radius — an accidental `"high"` on a huge list can't explode into thousands of sub-requests. There is no "off" level: not entering the block *is* off. It is deliberately a scoped `with` block rather than an automatic behavior or an environment variable, since the library can't tell in advance whether a query is large — a short-window query might fit in a single page, where extra chunks would only burn quota. Also exported at the top level as `dataretrieval.chunk_granularity`.
+**07/01/2026:** Added `waterdata.chunk_granularity(...)` — a context manager to control how finely the OGC `waterdata` getters split multi-value requests. By default the getters chunk only as much as the server's ~8 KB URL limit forces (the fewest sub-requests). But because every sub-request paginates, splitting a large result further costs little or no extra quota when each sub-request still spans many pages, so a caller who *knows* their pull is large can opt into a finer split for smoother progress, more even concurrency, and a smaller unit of retry/resume: `with waterdata.chunk_granularity("high"): df, md = waterdata.get_daily(monitoring_location_id=many_sites)`. The level is one of `"low"`, `"medium"`, or `"high"` — typed as `waterdata.GranularityLevel` (a `typing.Literal`), so a type checker rejects any other value and an invalid string raises `ValueError` at the `with`. Each level caps the *total* number of sub-requests the call is split into — `"low"` / `"medium"` / `"high"` cap at 2 / 8 / 32 overall, across every multi-value argument combined, not per argument. That ceiling is a granularity constant, deliberately independent of the fan-out concurrency (`API_USGS_CONCURRENT`): how finely a query splits is orthogonal to how many sub-requests run at once. Capping the aggressive end at 32 bounds the blast radius — an accidental `"high"` on a huge list, or a call with several multi-value arguments, can't explode into thousands of sub-requests. There is no "off" level: not entering the block *is* off. It is deliberately a scoped `with` block rather than an automatic behavior or an environment variable, since the library can't tell in advance whether a query is large — a short-window query might fit in a single page, where extra chunks would only burn quota. Also exported at the top level as `dataretrieval.chunk_granularity`.
 
 **06/23/2026:** **Breaking change (1.2.0):** the minimum supported Python is now **3.10** (`requires-python = ">=3.10"`). 3.9 support was already effectively broken — the `waterdata` module's dependencies (`anyio`, the test stack) require 3.10+, and the `waterdata` test modules already skipped on <3.10. `anyio` is now declared as a direct dependency (it is imported directly by `waterdata`), and the CI/ruff/mypy targets move to 3.10. Also fully removed the deprecated `variable_info` metadata property: the `NWIS_Metadata` override only warned and returned `None` (it relied on the defunct `get_pmcodes`), and the `BaseMetadata` abstract is gone too since nothing implemented it — accessing `.variable_info` now raises `AttributeError`. `site_info` is unaffected.
 

README.md

@@ -105,6 +105,54 @@ df, metadata = waterdata.get_continuous(
 print(f"Retrieved {len(df)} continuous gage height measurements")
 ```
 
+#### Speeding up large downloads with `chunk_granularity`
+
+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**. `chunk_granularity` opts a single call into that
+finer split. 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.chunk_granularity("high"):        # "low" | "medium" | "high"
+    df, md = waterdata.get_daily(
+        monitoring_location_id=sites["monitoring_location_id"].tolist(),
+        parameter_code="00060",                  # discharge
+        time="2004-01-01/2023-12-31",
+    )
+```
+
+`"high"` fans a call out into up to 32 sub-requests, `"medium"` up to 8, and
+`"low"` up to 2 — a fixed ceiling, so an accidental setting on a huge list can't
+explode into thousands of requests.
+
+Benchmark — 271 Ohio discharge sites (`get_daily`, `parameter_code="00060"`),
+cold cache, each level run against its own time window so results are not
+cache-served, with the page size fixed so every level fetches roughly the same
+number of pages (isolating the effect of parallelism):
+
+| level          | parallelism | pages | wall-clock                | speedup |
+| -------------- | ----------- | ----- | ------------------------- | ------- |
+| `default`      | 1           | 32    | 23.9 s / 27.4 s (2 runs)  | 1×      |
+| `"medium"` (8) | 8           | 37    | 4.1 s / 4.3 s             | ~6×     |
+| `"high"` (32)  | 32          | 44    | 2.0 s                     | ~12×    |
+
+The gain comes from overlapping each sub-request's per-page latency and
+server-side work — a genuinely large pull at the default page size shows a
+similar multi-fold speedup. The extra sub-requests each cost one request against
+your hourly [rate limit](https://api.waterdata.usgs.gov/signup/), so reserve the
+aggressive levels 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/ogc/chunking.py

@@ -182,8 +182,9 @@ def get_active_client() -> httpx.AsyncClient | None:
 # Chunk-granularity dial: opt-in to fan a query out *more finely* than the byte
 # limit alone requires. Scoped to a ``with chunk_granularity(...):`` block (a
 # ContextVar), deliberately NOT an env var (see :func:`chunk_granularity` for
-# why). The ambient holds the resolved cap on sub-chunks per axis; ``0`` (the
-# default, outside any block) means "chunk only as much as the byte limit needs".
+# why). The ambient holds the resolved cap on the plan's total sub-request
+# count; ``0`` (the default, outside any block) means "chunk only as much as
+# the byte limit needs".
 _granularity: Ambient[int] = Ambient("ogc_chunk_granularity", 0)
 
 #: The three accepted granularity levels, as a typing ``Literal`` so a type
@@ -196,12 +197,14 @@ def get_active_client() -> httpx.AsyncClient | None:
 #: modules).
 _VALID_LEVELS: tuple[GranularityLevel, ...] = get_args(GranularityLevel)
 
-# Granularity's own ceiling on sub-chunks per axis — deliberately NOT tied to the
-# concurrency default: fan-out *volume* (how many sub-requests a query becomes)
-# is orthogonal to how many run at once (``API_USGS_CONCURRENT``). 32 is a sane
-# single-axis ceiling that also bounds the blast radius of an accidental
-# ``"high"`` on a very long list; the milder levels are a quarter and a
-# sixteenth of it (the three are spaced 4x apart).
+# Granularity's own ceiling on the plan's total sub-request count —
+# deliberately NOT tied to the concurrency default: fan-out *volume* (how many
+# sub-requests a query becomes) is orthogonal to how many run at once
+# (``API_USGS_CONCURRENT``). 32 is a sane ceiling on the whole call (across
+# every axis, not per axis — see ``ChunkPlan._refine``) that bounds the blast
+# radius of an accidental ``"high"`` on a very long list or several multi-value
+# arguments at once; the milder levels are a quarter and a sixteenth of it (the
+# three are spaced 4x apart).
 _GRANULARITY_MAX_CHUNKS = 32
 _GRANULARITY_LEVELS: dict[str, int] = {
     "low": _GRANULARITY_MAX_CHUNKS // 16,  # 2
@@ -212,7 +215,7 @@ def get_active_client() -> httpx.AsyncClient | None:
 
 def _resolve_granularity(level: GranularityLevel) -> int:
     """
-    Map a granularity level name to its per-axis sub-chunk cap.
+    Map a granularity level name to its total sub-request cap.
 
     Parameters
     ----------
@@ -222,7 +225,7 @@ def _resolve_granularity(level: GranularityLevel) -> int:
     Returns
     -------
     int
-        The maximum sub-chunks per multi-value axis for that level
+        The maximum total sub-requests for the whole call at that level
         (``2`` / ``8`` / ``32``).
 
     Raises
@@ -247,12 +250,16 @@ def chunk_granularity(level: GranularityLevel) -> Iterator[None]:
     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 is
-    usually quota-neutral (ten states pulled as one under-limit request page
-    just as many times as ten per-state requests would). This context manager
-    lets a caller who *knows* their pull is large ask for that finer split —
-    trading the same pages for more, smaller sub-requests, which gives smoother
-    progress, more even concurrency, and a smaller unit of retry/resume.
+    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
@@ -267,15 +274,19 @@ def chunk_granularity(level: GranularityLevel) -> Iterator[None]:
     Parameters
     ----------
     level : {"low", "medium", "high"}
-        How aggressively to chunk within the block. Each level caps how many
-        sub-chunks a single multi-value argument is split into — ``2`` / ``8`` /
-        ``32`` for ``"low"`` / ``"medium"`` / ``"high"`` — or one value per
-        sub-request once the argument has fewer values than the cap. The ceiling
-        is fixed (it is *not* tied to ``API_USGS_CONCURRENT``: how finely a query
-        splits is orthogonal to how many sub-requests run at once); capping
-        ``"high"`` at ``32`` keeps an accidental aggressive level on a very long
-        list from exploding into thousands of sub-requests. (With several
-        multi-value arguments the per-argument counts still multiply.)
+        How aggressively to chunk within the block. Each level caps the
+        *total* number of sub-requests the call is split into — ``2`` / ``8``
+        / ``32`` for ``"low"`` / ``"medium"`` / ``"high"`` — or one
+        sub-request per remaining atom once that's fewer than the cap. The
+        cap applies to the whole call, not per multi-value argument: with
+        several multi-value arguments the axes are refined together against
+        the same shared ceiling (their cartesian product does not multiply
+        past it). The ceiling is fixed (it is *not* tied to
+        ``API_USGS_CONCURRENT``: how finely a query splits is orthogonal to
+        how many sub-requests run at once); capping ``"high"`` at ``32`` keeps
+        an accidental aggressive level on a very long list — or several
+        multi-value arguments at once — from exploding into thousands of
+        sub-requests.
 
     Yields
     ------

dataretrieval/ogc/planning.py

@@ -330,12 +330,13 @@ class ChunkPlan:
         Byte budget for the request (URL + body) — a hard ceiling every
         sub-request must fit.
     max_chunks_per_axis : int, optional
-        Soft cap on sub-chunks per multi-value axis (default ``0`` = off).
+        Soft cap on the plan's total sub-request count (default ``0`` = off).
         ``0`` chunks only as much as ``url_limit`` requires — the most
-        conservative plan, fewest sub-requests. A positive cap fans each axis
-        out into up to ``min(len(atoms), max_chunks_per_axis)`` pieces (never
-        fewer than the byte budget already forces), so a large multi-page pull
-        is issued as more, smaller sub-requests. Set from the resolved
+        conservative plan, fewest sub-requests. A positive cap fans the plan
+        out to up to ``max_chunks_per_axis`` 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. Set from the resolved
         :func:`~dataretrieval.ogc.chunking.chunk_granularity` level; see
         :meth:`_refine`.
 
@@ -492,38 +493,50 @@ def _plan(
 
     def _refine(self, max_chunks_per_axis: int) -> None:
         """
-        Fan each axis out more finely than the byte budget alone requires —
+        Fan the plan out more finely than the byte budget alone requires —
         the granularity dial (see
         :func:`~dataretrieval.ogc.chunking.chunk_granularity` for why a caller
         would want this).
 
-        Each axis is split until it holds at least
-        ``min(len(atoms), max_chunks_per_axis)`` chunks — up to the cap, or one
-        atom per chunk for a shorter axis. Purely additive — only ever *splits*
-        existing chunks, so the byte pass's work and the ``url_limit`` invariant
-        are both preserved (an axis the byte pass already split past the cap is
-        left alone), and it never raises. A no-op at ``max_chunks_per_axis <= 0``.
+        Caps the plan's *total* sub-request count (:attr:`total`, the
+        cartesian product across all axes) at ``max_chunks_per_axis``, not
+        each axis independently — with several multi-value axes, a cap of 32
+        still means at most 32 sub-requests overall, not ``32 ** n_axes``.
+        Each split picks the single largest splittable chunk across *every*
+        axis (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_per_axis <= 0``.
 
         Parameters
         ----------
         max_chunks_per_axis : int
-            Soft cap on sub-chunks per axis (``0`` = off), the resolved
-            granularity level. A shorter axis simply saturates at one atom per
-            chunk.
+            Soft cap on the plan's total sub-request count (``0`` = off),
+            the resolved granularity level. Despite the name — kept for the
+            public dial's "sub-chunks per multi-value argument" framing,
+            which matches this cap exactly in the common single-axis case —
+            multi-axis plans are capped on the product, not per axis.
         """
         if max_chunks_per_axis <= 0:
             return
-        for axis in self.axes:
-            chunks = self.chunks[axis.arg_key]
-            target = min(len(axis.atoms), max_chunks_per_axis)
-            # ``target <= len(atoms)`` guarantees a splittable chunk each pass, so
-            # this reaches exactly ``target`` and terminates. Split the chunk with
-            # the most *atoms* (``_plan`` splits by *bytes*): here we even out
-            # cardinality for smooth fan-out, not URL size. Ties take the lowest
-            # index, keeping the split deterministic.
-            while len(chunks) < target:
-                idx, _ = max(enumerate(chunks), key=lambda kv: len(kv[1]))
-                _split_at(chunks, idx)
+        while self.total < max_chunks_per_axis:
+            # Largest splittable chunk across every axis; a chunk of size 1
+            # can't be split further. ``max`` with a 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:
+                for idx, chunk in enumerate(self.chunks[axis.arg_key]):
+                    if len(chunk) <= 1:
+                        continue
+                    if len(chunk) > candidate_size:
+                        candidate, candidate_size = (axis, idx), len(chunk)
+            if candidate is None:
+                return  # every axis saturated at one atom per chunk
+            axis, idx = candidate
+            _split_at(self.chunks[axis.arg_key], idx)
 
     def _worst_case_args(self) -> dict[str, Any]:
         """