Use tree-reduction in mosaic (#126)

This should hopefully result in lower cluster memory usage, since intermediate results can get processed sooner (there's more parallelism, so less vulnerable to root task overproduction)

conftest.py

@@ -0,0 +1,12 @@
+import os
+from hypothesis import settings
+import pytest
+
+pytest.register_assert_rewrite(
+    "dask.array.utils", "dask.dataframe.utils", "dask.bag.utils"
+)
+
+settings.register_profile("default", settings.default, print_blob=True)
+settings.register_profile("ci", max_examples=1000, derandomize=True)
+
+settings.load_profile(os.getenv("HYPOTHESIS_PROFILE", "default"))

stackstac/ops.py

@@ -1,29 +1,166 @@
-from typing import Hashable, Sequence, Union
+from __future__ import annotations
+
+from functools import partial
+from typing import Hashable, Tuple, List, Union, Optional
+
 import numpy as np
+import dask.array as da
 import xarray as xr
 
 
-def _mosaic(arr, axis, reverse: bool = False, nodata: Union[int, float] = np.nan):
-    ax_length = arr.shape[axis]
+def _mosaic_base(
+    arr: np.ndarray,
+    axis: int,
+    *,
+    reverse: bool,
+    nodata: Union[int, float],
+    keepdims: bool = False,
+    initial: Optional[np.ndarray] = None,
+) -> Tuple[np.ndarray, bool]:
+    """
+    Internal function to mosaic a NumPy array along an axis. Does some extra funky stuff to be useful with dask.
+
+    Parameters
+    ----------
+    arr: array to mosaic, must be at least 1D
+    axis: axis to reduce
+    reverse: False = last on top, True = first on top
+    nodata: the value to treat as invalid
+    keepdims: whether to preserve the reduced axis as 1-length
+    initial:
+        Array to use as the initial layer, instead of the first index into ``arr``.
+        Used when chaning mosaics together.
+
+    Returns
+    -------
+    result, done
+        The mosaicked array, and a bool of whether mosaicing is complete (every pixel in ``result`` is valid)
+    """
+    if isinstance(axis, tuple):
+        # `da.reduction` likes to turn int axes into tuples
+        assert len(axis) == 1, f"Multiple axes to mosaic not supported: {axis=}"
+        axis = axis[0]
+
+    ax_length = arr.shape[axis] if np.ndim(arr) else 1
 
     # "normal" means last -> first, "reversed" means first -> last,
     # so we apply `reversed` in the opposite case you'd expect
     indices = iter(range(ax_length) if reverse else reversed(range(ax_length)))
-    out = np.take(arr, next(indices), axis=axis)
 
+    if initial is None:
+        i = next(indices)
+        out = np.take(arr, [i] if keepdims else i, axis=axis)
+    else:
+        out = initial
+
+    # TODO implement this with Numba, would be much faster
+    done = False
     for i in indices:
-        layer = np.take(arr, i, axis=axis)
         where = np.isnan(out) if np.isnan(nodata) else out == nodata
+        if not where.any():
+            # Short-circuit when mosaic is complete (no invalid pixels left)
+            done = True
+            break
+        layer = np.take(arr, [i] if keepdims else i, axis=axis)
         out = np.where(where, layer, out)
+
+    # Final `where` may have completed the mosaic, but we wouldn't have checked
+    if not done:
+        where = np.isnan(out) if np.isnan(nodata) else out == nodata
+        done = not where.any()
+
+    return out, done
+
+
+def _mosaic_np(
+    arr: np.ndarray,
+    axis: int,
+    *,
+    reverse: bool,
+    nodata: Union[int, float],
+    keepdims: bool = False,
+) -> np.ndarray:
+    "Mosaic a NumPy array, without returning ``done``"
+    return _mosaic_base(arr, axis, reverse=reverse, nodata=nodata, keepdims=keepdims)[0]
+
+
+def _mosaic_dask_aggregate(
+    arrs: Union[List[np.ndarray], np.ndarray],
+    axis: Union[int, Tuple[int]],
+    keepdims: bool,
+    *,
+    reverse: bool,
+    nodata: Union[int, float],
+) -> np.ndarray | np.ndarray:
+    """
+    Mosaic a list of NumPy arrays (or a single one) without concatenating them all first.
+
+    Avoiding concatenation lets us avoid unnecessary copies and memory spikes.
+    """
+    if not isinstance(arrs, list):
+        arrs = [arrs]
+
+    if isinstance(axis, tuple):
+        # `da.reduction` likes to turn int axes into tuples
+        assert len(axis) == 1, f"Multiple axes to mosaic not supported: {axis=}"
+        axis = axis[0]
+
+    out: Optional[np.ndarray] = None
+    for arr in arrs if reverse else arrs[::-1]:
+        # ^ Remember, `reverse` is backwards of what you'd think
+        if out is None and (arr.ndim == 0 or arr.shape[axis] <= 1):
+            # There's nothing to mosaic in this chunk (it's effectively 1-length along `axis` already).
+            # Skip mosaicing, just drop the axis if we're supposed to.
+            out = np.take(arr, 0, axis=axis) if not keepdims else arr
+        else:
+            # Multiple entries along the axis: mosaic them, using any
+            # mosaic we've already built up as a starting point.
+            out, done = _mosaic_base(
+                arr,
+                axis,
+                reverse=reverse,
+                nodata=nodata,
+                keepdims=keepdims,
+                initial=out,
+            )
+            if done:
+                break
+
+    assert out is not None, "Cannot mosaic zero arrays"
     return out
 
 
+def _mosaic_dask(
+    arr: da.Array,
+    axis: int,
+    *,
+    reverse: bool,
+    nodata: Union[int, float],
+    split_every: Union[None, int],
+) -> da.Array:
+    "Tree-reduction-based mosaic for Dask arrays."
+    return da.reduction(
+        arr,
+        partial(_mosaic_np, reverse=reverse, nodata=nodata),
+        partial(_mosaic_dask_aggregate, reverse=reverse, nodata=nodata),
+        axis=axis,
+        keepdims=False,
+        dtype=arr.dtype,
+        split_every=split_every,
+        name="mosaic",
+        meta=da.utils.meta_from_array(arr, ndim=arr.ndim - 1),
+        concatenate=False,
+    )
+
+
 def mosaic(
     arr: xr.DataArray,
-    dim: Union[None, Hashable, Sequence[Hashable]] = None,
-    axis: Union[None, int, Sequence[int]] = 0,
+    dim: Union[None, Hashable] = None,
+    axis: Union[None, int] = 0,
+    *,
     reverse: bool = False,
     nodata: Union[int, float] = np.nan,
+    split_every: Union[None, int] = None,
 ):
     """
     Flatten a dimension of a `~xarray.DataArray` by picking the first valid pixel.
@@ -49,6 +186,13 @@ def mosaic(
         is used when the array has an integer or boolean dtype. Since NaN
         cannot exist in those arrays, this indicates a different ``nodata``
         value needs to be used.
+    split_every:
+        For Dask arrays: how many *chunks* to mosaic together at once.
+
+        The Dask default is 4. A higher number will mean a smaller graph,
+        but higher peak memory usage.
+
+        Ignored for NumPy arrays.
 
     Returns
     -------
@@ -60,8 +204,14 @@ def mosaic(
         raise ValueError(
             f"Cannot use {nodata=} when mosaicing a {arr.dtype} array, since {nodata} cannot exist in the array."
         )
+
+    func = (
+        partial(_mosaic_dask, split_every=split_every)
+        if isinstance(arr.data, da.Array)
+        else _mosaic_np
+    )
     return arr.reduce(
-        _mosaic,
+        func,
         dim=dim,
         axis=axis,
         keep_attrs=True,

stackstac/tests/test_mosaic.py

@@ -0,0 +1,81 @@
+from typing import Tuple
+from hypothesis import given, assume, strategies as st
+from hypothesis.extra import numpy as st_np
+import pytest
+
+import xarray as xr
+import dask.array as da
+from dask.array.utils import assert_eq
+import numpy as np
+
+from stackstac.ops import mosaic
+from stackstac.testing import strategies as st_stc
+
+
+@pytest.mark.parametrize("dask", [False, True])
+def test_mosaic_basic(dask: bool):
+    arr = np.array(
+        [
+            [np.nan, 1, 2, np.nan],
+            [np.nan, 10, 20, 30],
+            [np.nan, 100, 200, np.nan],
+        ]
+    )
+    if dask:
+        arr = da.from_array(arr, chunks=1)
+
+    xarr = xr.DataArray(arr)
+
+    fwd = mosaic(xarr, axis=0)
+    # Remember, forward means "last item on top" (because last item == last date == newest)
+    assert_eq(fwd.data, np.array([np.nan, 100, 200, 30]), equal_nan=True)
+
+    rev = mosaic(xarr, axis=0, reverse=True)
+    assert_eq(rev.data, np.array([np.nan, 1, 2, 30]), equal_nan=True)
+
+
+@pytest.mark.parametrize("dtype", [np.dtype("bool"), np.dtype("int"), np.dtype("uint")])
+def test_mosaic_dtype_error(dtype: np.dtype):
+    arr = xr.DataArray(np.arange(3).astype(dtype))
+    with pytest.raises(ValueError, match="Cannot use"):
+        mosaic(arr)
+
+
+@given(
+    st.data(),
+    st_stc.raster_dtypes,
+    st_np.array_shapes(max_dims=4, max_side=5),
+    st.booleans(),
+)
+def test_fuzz_mosaic(
+    data: st.DataObject, dtype: np.dtype, shape: Tuple[int, ...], reverse: bool
+):
+    """
+    See if we can break mosaic.
+
+    Not testing correctness much here, since that's hard to do without rewriting a mosaic implementation.
+    """
+    # `np.where` doesn't seem to preserve endianness.
+    # Even with our best efforts to add an `astype` at the end, this will fail (on little-endian systems at least?)
+    # with big-endian dtypes.
+    assume(dtype.byteorder != ">")
+
+    fill_value = data.draw(st_np.from_dtype(dtype), label="fill_value")
+    arr = data.draw(st_np.arrays(dtype, shape, fill=st.just(fill_value)), label="arr")
+    chunkshape = data.draw(
+        st_np.array_shapes(
+            min_dims=arr.ndim, max_dims=arr.ndim, max_side=max(arr.shape)
+        ), label="chunkshape"
+    )
+    axis = data.draw(st.integers(-arr.ndim + 1, arr.ndim - 1), label="axis")
+
+    darr = da.from_array(arr, chunks=chunkshape)
+    split_every = data.draw(st.integers(1, darr.numblocks[axis]), label="split_every")
+    xarr = xr.DataArray(darr)
+
+    result = mosaic(
+        xarr, axis=axis, reverse=reverse, nodata=fill_value, split_every=split_every
+    )
+    assert result.dtype == arr.dtype
+    result_np = mosaic(xr.DataArray(arr), axis=axis, reverse=reverse, nodata=fill_value)
+    assert_eq(result.data, result_np.data, equal_nan=True)