fix: do not transpose the lazy array for the rgb variant (#289)

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

Author: dcherian

src/xpublish_tiles/pipeline.py

@@ -1168,7 +1168,9 @@ def apply_query(
                 if band_dim in array.coords
                 else None,
             )
-            array = array.transpose(band_dim, ...)
+            # Do not transpose here: xarray turns a lazy transpose into a
+            # vectorized indexer over every element, and zarr then gathers
+            # pointwise. ``subset_to_bbox`` moves the band dim first after load.
             datatype = RGBData(band_dim=band_dim, valid_range=_rgb_valid_range(array))
             extra_dims.discard(band_dim)
         else:
@@ -1201,6 +1203,7 @@ async def subset_to_bbox(
     # returns. ``pending`` carries the per-var bookkeeping needed to build
     # the ``PopulatedRenderContext`` once data has loaded.
     plan_patches: list[Patch] = []
+    plan_datatypes: list[DataType] = []
     pending: list[tuple[str, GridSystem, DataType, list[Patch]]] = []
 
     for var_name, array in validated.items():
@@ -1325,11 +1328,15 @@ async def subset_to_bbox(
             for patch in patches
         )
         plan_patches.extend(patches)
+        plan_datatypes.extend([array.datatype] * len(patches))
         pending.append((var_name, grid, array.datatype, patches))
 
     loaded = await load_plans(plans) if plans else []
 
-    async def _post_load(patch: Patch, ld: xr.DataArray) -> None:
+    async def _post_load(patch: Patch, datatype: DataType, ld: xr.DataArray) -> None:
+        if isinstance(datatype, RGBData):
+            # Downstream expects the band dim to lead; cheap on loaded data.
+            ld = ld.transpose(datatype.band_dim, ...)
         if isinstance(patch.grid, Polar):
             ld = patch.grid.assign_index(ld)
         if patch.coarsen_factors:
@@ -1339,7 +1346,10 @@ async def _post_load(patch: Patch, ld: xr.DataArray) -> None:
         patch.da = ld
 
     await asyncio.gather(
-        *(_post_load(p, ld) for p, ld in zip(plan_patches, loaded, strict=True))
+        *(
+            _post_load(p, dt, ld)
+            for p, dt, ld in zip(plan_patches, plan_datatypes, loaded, strict=True)
+        )
     )
 
     for var_name, grid, datatype, patches in pending:

tests/test_pipeline.py

@@ -3,6 +3,7 @@
 
 import io
 from dataclasses import replace
+from typing import Any
 
 import cf_xarray  # noqa: F401 - Enable cf accessor
 import morecantile
@@ -28,11 +29,13 @@
     MissingParameterError,
     VariableNotFoundError,
     check_transparent_pixels,
+    max_render_shape,
 )
 from xpublish_tiles.pipeline import (
     apply_query,
     bbox_overlap,
     pipeline,
+    subset_to_bbox,
 )
 from xpublish_tiles.testing.datasets import (
     CUBED_SPHERE,
@@ -74,7 +77,14 @@
     WGS84_TMS,
     TileTestParam,
 )
-from xpublish_tiles.types import ImageFormat, OutputBBox, OutputCRS, QueryParams, RGBData
+from xpublish_tiles.types import (
+    ImageFormat,
+    OutputBBox,
+    OutputCRS,
+    PopulatedRenderContext,
+    QueryParams,
+    RGBData,
+)
 
 
 @st.composite
@@ -1250,13 +1260,50 @@ def test_apply_query_rgb_keeps_band_dim():
     array = validated["foo"]
     assert isinstance(array.datatype, RGBData)
     assert array.datatype.band_dim == "rgb"
-    assert array.da.dims[0] == "rgb"
+    assert "rgb" in array.da.dims
 
     # without the rgb variant the band dim is squeezed like any extra dim
     validated = apply_query(ds, variables=["foo"], selectors={"rgb": "red"})
     assert validated["foo"].da.dims == ("latitude", "longitude")
 
 
+@pytest.mark.asyncio
+async def test_rgb_lazy_array_stays_sliceable(tmp_path):
+    """The rgb variant must not transpose the lazy array: xarray rewrites a
+    lazy transpose as a vectorized indexer over every element, and zarr then
+    gathers pointwise with full-size int64 index arrays. The band dim is moved
+    first only once the subset is in memory."""
+    from xarray.core.indexing import LazilyIndexedArray, LazilyVectorizedIndexedArray
+
+    RGB.create().to_zarr(
+        tmp_path / "rgb.zarr", mode="w", zarr_format=3, consolidated=False
+    )
+    ds = xr.open_zarr(
+        tmp_path / "rgb.zarr", chunks=None, consolidated=False, zarr_format=3
+    )
+
+    validated = apply_query(ds, variables=["foo"], selectors={}, rgb=True)
+    data: Any = validated["foo"].da.variable._data
+    chain: list[Any] = [data]
+    while hasattr(data, "array"):
+        data = data.array
+        chain.append(data)
+    assert any(isinstance(a, LazilyIndexedArray) for a in chain), chain
+    assert not any(isinstance(a, LazilyVectorizedIndexedArray) for a in chain), chain
+
+    query = create_query_params(Tile(x=0, y=0, z=0), WEBMERC_TMS, variant="rgb")
+    contexts = await subset_to_bbox(
+        validated,
+        bbox=query.bbox,
+        crs=query.crs,
+        max_shape=max_render_shape(style="raster", width=256, height=256),
+    )
+    context = contexts["foo"]
+    assert isinstance(context, PopulatedRenderContext)
+    (patch,) = context.patches
+    assert patch.da.dims[0] == "rgb"
+
+
 @pytest.mark.asyncio
 async def test_rgb_band_selector_renders_single_band(png_snapshot):
     """Colormap variants pick one band with a selector."""