perf(zarr-transforms): partition sorted 1-D array maps by chunk

packages/zarr-transforms/src/zarr_transforms/chunk_resolution.py

@@ -21,6 +21,10 @@
 4. **Yield** — produce `(chunk_coords, local_transform, surviving_indices)`
    triples that the codec pipeline consumes.
 
+Sorted one-dimensional correlated array maps can be partitioned directly
+because every touched chunk owns a contiguous slice of the index array. That
+case bypasses candidate enumeration and repeated intersection.
+
 `sub_transform_to_selections` bridges from the transform representation
 back to the raw `(chunk_selection, out_selection, drop_axes)` tuples that
 the current codec pipeline expects. This bridge will go away when the codec
@@ -40,7 +44,7 @@
 if TYPE_CHECKING:
     from collections.abc import Iterator, Sequence
 
-    from zarr_transforms.grid import ChunkGridLike
+    from zarr_transforms.grid import ChunkGridLike, DimensionGridLike
 
 OutIndices = (
     dict[int, np.ndarray[Any, np.dtype[np.intp]]] | np.ndarray[Any, np.dtype[np.intp]] | None
@@ -53,6 +57,61 @@
 ]
 
 
+def _one_dimensional_correlated_array_map(
+    transform: IndexTransform,
+) -> tuple[ArrayMap, np.ndarray[Any, np.dtype[np.intp]]] | None:
+    """Return a nonempty correlated 1-D ArrayMap and its storage coordinates.
+
+    A one-dimensional array selection has no cross-dimensional correlation to
+    preserve. The computed storage coordinates are also reused by general
+    resolution when they are unsorted.
+    """
+    if transform.input_rank != 1 or transform.output_rank != 1:
+        return None
+
+    m = transform.output[0]
+    if (
+        not isinstance(m, ArrayMap)
+        or m.input_dimension is not None
+        or m.index_array.ndim != 1
+        or m.index_array.size == 0
+    ):
+        return None
+
+    return m, m.offset + m.stride * m.index_array
+
+
+def _iter_sorted_1d_array_map(
+    m: ArrayMap,
+    storage: np.ndarray[Any, np.dtype[np.intp]],
+    dim_grid: DimensionGridLike,
+) -> Iterator[ChunkTransformResult]:
+    """Resolve a sorted 1-D ArrayMap one touched chunk at a time."""
+    start = 0
+    while start < storage.size:
+        chunk = dim_grid.index_to_chunk(int(storage[start]))
+        chunk_start = dim_grid.chunk_offset(chunk)
+        chunk_stop = chunk_start + dim_grid.chunk_size(chunk)
+        stop = int(np.searchsorted(storage, chunk_stop, side="left"))
+
+        restricted = IndexTransform(
+            domain=IndexDomain(inclusive_min=(0,), exclusive_max=(stop - start,)),
+            output=(
+                ArrayMap(
+                    index_array=m.index_array[start:stop],
+                    offset=m.offset,
+                    stride=m.stride,
+                    input_dimension=m.input_dimension,
+                ),
+            ),
+        )
+        local = restricted.translate((-chunk_start,))
+        surviving = np.arange(start, stop, dtype=np.intp)
+
+        yield (chunk,), local, surviving
+        start = stop
+
+
 def iter_chunk_transforms(
     transform: IndexTransform,
     chunk_grid: ChunkGridLike,
@@ -68,6 +127,16 @@ def iter_chunk_transforms(
     """
     dim_grids = chunk_grid._dimensions
 
+    array_map_1d = _one_dimensional_correlated_array_map(transform)
+    if array_map_1d is not None:
+        m, storage = array_map_1d
+        if storage[0] <= storage[-1] and bool(np.all(storage[1:] >= storage[:-1])):
+            dim_grid = dim_grids[0]
+            first_chunk = dim_grid.index_to_chunk(int(storage[0]))
+            if dim_grid.chunk_size(first_chunk) > 0:
+                yield from _iter_sorted_1d_array_map(m, 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
@@ -113,7 +182,10 @@ def iter_chunk_transforms(
             last = dg.index_to_chunk(s_max)
             chunk_candidates.append(range(first, last + 1))
         elif isinstance(m, ArrayMap):
-            storage = m.offset + m.stride * m.index_array
+            # already computed these storage coordinates 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:
                 # Empty fancy selection: no coordinates, so no chunks are touched.

packages/zarr-transforms/tests/test_chunk_resolution.py

@@ -3,8 +3,9 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
-from zarr.core.chunk_grids import ChunkGrid, FixedDimension
+from zarr.core.chunk_grids import ChunkGrid, FixedDimension, VaryingDimension
 
+import zarr_transforms.chunk_resolution as chunk_resolution
 from zarr_transforms.chunk_resolution import iter_chunk_transforms, sub_transform_to_selections
 from zarr_transforms.domain import IndexDomain
 from zarr_transforms.output_map import ArrayMap, ConstantMap, DimensionMap
@@ -110,6 +111,150 @@ def test_array_index(self) -> None:
         assert (2,) in coords_list
 
 
+class TestChunkResolutionSorted1D:
+    def test_matches_general_resolution_for_randomized_sorted_selections(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Direct partitioning matches the original resolver across varied inputs."""
+        rng = np.random.default_rng(0)
+        grids = (
+            ChunkGrid(dimensions=(FixedDimension(size=7, extent=30),)),
+            ChunkGrid(dimensions=(VaryingDimension(edges=(3, 4, 8, 5, 10), extent=30),)),
+        )
+
+        for grid in grids:
+            for _ in range(50):
+                idx = np.sort(rng.integers(0, 30, size=int(rng.integers(1, 80)))).astype(np.intp)
+                transform = IndexTransform.from_shape((30,)).vindex[idx]
+                direct = list(iter_chunk_transforms(transform, grid))
+
+                with monkeypatch.context() as context:
+                    context.setattr(
+                        chunk_resolution,
+                        "_one_dimensional_correlated_array_map",
+                        lambda _transform: None,
+                    )
+                    general = list(iter_chunk_transforms(transform, grid))
+
+                assert [result[0] for result in direct] == [result[0] for result in general]
+                for direct_result, general_result in zip(direct, general, strict=True):
+                    _, direct_t, direct_out = direct_result
+                    _, general_t, general_out = general_result
+                    assert direct_t.domain == general_t.domain
+
+                    direct_chunk_sel, direct_out_sel, direct_drop = sub_transform_to_selections(
+                        direct_t, direct_out
+                    )
+                    general_chunk_sel, general_out_sel, general_drop = sub_transform_to_selections(
+                        general_t, general_out
+                    )
+                    assert direct_drop == general_drop
+                    np.testing.assert_array_equal(direct_chunk_sel[0], general_chunk_sel[0])
+                    np.testing.assert_array_equal(direct_out_sel[0], general_out_sel[0])
+
+    def test_sorted_vindex_partitions_chunks_without_intersection(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Sorted vectorized coordinates are sliced directly per touched chunk."""
+        idx = np.array([0, 3, 4, 4, 9, 11], dtype=np.intp)
+        t = IndexTransform.from_shape((12,)).vindex[idx]
+        grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=12),))
+
+        calls = _count_intersect_calls(monkeypatch)
+        results = list(iter_chunk_transforms(t, grid))
+
+        assert [result[0] for result in results] == [(0,), (1,), (2,)]
+        assert calls["n"] == 0
+
+        expected_chunk_indices = ([0, 3], [0, 0], [1, 3])
+        expected_out_indices = ([0, 1], [2, 3], [4, 5])
+        for result, expected_chunk, expected_out in zip(
+            results, expected_chunk_indices, expected_out_indices, strict=True
+        ):
+            _, sub_t, out_indices = result
+            chunk_sel, out_sel, drop_axes = sub_transform_to_selections(sub_t, out_indices)
+            np.testing.assert_array_equal(chunk_sel[0], expected_chunk)
+            np.testing.assert_array_equal(out_sel[0], expected_out)
+            assert drop_axes == ()
+
+    def test_sorted_array_map_preserves_offset_and_stride(self) -> None:
+        """Storage partitioning retains the ArrayMap's offset and stride."""
+        t = IndexTransform(
+            domain=IndexDomain.from_shape((3,)),
+            output=(
+                ArrayMap(
+                    index_array=np.array([0, 1, 2], dtype=np.intp),
+                    offset=1,
+                    stride=3,
+                ),
+            ),
+        )
+        grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=8),))
+
+        results = list(iter_chunk_transforms(t, grid))
+
+        assert [result[0] for result in results] == [(0,), (1,)]
+        expected_chunk_indices = ([1], [0, 3])
+        expected_out_indices = ([0], [1, 2])
+        for result, expected_chunk, expected_out in zip(
+            results, expected_chunk_indices, expected_out_indices, strict=True
+        ):
+            _, sub_t, out_indices = result
+            chunk_sel, out_sel, _ = sub_transform_to_selections(sub_t, out_indices)
+            np.testing.assert_array_equal(chunk_sel[0], expected_chunk)
+            np.testing.assert_array_equal(out_sel[0], expected_out)
+
+    def test_sorted_vindex_with_varying_chunks(self) -> None:
+        """Touched-boundary searches also support a non-uniform 1-D grid."""
+        idx = np.array([0, 1, 2, 3, 5, 9], dtype=np.intp)
+        t = IndexTransform.from_shape((10,)).vindex[idx]
+        grid = ChunkGrid(dimensions=(VaryingDimension(edges=(2, 3, 5), extent=10),))
+
+        results = list(iter_chunk_transforms(t, grid))
+
+        assert [result[0] for result in results] == [(0,), (1,), (2,)]
+        expected_chunk_indices = ([0, 1], [0, 1], [0, 4])
+        for result, expected_chunk in zip(results, expected_chunk_indices, strict=True):
+            _, sub_t, out_indices = result
+            chunk_sel, _, _ = sub_transform_to_selections(sub_t, out_indices)
+            np.testing.assert_array_equal(chunk_sel[0], expected_chunk)
+
+    def test_sorted_vindex_with_zero_sized_dimension_uses_general_resolution(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """A zero-sized grid cannot be partitioned by touched boundaries."""
+        t = IndexTransform.from_shape((10,)).vindex[np.array([1], dtype=np.intp)]
+        grid = ChunkGrid(dimensions=(FixedDimension(size=0, extent=10),))
+
+        calls = _count_intersect_calls(monkeypatch)
+        results = list(iter_chunk_transforms(t, grid))
+
+        assert results == []
+        assert calls["n"] == 1
+
+    def test_unsorted_vindex_uses_general_resolution(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        """Unsorted coordinates continue through the general intersection logic."""
+        t = IndexTransform.from_shape((12,)).vindex[np.array([9, 0, 4], dtype=np.intp)]
+        grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=12),))
+
+        calls = _count_intersect_calls(monkeypatch)
+        results = list(iter_chunk_transforms(t, grid))
+
+        assert [result[0] for result in results] == [(0,), (1,), (2,)]
+        assert calls["n"] == 3
+
+    def test_sorted_oindex_uses_general_resolution(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        """Orthogonal ArrayMaps retain their existing domain-aware resolution."""
+        t = IndexTransform.from_shape((12,)).oindex[np.array([0, 4, 9], dtype=np.intp)]
+        grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=12),))
+
+        calls = _count_intersect_calls(monkeypatch)
+        results = list(iter_chunk_transforms(t, grid))
+
+        assert [result[0] for result in results] == [(0,), (1,), (2,)]
+        assert calls["n"] == 3
+
+
 def _count_intersect_calls(monkeypatch: pytest.MonkeyPatch) -> dict[str, int]:
     """Wrap `IndexTransform.intersect` with a call counter.
 
@@ -154,8 +299,9 @@ def test_1d_sparse_vindex_enumerates_only_touched_chunks(
 
         coords = sorted(r[0] for r in results)
         assert coords == [(0,), (999,)]
-        # Exactly the touched chunks, independent of the 1000-chunk grid size.
-        assert calls["n"] == 2
+        # Sorted 1-D coordinates are partitioned directly, without intersecting
+        # either the touched chunks or the 998 empty chunks between them.
+        assert calls["n"] == 0
 
     def test_2d_orthogonal_enumerates_only_touched_chunks(
         self, monkeypatch: pytest.MonkeyPatch