Feat/fast subslice reads

bioio_base/reader.py

@@ -727,6 +727,32 @@ def get_image_dask_data(
             **kwargs,
         )
 
+    def _read_indexed(
+        self, given_dims: str, dim_specs: List[types.DimSpec]
+    ) -> np.ndarray:
+        """
+        Return the native-order array with ``dim_specs`` applied.
+
+        The default behavior here is to call ``self.data`` which can cause a
+        delayed-load of the whole image. Readers should override this function
+        in order to do more optimal finegrained data reads based on the
+        ``dim_specs``.
+
+        Parameters
+        ----------
+        given_dims: str
+            The native dimension ordering of the image (``self.dims.order``).
+        dim_specs: List[types.DimSpec]
+            One getitem operation per dimension in ``given_dims``, as produced by
+            ``transforms.compute_dim_specs``.
+
+        Returns
+        -------
+        data: np.ndarray
+            The indexed image data in native (reduced) dimension order.
+        """
+        return self.data[tuple(dim_specs)]
+
     def get_image_data(
         self, dimension_order_out: Optional[str] = None, **kwargs: Any
     ) -> np.ndarray:
@@ -789,21 +815,27 @@ def get_image_data(
         -----
         * If a requested dimension is not present in the data the dimension is
           added with a depth of 1.
-        * This will preload the entire image before returning the requested data.
+        * Only the requested sub-region is read when the reader supports it (see
+          ``_read_indexed``); otherwise the full image is read then sliced.
 
         See `aicsimageio.transforms.reshape_data` for more details.
         """
-        # If no out orientation, simply return current data as dask array
+        # If no out orientation, simply return current data
         if dimension_order_out is None:
             return self.data
 
-        # Transform and return
-        return transforms.reshape_data(
-            data=self.data,
-            given_dims=self.dims.order,
-            return_dims=dimension_order_out,
+        # Compute the indexer (reads no pixels), let the reader materialize only
+        # the requested sub-region, then pad/transpose to the requested order.
+        dim_specs, new_dims = transforms.compute_dim_specs(
+            self.shape,
+            self.dims.order,
+            dimension_order_out,
             **kwargs,
         )
+        indexed = self._read_indexed(self.dims.order, dim_specs)
+        return transforms.finalize_dims(
+            indexed, new_dims, self.dims.order, dimension_order_out
+        )
 
     @property
     def metadata(self) -> Any:

bioio_base/tests/test_reader.py

@@ -0,0 +1,39 @@
+from typing import Any, Dict
+
+import numpy as np
+import pytest
+
+from bioio_base import noop_reader, transforms
+
+
+# NoopReader exposes mock data with dims TCZYX (shape 4, 5, 6, 7, 8).
+@pytest.mark.parametrize(
+    "order, kwargs",
+    [
+        ("ZYX", {"T": 0, "C": 1}),
+        ("TCZYX", {"C": [0, 2], "Z": slice(0, 4, 2)}),
+    ],
+)
+def test_get_image_data_matches_full_read_then_slice(
+    order: str, kwargs: Dict[str, Any]
+) -> None:
+    reader = noop_reader.NoopReader("anything")
+    expected = transforms.reshape_data(reader.data, reader.dims.order, order, **kwargs)
+    actual = reader.get_image_data(order, **kwargs)
+    np.testing.assert_array_equal(actual, expected)
+
+
+def test_get_image_data_no_order_returns_full_data() -> None:
+    reader = noop_reader.NoopReader("anything")
+    np.testing.assert_array_equal(reader.get_image_data(), reader.data)
+
+
+def test_read_indexed_default_matches_getitem() -> None:
+    reader = noop_reader.NoopReader("anything")
+    specs, _ = transforms.compute_dim_specs(
+        reader.shape, reader.dims.order, "CZYX", T=1
+    )
+    np.testing.assert_array_equal(
+        reader._read_indexed(reader.dims.order, specs),
+        reader.data[tuple(specs)],
+    )

bioio_base/tests/test_transforms.py

@@ -520,3 +520,70 @@ def test_convert_list_to_slice(
     list_to_test: Union[List, Tuple], expected: Union[int, List, slice, Tuple]
 ) -> None:
     assert transforms.reduce_to_slice(list_to_test) == expected
+
+
+# slice(None, None, None) is the "select everything" spec for an untouched axis.
+_ALL = slice(None, None, None)
+
+
+@pytest.mark.parametrize(
+    "shape, given_dims, return_dims, kwargs, expected_specs, expected_new_dims",
+    [
+        # Fixed integer not in return_dims -> bare-int spec, axis dropped.
+        ((10, 100, 100), "ZYX", "YX", {"Z": 1}, [1, _ALL, _ALL], "YX"),
+        # Slice in return_dims is kept as-is, axis stays.
+        (
+            (10, 100, 100),
+            "ZYX",
+            "ZYX",
+            {"Z": slice(0, 4, 2)},
+            [slice(0, 4, 2), _ALL, _ALL],
+            "ZYX",
+        ),
+        # Evenly-spaced list collapses to an equivalent slice.
+        (
+            (6, 100, 100),
+            "CYX",
+            "CYX",
+            {"C": [0, 2, 4]},
+            [slice(0, 5, 2), _ALL, _ALL],
+            "CYX",
+        ),
+        # Unevenly-spaced list cannot reduce, so it is kept as a list.
+        ((6, 100, 100), "CYX", "CYX", {"C": [0, 3, 1]}, [[0, 3, 1], _ALL, _ALL], "CYX"),
+    ],
+    ids=[
+        "fixed_integer_drops_axis",
+        "slice_kept",
+        "list_reduced_to_slice",
+        "list_kept",
+    ],
+)
+def test_compute_dim_specs(
+    shape: Tuple[int, ...],
+    given_dims: str,
+    return_dims: str,
+    kwargs: Mapping[str, Any],
+    expected_specs: List[Any],
+    expected_new_dims: str,
+) -> None:
+    specs, new_dims = transforms.compute_dim_specs(
+        shape, given_dims, return_dims, **kwargs
+    )
+    assert specs == expected_specs
+    assert new_dims == expected_new_dims
+
+
+def test_finalize_dims_pads_and_transposes() -> None:
+    data = np.arange(12).reshape(3, 4)  # YX after indexing
+    out = transforms.finalize_dims(data, "YX", "YX", "TXY")
+    assert out.shape == (1, 4, 3)
+    np.testing.assert_array_equal(out[0], data.T)
+
+
+def test_reshape_data_still_matches_helpers() -> None:
+    data = np.random.rand(10, 100, 100)
+    direct = transforms.reshape_data(data, "ZYX", "YX", Z=3)
+    specs, new_dims = transforms.compute_dim_specs(data.shape, "ZYX", "YX", Z=3)
+    via_helpers = transforms.finalize_dims(data[tuple(specs)], new_dims, "ZYX", "YX")
+    np.testing.assert_array_equal(direct, via_helpers)