Block-mapped resample with the help of flox

pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["flit_core >=3.9,<4"]
+requires = ["flit_core >=3.8,<4"]
 build-backend = "flit_core.buildapi"
 
 [project]
@@ -177,7 +177,7 @@ pep621_dev_dependency_groups = ["all", "dev", "docs"]
 "pyyaml" = "yaml"
 
 [tool.deptry.per_rule_ignores]
-DEP001 = ["SBCK"]
+DEP001 = ["SBCK", "flox"]
 DEP002 = ["bottleneck", "pyarrow"]
 DEP004 = ["matplotlib", "pytest_socket"]
 

xclim/core/indicator.py

@@ -162,6 +162,7 @@
     is_percentile_dataarray,
     load_module,
     raise_warn_or_log,
+    split_auxiliary_coordinates,
 )
 
 # Indicators registry
@@ -1450,12 +1451,11 @@ def _postprocess(self, outs, das, params):
             # Reduce by or and broadcast to ensure the same length in time
             # When indexing is used and there are no valid points in the last period, mask will not include it
             mask = reduce(np.logical_or, miss)
-            if (
-                isinstance(mask, DataArray)
-                and "time" in mask.dims
-                and mask.time.size < outs[0].time.size
-            ):
-                mask = mask.reindex(time=outs[0].time, fill_value=True)
+            if isinstance(mask, DataArray):  # mask might be a bool in some cases
+                if "time" in mask.dims and mask.time.size < outs[0].time.size:
+                    mask = mask.reindex(time=outs[0].time, fill_value=True)
+                # Remove any aux coord to avoid any unwanted dask computation in the alignment within "where"
+                mask, _ = split_auxiliary_coordinates(mask)
             outs = [out.where(~mask) for out in outs]
 
         return outs

xclim/core/options.py

@@ -25,6 +25,7 @@
 SDBA_ENCODE_CF = "sdba_encode_cf"
 KEEP_ATTRS = "keep_attrs"
 AS_DATASET = "as_dataset"
+MAP_BLOCKS = "resample_map_blocks"
 
 MISSING_METHODS: dict[str, Callable] = {}
 
@@ -39,6 +40,7 @@
     SDBA_ENCODE_CF: False,
     KEEP_ATTRS: "xarray",
     AS_DATASET: False,
+    MAP_BLOCKS: False,
 }
 
 _LOUDNESS_OPTIONS = frozenset(["log", "warn", "raise"])
@@ -71,6 +73,7 @@ def _valid_missing_options(mopts):
     SDBA_ENCODE_CF: lambda opt: isinstance(opt, bool),
     KEEP_ATTRS: _KEEP_ATTRS_OPTIONS.__contains__,
     AS_DATASET: lambda opt: isinstance(opt, bool),
+    MAP_BLOCKS: lambda opt: isinstance(opt, bool),
 }
 
 
@@ -185,6 +188,9 @@ class set_options:
         Note that xarray's "default" is equivalent to False. Default: ``"xarray"``.
     as_dataset : bool
         If True, indicators output datasets. If False, they output DataArrays. Default :``False``.
+    resample_map_blocks: bool
+        If True, some indicators will wrap their resampling operations with `xr.map_blocks`, using :py:func:`xclim.indices.utils.resample_map`.
+        This requires `flox` to be installed in order to ensure the chunking is appropriate.
 
     Examples
     --------

xclim/core/utils.py

@@ -862,3 +862,43 @@ def _chunk_like(*inputs: xr.DataArray | xr.Dataset, chunks: dict[str, int] | Non
                 da.chunk(**{d: c for d, c in chunks.items() if d in da.dims})
             )
     return tuple(outputs)
+
+
+def split_auxiliary_coordinates(
+    obj: xr.DataArray | xr.Dataset,
+) -> tuple[xr.DataArray | xr.Dataset, xr.Dataset]:
+    """Split auxiliary coords from the dataset.
+
+    An auxiliary coordinate is a coordinate variable that does not define a dimension and thus is not necessarily needed for dataset alignment.
+    Any coordinate that has a name different than its dimension(s) is flagged as auxiliary. All scalar coordinates are flagged as auxiliary.
+
+    Parameters
+    ----------
+    obj : DataArray or Dataset
+        Xarray object
+
+    Returns
+    -------
+    clean_obj :
+        Same as `obj` but without any auxiliary coordinate.
+    aux_coords : Dataset
+        The auxiliary coordinates as a dataset. Might be empty.
+
+    Note
+    ----
+    This is useful to circumvent xarray's alignment checks that will sometimes look the auxiliary coordinate's data, which can trigger
+    unwanted dask computations.
+
+    The auxiliary coordinates can be merged back with the dataset with
+    :py:meth:`xarray.Dataset.assign_coords` or :py:meth:`xarray.DataArray.assign_coords`.
+
+    >>> clean, aux = split_auxiliary_coordinates(ds)
+    >>> merged = clean.assign_coords(da.coords)
+    >>> merged.identical(ds)  # True
+    """
+    aux_crd_names = [
+        nm for nm, crd in obj.coords.items() if len(crd.dims) != 1 or crd.dims[0] != nm
+    ]
+    aux_crd_ds = obj.coords.to_dataset()[aux_crd_names]
+    clean_obj = obj.drop_vars(aux_crd_names)
+    return clean_obj, aux_crd_ds

xclim/indices/_threshold.py

@@ -30,6 +30,7 @@
     spell_length_statistics,
     threshold_count,
 )
+from .helpers import resample_map
 
 # Frequencies : YS: year start, QS-DEC: seasons starting in december, MS: month start
 # See http://pandas.pydata.org/pandas-docs/stable/timeseries.html#offset-aliases
@@ -1488,12 +1489,17 @@ def last_spring_frost(
     thresh = convert_units_to(thresh, tasmin)
     cond = compare(tasmin, op, thresh, constrain=("<", "<="))
 
-    out = cond.resample(time=freq).map(
+    out = resample_map(
+        cond,
+        "time",
+        freq,
         rl.last_run_before_date,
-        window=window,
-        date=before_date,
-        dim="time",
-        coord="dayofyear",
+        map_kwargs=dict(
+            window=window,
+            date=before_date,
+            dim="time",
+            coord="dayofyear",
+        ),
     )
     out.attrs.update(units="", is_dayofyear=np.int32(1), calendar=get_calendar(tasmin))
     return out
@@ -1659,11 +1665,12 @@ def first_snowfall(
     thresh = convert_units_to(thresh, prsn, context="hydro")
     cond = prsn >= thresh
 
-    out = cond.resample(time=freq).map(
+    out = resample_map(
+        cond,
+        "time",
+        freq,
         rl.first_run,
-        window=1,
-        dim="time",
-        coord="dayofyear",
+        map_kwargs=dict(window=1, dim="time", coord="dayofyear"),
     )
     out.attrs.update(units="", is_dayofyear=np.int32(1), calendar=get_calendar(prsn))
     return out
@@ -1714,11 +1721,12 @@ def last_snowfall(
     thresh = convert_units_to(thresh, prsn, context="hydro")
     cond = prsn >= thresh
 
-    out = cond.resample(time=freq).map(
+    out = resample_map(
+        cond,
+        "time",
+        freq,
         rl.last_run,
-        window=1,
-        dim="time",
-        coord="dayofyear",
+        map_kwargs=dict(window=1, dim="time", coord="dayofyear"),
     )
     out.attrs.update(units="", is_dayofyear=np.int32(1), calendar=get_calendar(prsn))
     return out
@@ -3093,7 +3101,7 @@ def _exceedance_date(grp):
         )
         return xarray.where((cumsum <= sum_thresh).all("time"), never_reached_val, out)
 
-    dded = c.clip(0).resample(time=freq).map(_exceedance_date)
+    dded = resample_map(c.clip(0), "time", freq, _exceedance_date)
     dded = dded.assign_attrs(
         units="", is_dayofyear=np.int32(1), calendar=get_calendar(tas)
     )

xclim/indices/generic.py

@@ -29,6 +29,7 @@
 from xclim.core.utils import DayOfYearStr, Quantified, Quantity
 
 from . import run_length as rl
+from .helpers import resample_map
 
 __all__ = [
     "aggregate_between_dates",
@@ -88,14 +89,15 @@ def select_resample_op(
         The maximum value for each period.
     """
     da = select_time(da, **indexer)
-    r = da.resample(time=freq)
     if op in _xclim_ops:
         op = _xclim_ops[op]
     if isinstance(op, str):
-        out = getattr(r, op.replace("integral", "sum"))(dim="time", keep_attrs=True)
+        out = getattr(da.resample(time=freq), op.replace("integral", "sum"))(
+            dim="time", keep_attrs=True
+        )
     else:
         with xr.set_options(keep_attrs=True):
-            out = r.map(op)
+            out = resample_map(da, "time", freq, op)
         op = op.__name__
     if out_units is not None:
         return out.assign_attrs(units=out_units)
@@ -544,7 +546,7 @@ def season(
     map_kwargs = dict(window=window, date=mid_date)
     if stat in ["start", "end"]:
         map_kwargs["coord"] = "dayofyear"
-    out = cond.resample(time=freq).map(FUNC[stat], **map_kwargs)
+    out = resample_map(cond, "time", freq, FUNC[stat], map_kwargs=map_kwargs)
     if stat == "length":
         return to_agg_units(out, data, "count")
     # else, a date
@@ -705,11 +707,12 @@ def first_occurrence(
 
     cond = compare(data, op, threshold, constrain)
 
-    out = cond.resample(time=freq).map(
+    out = resample_map(
+        cond,
+        "time",
+        freq,
         rl.first_run,
-        window=1,
-        dim="time",
-        coord="dayofyear",
+        map_kwargs=dict(window=1, dim="time", coord="dayofyear"),
     )
     out.attrs["units"] = ""
     return out
@@ -750,11 +753,12 @@ def last_occurrence(
 
     cond = compare(data, op, threshold, constrain)
 
-    out = cond.resample(time=freq).map(
+    out = resample_map(
+        cond,
+        "time",
+        freq,
         rl.last_run,
-        window=1,
-        dim="time",
-        coord="dayofyear",
+        map_kwargs=dict(window=1, dim="time", coord="dayofyear"),
     )
     out.attrs["units"] = ""
     return out
@@ -795,11 +799,12 @@ def spell_length(
 
     cond = compare(data, op, threshold)
 
-    out = cond.resample(time=freq).map(
+    out = resample_map(
+        cond,
+        "time",
+        freq,
         rl.rle_statistics,
-        reducer=reducer,
-        window=1,
-        dim="time",
+        map_kwargs=dict(reducer=reducer, window=1, dim="time"),
     )
     return to_agg_units(out, data, "count")
 
@@ -1139,12 +1144,12 @@ def first_day_threshold_reached(
 
     cond = compare(data, op, threshold, constrain=constrain)
 
-    out: xarray.DataArray = cond.resample(time=freq).map(
+    out = resample_map(
+        cond,
+        "time",
+        freq,
         rl.first_run_after_date,
-        window=window,
-        date=after_date,
-        dim="time",
-        coord="dayofyear",
+        map_kwargs=dict(window=window, date=after_date, dim="time", coord="dayofyear"),
     )
     out.attrs.update(units="", is_dayofyear=np.int32(1), calendar=get_calendar(data))
     return out