refactor(waterdata): tidy parallel_chunks (rename param, align validation, dedup test)

/simplify cleanup: (1) rename ChunkPlan's max_chunks_per_axis -> max_chunks — the cap is on the plan's TOTAL sub-request count, not per axis, so the old name needed apologetic docstrings; dropped them. (2) Validate parallel_chunks(n) with numbers.Integral (matching the max_rows guard in engine.py), so numpy integers are accepted like the sibling validator. (3) Fold the n=1 no-op test into the parametrized arbitrary-n test, removing a duplicated fetch fixture.

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

Author: thodson-usgs

dataretrieval/ogc/chunking.py

@@ -73,6 +73,7 @@
 
 import asyncio
 import functools
+import numbers
 import os
 from collections.abc import Callable, Iterator
 from contextlib import contextmanager
@@ -258,8 +259,9 @@ def parallel_chunks(n: int) -> Iterator[None]:
     --------
     ChunkPlan._refine : the planning-side effect of ``n``.
     """
-    # ``bool`` is an ``int`` subclass but nonsensical here; reject it explicitly.
-    if not isinstance(n, int) or isinstance(n, bool) or n < 1:
+    # Accept any integer type (incl. numpy ints, mirroring ``max_rows``); ``bool``
+    # is an ``Integral`` subclass but nonsensical here, so reject it explicitly.
+    if not isinstance(n, numbers.Integral) or isinstance(n, bool) or n < 1:
         raise ValueError(
             f"parallel_chunks(n): n must be a positive integer (e.g. 2, 8, 32); "
             f"got {n!r}."
@@ -739,7 +741,7 @@ def wrapper(
             # 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_per_axis=_parallel_chunks.get()
+                args, build_request, limit, max_chunks=_parallel_chunks.get()
             )
             retry_policy = RetryPolicy.from_env()
             # The concurrency cap is resolved inside ``resume()`` from

dataretrieval/ogc/planning.py

@@ -329,11 +329,11 @@ class ChunkPlan:
     url_limit : int
         Byte budget for the request (URL + body) — a hard ceiling every
         sub-request must fit.
-    max_chunks_per_axis : int, optional
+    max_chunks : int, optional
         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 the plan
-        out to up to ``max_chunks_per_axis`` sub-requests overall (the
+        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. Set from the
@@ -370,7 +370,7 @@ def __init__(
         args: dict[str, Any],
         build_request: Callable[..., httpx.Request],
         url_limit: int,
-        max_chunks_per_axis: int = 0,
+        max_chunks: int = 0,
     ) -> None:
         self.args = args
         self.axes: list[_Axis] = []
@@ -423,9 +423,9 @@ def __init__(
         # 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. Only when ``max_chunks_per_axis`` asks for extra fan-out do
+        # verbatim. Only when ``max_chunks`` asks for extra fan-out do
         # we set the axes up to be refined below.
-        if fits and max_chunks_per_axis <= 0:
+        if fits and max_chunks <= 0:
             return
 
         self.axes = axes
@@ -436,8 +436,8 @@ def __init__(
             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_per_axis <= 0``.
-        self._refine(max_chunks_per_axis)
+        # satisfied; a no-op at ``max_chunks <= 0``.
+        self._refine(max_chunks)
 
         if self.canonical_url is None:
             # Original URL was un-constructable (httpx.InvalidURL); fall
@@ -491,15 +491,15 @@ def _plan(
                 )
             _split_at(self.chunks[biggest_axis.arg_key], biggest_idx)
 
-    def _refine(self, max_chunks_per_axis: int) -> None:
+    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).
 
         Caps the plan's *total* sub-request count (:attr:`total`, the
-        cartesian product across all axes) at ``max_chunks_per_axis``, not
+        cartesian product across all axes) at ``max_chunks``, 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*
@@ -508,20 +508,19 @@ def _refine(self, max_chunks_per_axis: int) -> None:
         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``.
+        ``max_chunks <= 0``.
 
         Parameters
         ----------
-        max_chunks_per_axis : int
+        max_chunks : int
             Soft cap on the plan's total sub-request count (``0`` = off) — the
-            ``parallel_chunks(n)`` value. The parameter name is legacy: the cap
-            is on the whole plan (the cartesian product across axes), which
-            matches "sub-chunks per argument" only in the common single-axis
-            case — multi-axis plans are capped on the product, not per axis.
+            ``parallel_chunks(n)`` value. The cap is on the whole plan (the
+            cartesian product across axes), not per axis, so several multi-value
+            axes can't multiply past it.
         """
