perf(zarr-indexing): joint chunk enumeration for correlated vindex maps

Candidate-chunk enumeration took the cartesian product of each correlated
ArrayMap's per-dimension distinct chunk ids and relied on intersect() to
filter untouched combinations. For a diagonal selection of P scattered
points that is P**2 intersect calls — quadratic in the number of selected
points, the same workload shape as zarr-developers#4174 (400 points:
~2.6s; 10k points: ~30min).

Group correlated maps jointly instead: broadcast their per-point chunk
ids, take the distinct rows (np.unique(axis=0), O(P log P)), and
enumerate exactly the touched combinations. Candidate slots now carry
chunk-coordinate tuples covering one or more output dimensions;
orthogonal/constant/slice dimensions keep their existing per-dimension
candidates. 400-point diagonal resolution drops from 2628ms to 14ms and
scales linearly.

Assisted-by: ClaudeCode:claude-fable-5

Author: d-v-b

packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py

@@ -138,35 +138,42 @@ def iter_chunk_transforms(
                 yield from _iter_sorted_1d_array_map(sorted_map, storage, dim_grid)
                 return
 
-    # Enumerate candidate chunks via the cartesian product of per-dimension
-    # candidate chunk ids, then for each candidate intersect the transform with
-    # the chunk domain (`transform.intersect` handles orthogonal and vectorized
-    # cases alike, filtering out combinations it does not actually touch).
+    # Enumerate candidate chunks via the cartesian product of per-slot candidate
+    # chunk ids, then for each candidate intersect the transform with the chunk
+    # domain (`transform.intersect` handles orthogonal and vectorized cases
+    # alike, filtering out combinations it does not actually touch).
     #
-    # Each dimension contributes exactly the chunk ids it can touch:
+    # A slot covers one or more output dimensions and contributes exactly the
+    # chunk-coordinate tuples those dimensions can touch:
     #
-    # - `ConstantMap`/`DimensionMap` dims contribute a contiguous `range` — a
-    #   single chunk for a constant, and the span between the first and last
-    #   chunk for a slice. These are already tight (or nearly so).
-    # - `ArrayMap` (fancy) dims contribute only the *distinct* chunk ids the
-    #   index array actually lands in (`np.unique`), never the dense
-    #   `range(min_chunk, max_chunk + 1)` between them. A sparse fancy selection
-    #   (e.g. two far-apart coordinates) would otherwise enumerate every chunk
-    #   in the bounding box, making resolution scale with grid size instead of
-    #   with the number of selected coordinates.
-    #
-    # For >= 2 correlated (vindex) ArrayMaps the per-dimension distinct sets
-    # over-approximate the *joint* touched set (their cartesian product includes
-    # combinations no single point lands in), but `intersect` filters those out,
-    # so the yielded chunks are identical either way — and the work stays bounded
-    # by the per-dimension distinct chunk counts, not the grid size.
-    chunk_candidates: list[Sequence[int]] = []
+    # - `ConstantMap`/`DimensionMap` dims each form their own slot with a
+    #   contiguous range — a single chunk for a constant, and the span between
+    #   the first and last chunk for a slice. These are already tight (or
+    #   nearly so).
+    # - Orthogonal `ArrayMap` (fancy) dims each form their own slot with only
+    #   the *distinct* chunk ids the index array actually lands in
+    #   (`np.unique`), never the dense `range(min_chunk, max_chunk + 1)`
+    #   between them. A sparse fancy selection (e.g. two far-apart coordinates)
+    #   would otherwise enumerate every chunk in the bounding box, making
+    #   resolution scale with grid size instead of with the number of selected
+    #   coordinates.
+    # - Correlated (vindex) `ArrayMap` dims share one *joint* slot holding the
+    #   distinct chunk-coordinate tuples the points actually land in. The
+    #   cartesian product of their per-dimension distinct sets would include
+    #   combinations no point touches — quadratic in the number of selected
+    #   points for a diagonal selection — while the joint distinct set is
+    #   bounded by the point count (see zarr-python gh-4174).
+    correlated_dims: list[int] = []
+    correlated_chunk_ids: list[np.ndarray[Any, np.dtype[np.intp]]] = []
+    slot_dims: list[tuple[int, ...]] = []
+    slot_candidates: list[Sequence[tuple[int, ...]]] = []
     for out_dim, m in enumerate(transform.output):
         dg = dim_grids[out_dim]
         if isinstance(m, ConstantMap):
             # Single chunk
             c = dg.index_to_chunk(m.offset)
-            chunk_candidates.append((c,))
+            slot_dims.append((out_dim,))
+            slot_candidates.append(((c,),))
         elif isinstance(m, DimensionMap):
             d = m.input_dimension
             dim_lo = transform.domain.inclusive_min[d]
@@ -181,25 +188,48 @@ def iter_chunk_transforms(
                 s_max = m.offset + m.stride * dim_lo
             first = dg.index_to_chunk(s_min)
             last = dg.index_to_chunk(s_max)
-            chunk_candidates.append(range(first, last + 1))
+            slot_dims.append((out_dim,))
+            slot_candidates.append([(c,) for c in range(first, last + 1)])
         else:
             # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap).
             # Storage coordinates were already computed for a correlated 1-D map.
             storage = (
                 array_map_1d[1] if array_map_1d is not None else m.offset + m.stride * m.index_array
             )
-            flat = storage.ravel().astype(np.intp)
-            if flat.size == 0:
+            if storage.size == 0:
                 # Empty fancy selection: no coordinates, so no chunks are touched.
                 return
-            chunk_ids = dg.indices_to_chunks(flat)
-            # Enumerate only the distinct chunks the coordinates land in.
-            chunk_candidates.append([int(c) for c in np.unique(chunk_ids)])
+            # Keep the index-array shape: correlated maps broadcast against each
+            # other below, and raveling first would lose the singleton axes.
+            chunk_ids = dg.indices_to_chunks(storage.astype(np.intp))
+            if m.input_dimension is None:
+                correlated_dims.append(out_dim)
+                correlated_chunk_ids.append(chunk_ids)
+            else:
+                slot_dims.append((out_dim,))
+                slot_candidates.append([(int(c),) for c in np.unique(chunk_ids)])
+
+    if len(correlated_dims) == 1:
+        slot_dims.append((correlated_dims[0],))
+        slot_candidates.append([(int(c),) for c in np.unique(correlated_chunk_ids[0])])
+    elif len(correlated_dims) >= 2:
+        # Group the points jointly: distinct rows of the per-point chunk
+        # coordinates, O(points log points) regardless of grid size.
+        broadcast = np.broadcast_arrays(*correlated_chunk_ids)
+        stacked = np.stack([b.ravel() for b in broadcast], axis=1)
+        joint = np.unique(stacked, axis=0)
+        slot_dims.append(tuple(correlated_dims))
+        slot_candidates.append([tuple(int(c) for c in row) for row in joint])
 
     import itertools
 
-    for chunk_coords_tuple in itertools.product(*chunk_candidates):
-        chunk_coords = tuple(int(c) for c in chunk_coords_tuple)
+    output_rank = len(transform.output)
+    for combo in itertools.product(*slot_candidates):
+        chunk_coords_list = [0] * output_rank
+        for dims, part in zip(slot_dims, combo, strict=True):
+            for d, c in zip(dims, part, strict=True):
+                chunk_coords_list[d] = c
+        chunk_coords = tuple(chunk_coords_list)
 
         # Build the chunk domain in storage space
         chunk_min: list[int] = []

packages/zarr-indexing/tests/test_chunk_resolution.py

@@ -329,17 +329,16 @@ def test_2d_orthogonal_enumerates_only_touched_chunks(
         assert coords == [(0, 0), (0, 999), (999, 0), (999, 999)]
         assert calls["n"] == 4
 
-    def test_2d_correlated_vindex_enumerates_per_dim_distinct_chunks(
+    def test_2d_correlated_vindex_enumerates_joint_touched_chunks(
         self, monkeypatch: pytest.MonkeyPatch
     ) -> None:
         """Two correlated (vindex) coordinate arrays scatter to 2 diagonal chunks.
 
         The two points (1, 2) and (3997, 3998) touch chunks (0, 0) and
-        (999, 999). Per-dimension distinct touched chunks are {0, 999} on each
-        axis, so enumeration intersects the 2x2 = 4 combinations; the two
-        off-diagonal combinations are filtered out by `intersect`, leaving 2
-        surviving chunks. The key guarantee is that the work is bounded by the
-        per-dimension distinct touched chunks (4), not the dense 1e6 grid.
+        (999, 999). Correlated coordinate arrays are grouped *jointly*, so
+        enumeration intersects exactly the 2 touched chunks — never the 2x2
+        cartesian product of per-dimension distinct chunks, and never the dense
+        1e6 grid.
         """
         grid = ChunkGrid(
             dimensions=(
@@ -356,8 +355,34 @@ def test_2d_correlated_vindex_enumerates_per_dim_distinct_chunks(
 
         coords = sorted(r[0] for r in results)
         assert coords == [(0, 0), (999, 999)]
-        # 2x2 per-dim-distinct combinations enumerated; 2 survive intersection.
-        assert calls["n"] == 4
+        assert calls["n"] == 2
+
+    def test_2d_correlated_vindex_diagonal_is_linear_in_points(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """A diagonal of P correlated points touches P chunks with O(P) intersections.
+
+        Enumerating the cartesian product of per-dimension distinct chunk sets
+        would cost P**2 intersections (2500 here) — quadratic in the number of
+        selected points for the scattered selections of zarr-python gh-4174.
+        Joint grouping keeps resolution work proportional to the touched chunks.
+        """
+        p = 50
+        grid = ChunkGrid(
+            dimensions=(
+                FixedDimension(size=4, extent=4000),
+                FixedDimension(size=4, extent=4000),
+            )
+        )
+        # point i lands in chunk (2i, 2i): all per-dimension chunks distinct
+        coords_1d = np.arange(p, dtype=np.intp) * 8
+        t = IndexTransform.from_shape((4000, 4000)).vindex[coords_1d, coords_1d]
+
+        calls = _count_intersect_calls(monkeypatch)
+        results = list(iter_chunk_transforms(t, grid._dimensions))
+
+        assert sorted(r[0] for r in results) == [(2 * i, 2 * i) for i in range(p)]
+        assert calls["n"] == p
 
 
 class TestSubTransformToSelections: