Close some long standing issues

.github/workflows/main.yml

@@ -288,7 +288,7 @@ jobs:
         if: |
           github.repository == 'opendatacube/odc-stac'
 
-        uses: codecov/codecov-action@v1
+        uses: codecov/codecov-action@v3
         with:
           fail_ci_if_error: false
           verbose: false
@@ -426,7 +426,7 @@ jobs:
       - name: Deploy to Netlify
         id: netlify
         if: github.event_name == 'pull_request'
-        uses: nwtgck/actions-netlify@v1.0
+        uses: nwtgck/actions-netlify@v2
         with:
           production-branch: "main"
           publish-dir: "docs/_build/html"
@@ -440,7 +440,7 @@ jobs:
           NETLIFY_SITE_ID: ${{ secrets.NETLIFY_SITE_ID }}
 
       - name: Print Notice
-        uses: actions/github-script@v5
+        uses: actions/github-script@v6
         if: github.event_name == 'pull_request'
         env:
           NETLIFY_URL: ${{ steps.netlify.outputs.deploy-url }}

binder/environment.yml

@@ -22,6 +22,7 @@ dependencies:
   # JupyterLab
   - jupytext
   - jupyter-server-proxy
+  - ipykernel
   - matplotlib
   - ipympl
   - dask

odc/stac/_load.py

@@ -1,4 +1,6 @@
 """stac.load - dc.load from STAC Items."""
+from __future__ import annotations
+
 import dataclasses
 import functools
 import itertools
@@ -12,6 +14,7 @@
     Iterable,
     Iterator,
     List,
+    Literal,
     Optional,
     Protocol,
     Sequence,
@@ -25,6 +28,7 @@
 import pystac.item
 import xarray as xr
 from dask import array as da
+from dask.array.core import normalize_chunks
 from dask.base import quote, tokenize
 from dask.utils import ndeepmap
 from numpy.typing import DTypeLike
@@ -85,7 +89,7 @@ def dst_roi(self):
     @property
     def dst_gbox(self) -> GeoBox:
         _, y, x = self.idx_tyx
-        return self.gbt[y, x]
+        return cast(GeoBox, self.gbt[y, x])
 
 
 class _DaskGraphBuilder:
@@ -98,13 +102,15 @@ def __init__(
         tyx_bins: Dict[Tuple[int, int, int], List[int]],
         gbt: GeoboxTiles,
         env: Dict[str, Any],
+        time_chunks: int = 1,
     ) -> None:
         self.cfg = cfg
         self.items = items
         self.tyx_bins = tyx_bins
         self.gbt = gbt
         self.env = env
