Adding DaskRasterReader protocol

Optional feature of the ReaderDriver.

odc/loader/_rio.py

@@ -33,7 +33,13 @@
     resolve_src_nodata,
     same_nodata,
 )
-from .types import MDParser, RasterLoadParams, RasterSource, ReaderSubsetSelection
+from .types import (
+    DaskRasterReader,
+    MDParser,
+    RasterLoadParams,
+    RasterSource,
+    ReaderSubsetSelection,
+)
 
 log = logging.getLogger(__name__)
 
@@ -143,7 +149,11 @@ def open(
         return RioReader(src, ctx)
 
     @property
-    def md_parser(self) -> Optional[MDParser]:
+    def md_parser(self) -> MDParser | None:
+        return None
+
+    @property
+    def dask_reader(self) -> DaskRasterReader | None:
         return None
 
 

odc/loader/testing/fixtures.py

@@ -22,6 +22,7 @@
 from .._reader import expand_selection
 from ..types import (
     BandKey,
+    DaskRasterReader,
     MDParser,
     RasterGroupMetadata,
     RasterLoadParams,
@@ -239,3 +240,7 @@ def open(self, src: RasterSource, ctx: FakeReader.LoadState) -> FakeReader:
     @property
     def md_parser(self) -> MDParser | None:
         return self._parser
+
+    @property
+    def dask_reader(self) -> DaskRasterReader | None:
+        return None

odc/loader/testing/mem_reader.py

@@ -14,6 +14,7 @@
 
 from ..types import (
     BandKey,
+    DaskRasterReader,
     FixedCoord,
     MDParser,
     RasterBandMetadata,
@@ -150,6 +151,10 @@ def open(self, src: RasterSource, ctx: Context) -> XrMemReader:
     def md_parser(self) -> MDParser:
         return XrMDPlugin(self.src)
 
+    @property
+    def dask_reader(self) -> DaskRasterReader | None:
+        return None
+
 
 def band_info(xx: xr.DataArray) -> RasterBandMetadata:
     """

odc/loader/types.py

@@ -445,6 +445,30 @@ def read(
     ) -> tuple[tuple[slice, slice], np.ndarray]: ...
 
 
+class DaskRasterReader(Protocol):
+    """
+    Protocol for raster readers that produce Dask sub-graphs.
+
+    ``.read`` method should return a Dask future evaluating to a numpy array of
+    pixels for a given geobox, alternatively dask future may evaluate to a
+    subset of the geobox overlapping with the source. In this case Dask future
+    should evaluate to a tuple: ``(yx_slice, pixels)``, such that
+    ``dst_geobox[yx_slice].shape == pixels.shape[ydim:ydim+2]``.
+    """
+
+    # pylint: disable=too-few-public-methods
+
+    def read(
+        self,
+        cfg: RasterLoadParams,
+        dst_geobox: GeoBox,
+        *,
+        selection: Optional[ReaderSubsetSelection] = None,
+    ) -> Any: ...
+
+    def open(self, src: RasterSource, ctx: Any) -> "DaskRasterReader": ...
+
+
 class ReaderDriver(Protocol):
     """
     Protocol for reader drivers.
@@ -470,6 +494,9 @@ def open(self, src: RasterSource, ctx: Any) -> RasterReader: ...
     @property
     def md_parser(self) -> MDParser | None: ...
 
+    @property
+    def dask_reader(self) -> DaskRasterReader | None: ...
+
 
 ReaderDriverSpec = Union[str, ReaderDriver]