From c0d9493d66c7dedf5bea9e76bc31198fd541589b Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Tue, 21 Jul 2026 22:43:05 +0200 Subject: [PATCH] add _iter_sorted_1d_array_map and add tests --- .../src/zarr_transforms/chunk_resolution.py | 76 ++++++++- .../tests/test_chunk_resolution.py | 152 +++++++++++++++++- 2 files changed, 223 insertions(+), 5 deletions(-) diff --git a/packages/zarr-transforms/src/zarr_transforms/chunk_resolution.py b/packages/zarr-transforms/src/zarr_transforms/chunk_resolution.py index 485341218a..edda3283ed 100644 --- a/packages/zarr-transforms/src/zarr_transforms/chunk_resolution.py +++ b/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. diff --git a/packages/zarr-transforms/tests/test_chunk_resolution.py b/packages/zarr-transforms/tests/test_chunk_resolution.py index 15ac8496e9..fc3205dbe9 100644 --- a/packages/zarr-transforms/tests/test_chunk_resolution.py +++ b/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