11# SPDX-License-Identifier: Apache-2.0
2- """O(rank) selection costs; lazy subdivision without a global chunk inventory."""
2+ """Bounded selection costs; lazy subdivision without a global chunk inventory."""
33
44from __future__ import annotations
55
@@ -34,13 +34,30 @@ def __post_init__(self) -> None:
3434 raise TypeError("strided_only must be a boolean")
3535
3636
37+ @dataclass(frozen=True)
38+ class GatherWindow:
39+ """Independent grid regions with one live decoded chunk per worker."""
40+
41+ region_shape: tuple[int, ...]
42+ workers: int
43+
44+ def __post_init__(self) -> None:
45+ if (isinstance(self.workers, bool) or not isinstance(self.workers, int)
46+ or not 1 <= self.workers <= 256):
47+ raise ValueError("Gather workers must be an integer from 1 to 256")
48+ if (not self.region_shape or any(isinstance(v, bool) or not isinstance(v, int)
49+ or v < 1 for v in self.region_shape)):
50+ raise ValueError("Gather regions require positive integer dimensions")
51+
52+
3753@dataclass(frozen=True)
3854class ArraySpec:
3955 shape: tuple[int, ...]
4056 chunks: tuple[int, ...]
4157 itemsize: int
4258 decode_window: DecodeWindow | None = None
4359 dense_gather: bool = False
60+ gather_window: GatherWindow | None = None
4461
4562 def __post_init__(self) -> None:
4663 if not isinstance(self.dense_gather, bool):
@@ -53,6 +70,17 @@ def __post_init__(self) -> None:
5370 if any(isinstance(v, bool) or not isinstance(v, int) or v < minimum for v in values):
5471 raise ValueError("shape, chunks and itemsize require valid integer dimensions")
5572
73+ gather = self.gather_window
74+ if gather is not None:
75+ if not isinstance(gather, GatherWindow):
76+ raise TypeError("gather_window must be a GatherWindow or None")
77+ if (self.dense_gather or self.decode_window is None
78+ or self.decode_window.max_chunks != 1):
79+ raise ValueError("Parallel gathering requires a single-chunk native decode window")
80+ if (len(gather.region_shape) != len(self.chunks)
81+ or any(r % c for r, c in zip(gather.region_shape, self.chunks, strict=True))):
82+ raise ValueError("Gather regions must be a whole number of chunks on each axis")
83+
5684 @property
5785 def logical_bytes(self) -> int:
5886 return math.prod(self.shape) * self.itemsize
@@ -67,8 +95,9 @@ class ReadPlan:
6795 decoded_bytes_upper_bound: int
6896 peak_decoded_bytes_upper_bound: int
6997 intermediate_bytes_upper_bound: int = 0
70- strategy: Literal["native", "dense-gather"] = "native"
98+ strategy: Literal["native", "dense-gather", "parallel-gather" ] = "native"
7199 read_selection: Selection | None = None
100+ gather_workers: int = 1
72101
73102 @property
74103 def decode_amplification(self) -> float:
@@ -129,9 +158,29 @@ def plan(spec: ArraySpec, selection: Selection, *, buffer_bytes: int | None = No
129158 touched * math.prod(spec.chunks) * spec.itemsize,
130159 live_chunks * math.prod(spec.chunks) * spec.itemsize,
131160 )
132- if (not spec.dense_gather or buffer_bytes is None or not request.output_bytes
161+ if (buffer_bytes is None or not request.output_bytes
133162 or not any(isinstance(s, slice) and s.step > 1 for s in normalized)):
134163 return request
164+ gather = spec.gather_window
165+ if gather is not None:
166+ regions = plan(ArraySpec(spec.shape, gather.region_shape, spec.itemsize),
167+ request.selection).chunks_touched
168+ region_output = spec.itemsize * math.prod(
169+ min(_length(s), (r + s.step - 1) // s.step)
170+ for s, r in zip(request.selection, gather.region_shape, strict=True)
171+ if isinstance(s, slice)
172+ )
173+ for workers in range(min(gather.workers, regions), 1, -1):
174+ candidate = replace(
175+ request, strategy="parallel-gather", gather_workers=workers,
176+ peak_decoded_bytes_upper_bound=workers * request.peak_decoded_bytes_upper_bound,
177+ intermediate_bytes_upper_bound=min(request.output_bytes, workers * region_output),
178+ )
179+ if candidate.buffer_bytes_upper_bound <= buffer_bytes:
180+ return candidate
181+ return request
182+ if not spec.dense_gather:
183+ return request
135184 bounds = tuple(
136185 s if isinstance(s, int) else slice(s.start, s.start + (_length(s) - 1) * s.step + 1, 1)
137186 for s in normalized
0 commit comments