-        if max_chunks_per_axis <= 0:
+        if max_chunks <= 0:
             return
-        while self.total < max_chunks_per_axis:
+        while self.total < max_chunks:
             # 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.

tests/waterdata_chunking_test.py

@@ -2115,30 +2115,30 @@ async def fetch(args):
 # ``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_per_axis``) directly;
+# ``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_zero_cap_preserves_passthrough():
-    """``max_chunks_per_axis=0`` (the default) must not perturb the existing
+    """``max_chunks=0`` (the default) 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, max_chunks_per_axis=0)
+    plan = ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=0)
     assert plan.axes == []
     assert plan.total == 1
     assert list(plan.iter_sub_args()) == [args]
 
 
 @pytest.mark.parametrize(
-    ("max_chunks_per_axis", "expected_pieces"),
+    ("max_chunks", "expected_pieces"),
     [(0, 1), (2, 2), (8, 8), (16, 10), (32, 10)],
 )
-def test_cap_ramps_then_saturates(max_chunks_per_axis, expected_pieces):
+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
@@ -2149,10 +2149,10 @@ def test_cap_ramps_then_saturates(max_chunks_per_axis, expected_pieces):
         {"monitoring_location_id": atoms},
         _fake_build,
         url_limit=8000,
-        max_chunks_per_axis=max_chunks_per_axis,
+        max_chunks=max_chunks,
     )
     assert plan.total == expected_pieces
-    if max_chunks_per_axis:
+    if max_chunks:
         flattened = [
             a for chunk in plan.chunks["monitoring_location_id"] for a in chunk
         ]
@@ -2170,7 +2170,7 @@ def test_cap_bounds_fan_out_for_a_long_axis():
         {"monitoring_location_id": atoms},
         _fake_build,
         url_limit=8000,
-        max_chunks_per_axis=high,
+        max_chunks=high,
     )
     assert plan.total == high
     flattened = [a for chunk in plan.chunks["monitoring_location_id"] for a in chunk]
@@ -2184,9 +2184,9 @@ def test_cap_below_byte_split_does_not_reduce_fan_out():
     # 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_per_axis=0)
+    baseline = ChunkPlan(args, _fake_build, url_limit=250, max_chunks=0)
     assert baseline.total > 2  # byte pass alone already fanned out past 2
-    refined = ChunkPlan(args, _fake_build, url_limit=250, max_chunks_per_axis=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
 
@@ -2197,8 +2197,8 @@ def test_cap_never_exceeds_the_byte_budget():
     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_per_axis=0)
-    plan = ChunkPlan(args, _fake_build, url_limit=limit, max_chunks_per_axis=32)
+    byte_only = ChunkPlan(args, _fake_build, url_limit=limit, max_chunks=0)
+    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
@@ -2210,7 +2210,7 @@ def test_cap_refines_the_filter_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_per_axis=4)
+    plan = ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=4)
     assert len(plan.chunks["filter"]) == 4  # min(8, 4)
     assert plan.total == 4
 
@@ -2225,7 +2225,7 @@ def test_cap_caps_the_total_across_axes():
         "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_per_axis=4)
+    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 (
@@ -2249,7 +2249,7 @@ def test_cap_bounds_fan_out_across_many_axes():
         "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_per_axis=high)
+    plan = ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=high)
     assert plan.total <= high
 
 
@@ -2259,7 +2259,7 @@ def test_cap_does_not_mask_unchunkable():
     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_per_axis=32)
+        ChunkPlan(args, _fake_build, url_limit=10, max_chunks=32)
 
 
 def test_parallel_chunks_publishes_n_on_the_ambient():
@@ -2298,22 +2298,6 @@ def test_parallel_chunks_rejects_non_positive_int(bad):
     assert _parallel_chunks.get() == 0
 
 
-def test_parallel_chunks_n1_is_no_extra_fan_out():
-    """``n=1`` is the explicit no-op: an under-limit request stays a single
-    passthrough call, exactly as if the block weren't entered."""
-    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(1):
-        fetch({"monitoring_location_id": sites})
-    assert calls == [8]  # one passthrough call carrying all sites
-
-
 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)``
@@ -2344,10 +2328,11 @@ async def fetch(args):
     assert sorted(df_fine["site"]) == sorted(sites)
 
 
-@pytest.mark.parametrize("n", [2, 3, 8])
+@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."""
+    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] = []