-        self._tk = tokenize(items, cfg, gbt, tyx_bins, env)
+        self._tk = tokenize(items, cfg, gbt, tyx_bins, env, time_chunks)
+        self.chunk_shape = (time_chunks, *self.gbt.chunk_shape((0, 0)).yx)
 
     def __call__(
         self,
@@ -119,6 +125,11 @@ def __call__(
         cfg = self.cfg[name]
         assert dtype == cfg.dtype
 
+        chunks = unpack_chunks(self.chunk_shape, shape)
+        tchunk_range = [
+            range(last - n, last) for last, n in zip(np.cumsum(chunks[0]), chunks[0])
+        ]
+
         cfg_key = f"cfg-{tokenize(cfg)}"
         gbt_key = f"grid-{tokenize(self.gbt)}"
 
@@ -129,19 +140,23 @@ def __call__(
         tk = self._tk
         band_key = f"{name}-{tk}"
         md_key = f"md-{name}-{tk}"
-        shape_in_blocks = (shape[0], *self.gbt.shape.yx)
+        shape_in_blocks = tuple(len(ch) for ch in chunks)
         for idx, item in enumerate(self.items):
             band = item.get(name, None)
             if band is not None:
                 dsk[md_key, idx] = band
 
         for ti, yi, xi in np.ndindex(shape_in_blocks):
-            tyx_idx = (ti, yi, xi)
-            srcs = [
-                (md_key, idx)
-                for idx in self.tyx_bins.get(tyx_idx, [])
-                if (md_key, idx) in dsk
-            ]
+            srcs = []
+            for _ti in tchunk_range[ti]:
+                srcs.append(
+                    [
+                        (md_key, idx)
+                        for idx in self.tyx_bins.get((_ti, yi, xi), [])
+                        if (md_key, idx) in dsk
+                    ]
+                )
+
             dsk[band_key, ti, yi, xi] = (
                 _dask_loader_tyx,
                 srcs,
@@ -151,9 +166,6 @@ def __call__(
                 self.env,
             )
 
-        chunk_shape = (1, *self.gbt.chunk_shape((0, 0)).yx)
-        chunks = unpack_chunks(chunk_shape, shape)
-
         return da.Array(dsk, band_key, chunks, dtype=dtype, shape=shape)
 
 
@@ -214,7 +226,7 @@ def load(
     groupby: Optional[Groupby] = "time",
     resampling: Optional[Union[str, Dict[str, str]]] = None,
     dtype: Union[DTypeLike, Dict[str, DTypeLike], None] = None,
-    chunks: Optional[Dict[str, int]] = None,
+    chunks: Optional[Dict[str, int | Literal["auto"]]] = None,
     pool: Union[ThreadPoolExecutor, int, None] = None,
     # Geo selection
     crs: MaybeCRS = Unset(),
@@ -493,14 +505,6 @@ def load(
     if gbox is None:
         raise ValueError("Failed to auto-guess CRS/resolution.")
 
-    if chunks is not None:
-        chunk_shape = _resolve_chunk_shape(gbox, chunks)
-    else:
-        chunk_shape = _resolve_chunk_shape(
-            gbox,
-            {dim: DEFAULT_CHUNK_FOR_LOAD for dim in gbox.dimensions},
-        )
-
     debug = kw.get("debug", False)
 
     # Check we have all the bands of interest
@@ -518,6 +522,14 @@ def load(
         fail_on_error=fail_on_error,
     )
 
+    if dtype is None:
+        _dtypes = sorted(
+            set(cfg.dtype for cfg in load_cfg.values() if cfg.dtype is not None),
+            key=lambda x: np.dtype(x).itemsize,
+            reverse=True,
+        )
+        dtype = "uint16" if len(_dtypes) == 0 else _dtypes[0]
+
     if patch_url is not None:
         _parsed = [patch_urls(item, edit=patch_url, bands=bands) for item in _parsed]
 
@@ -533,9 +545,19 @@ def load(
 
     tss = _extract_timestamps(ndeepmap(2, lambda idx: _parsed[idx], _grouped_idx))
 
+    if chunks is not None:
+        chunk_shape = _resolve_chunk_shape(len(tss), gbox, chunks, dtype)
+    else:
+        chunk_shape = _resolve_chunk_shape(
+            len(tss),
+            gbox,
+            {dim: DEFAULT_CHUNK_FOR_LOAD for dim in gbox.dimensions},
+            dtype,
+        )
+
     # Spatio-temporal binning
     assert isinstance(gbox.crs, CRS)
-    gbt = GeoboxTiles(gbox, chunk_shape)
+    gbt = GeoboxTiles(gbox, chunk_shape[1:])
     tyx_bins = dict(_tyx_bins(_grouped_idx, _parsed, gbt))
     _parsed = [item.strip() for item in _parsed]
 
@@ -560,15 +582,6 @@ def _with_debug_info(ds: xr.Dataset, **kw) -> xr.Dataset:
         )
         return ds
 
-    def _task_stream(bands: List[str]) -> Iterator[_LoadChunkTask]:
-        _shape = (len(_grouped_idx), *gbt.shape)
-        for band_name in bands:
-            cfg = load_cfg[band_name]
-            for ti, yi, xi in np.ndindex(_shape):
-                tyx_idx = (ti, yi, xi)
-                srcs = [(idx, band_name) for idx in tyx_bins.get(tyx_idx, [])]
-                yield _LoadChunkTask(band_name, srcs, cfg, gbt, tyx_idx)
-
     _rio_env = _capture_rio_env()
     if chunks is not None:
         # Dask case: dummy for now
@@ -578,9 +591,19 @@ def _task_stream(bands: List[str]) -> Iterator[_LoadChunkTask]:
             tyx_bins,
             gbt,
             _rio_env,
+            time_chunks=chunk_shape[0],
         )
         return _with_debug_info(_mk_dataset(gbox, tss, load_cfg, _loader))
 
+    def _task_stream(bands: List[str]) -> Iterator[_LoadChunkTask]:
+        _shape = (len(_grouped_idx), *gbt.shape)
+        for band_name in bands:
+            cfg = load_cfg[band_name]
+            for ti, yi, xi in np.ndindex(_shape):
+                tyx_idx = (ti, yi, xi)
+                srcs = [(idx, band_name) for idx in tyx_bins.get(tyx_idx, [])]
+                yield _LoadChunkTask(band_name, srcs, cfg, gbt, tyx_idx)
+
     ds = _mk_dataset(gbox, tss, load_cfg)
     ny, nx = gbt.shape.yx
     total_tasks = len(bands) * len(tss) * ny * nx
@@ -658,17 +681,19 @@ def _resolve(name: str, band: RasterBandMetadata) -> RasterLoadParams:
 
 
 def _dask_loader_tyx(
-    srcs: List[RasterSource],
+    srcs: List[List[RasterSource]],
     gbt: GeoboxTiles,
     iyx: Tuple[int, int],
     cfg: RasterLoadParams,
     env: Dict[str, Any],
 ):
     assert cfg.dtype is not None
-    gbox = gbt[iyx]
-    chunk = np.empty(gbox.shape.yx, dtype=cfg.dtype)
+    gbox = cast(GeoBox, gbt[iyx])
+    chunk = np.empty((len(srcs), *gbox.shape.yx), dtype=cfg.dtype)
     with rio_env(**env):
-        return _fill_2d_slice(srcs, gbox, cfg, chunk)[np.newaxis]
+        for i, plane in enumerate(srcs):
+            _fill_2d_slice(plane, gbox, cfg, chunk[i, :, :])
+        return chunk
 
 
 def _fill_2d_slice(
@@ -852,14 +877,17 @@ def _tyx_bins(
         yield from (((t_idx, *idx), ii_item) for idx, ii_item in _yx.items())
 
 
-def _resolve_chunk_shape(gbox: GeoBox, chunks: Dict[str, int]) -> Tuple[int, int]:
-    def _norm_dim(chunk: int, sz: int) -> int:
-        if chunk < 0 or chunk > sz:
-            return sz
-        return chunk
+def _resolve_chunk_shape(
+    nt: int, gbox: GeoBox, chunks: Dict[str, int | Literal["auto"]], dtype: Any
+) -> Tuple[int, int, int]:
+    tt = chunks.get("time", 1)
+    ty, tx = (
+        chunks.get(dim, chunks.get(fallback_dim, -1))
+        for dim, fallback_dim in zip(gbox.dimensions, ["y", "x"])
+    )
+    nt, ny, nx = (
+        ch[0]
+        for ch in normalize_chunks((tt, ty, tx), (nt, *gbox.shape.yx), dtype=dtype)
+    )
 
-    ny, nx = [
-        _norm_dim(chunks.get(dim, chunks.get(fallback_dim, -1)), n)
-        for dim, fallback_dim, n in zip(gbox.dimensions, ["y", "x"], gbox.shape.yx)
-    ]
-    return ny, nx
+    return nt, ny, nx