Support query optimization with Dask expression arrays

xarray/compat/dask_array_compat.py

@@ -1,5 +1,6 @@
 from typing import Any
 
+from xarray.namedarray.parallelcompat import get_chunked_array_type
 from xarray.namedarray.utils import module_available
 
 
@@ -8,10 +9,17 @@ def reshape_blockwise(
     shape: int | tuple[int, ...],
     chunks: tuple[tuple[int, ...], ...] | None = None,
 ):
-    if module_available("dask", "2024.08.2"):
-        from dask.array import reshape_blockwise
+    try:
+        array_api = get_chunked_array_type(x).array_api
+    except TypeError:
+        array_api = None
 
-        return reshape_blockwise(x, shape=shape, chunks=chunks)
+    if array_api is not None and hasattr(array_api, "reshape_blockwise"):
+        return array_api.reshape_blockwise(x, shape=shape, chunks=chunks)
+    elif module_available("dask", "2024.08.2"):
+        from dask.array import reshape_blockwise as dask_reshape_blockwise
+
+        return dask_reshape_blockwise(x, shape=shape, chunks=chunks)
     else:
         return x.reshape(shape)
 

xarray/compat/dask_array_ops.py

@@ -4,6 +4,7 @@
 
 from xarray.compat.dask_array_compat import reshape_blockwise
 from xarray.core import dtypes, nputils
+from xarray.namedarray.parallelcompat import get_chunked_array_type
 
 
 def dask_rolling_wrapper(moving_func, a, window, min_count=None, axis=-1):
@@ -20,7 +21,7 @@ def dask_rolling_wrapper(moving_func, a, window, min_count=None, axis=-1):
 
 
 def least_squares(lhs, rhs, rcond=None, skipna=False):
-    import dask.array as da
+    da = get_chunked_array_type(rhs).array_api
 
     # The trick here is that the core dimension is axis 0.
     # All other dimensions need to be reshaped down to one axis for `lstsq`
@@ -94,23 +95,25 @@ def push(array, n, axis, method="blelloch"):
     """
     Dask-aware bottleneck.push
     """
-    import dask.array as da
     import numpy as np
 
     from xarray.core.duck_array_ops import _push
     from xarray.core.nputils import nanlast
 
+    chunkmanager = get_chunked_array_type(array)
+    da = chunkmanager.array_api
+
     if n is not None and all(n <= size for size in array.chunks[axis]):
         return array.map_overlap(_push, depth={axis: (n, 0)}, n=n, axis=axis)
 
     # TODO: Replace all this function
     #  once https://github.com/pydata/xarray/issues/9229 being implemented
 
-    pushed_array = da.reductions.cumreduction(
+    pushed_array = chunkmanager.scan(
         func=_dtype_push,
         binop=_fill_with_last_one,
         ident=np.nan,
-        x=array,
+        arr=array,
         axis=axis,
         dtype=array.dtype,
         method=method,

xarray/core/dataarray.py

@@ -1113,6 +1113,13 @@ def __dask_graph__(self):
     def __dask_keys__(self):
         return self._to_temp_dataset().__dask_keys__()
 
+    def __dask_exprs__(self):
+        return self._to_temp_dataset().__dask_exprs__()
+
+    def __dask_rebuild_from_exprs__(self, exprs):
+        ds = self._to_temp_dataset().__dask_rebuild_from_exprs__(exprs)
+        return self._from_temp_dataset(ds)
+
     def __dask_layers__(self):
         return self._to_temp_dataset().__dask_layers__()
 

xarray/core/dataset.py

@@ -648,14 +648,17 @@ def __dask_graph__(self):
         if not graphs:
             return None
         else:
-            try:
-                from dask.highlevelgraph import HighLevelGraph
+            from dask.highlevelgraph import HighLevelGraph
 
+            if all(isinstance(graph, HighLevelGraph) for graph in graphs.values()):
                 return HighLevelGraph.merge(*graphs.values())
-            except ImportError:
-                from dask import sharedict
 
-                return sharedict.merge(*graphs.values())
+            from dask.utils import ensure_dict
+
+            merged = {}
+            for graph in graphs.values():
+                merged.update(ensure_dict(graph))
+            return merged
 
     def __dask_keys__(self):
         import dask
@@ -666,6 +669,56 @@ def __dask_keys__(self):
             if dask.is_dask_collection(v)
         ]
 
+    def __dask_exprs__(self):
+        from importlib import import_module
+
+        import dask
+
+        try:
+            DaskArray = import_module("dask_array").Array
+        except ImportError:
+            return None
+
+        exprs = []
+        for v in self.variables.values():
+            if dask.is_dask_collection(v):
+                if not isinstance(v._data, DaskArray):
+                    # Composite expressions must account for every Dask-backed
+                    # variable.  Returning None keeps Dask's collection APIs on
+                    # the existing HighLevelGraph path for mixed
+                    # legacy/expression datasets.
+                    return None
+                exprs.append(v._data.expr)
+        return exprs or None
+
+    def __dask_rebuild_from_exprs__(self, exprs):
+        import dask
+        from dask._collections import new_collection
+
+        dask_variables = [
+            (k, v) for k, v in self._variables.items() if dask.is_dask_collection(v)
+        ]
+        exprs = list(exprs)
+        if len(exprs) != len(dask_variables):
+            raise ValueError(
+                f"Expected {len(dask_variables)} expressions to rebuild Dataset, "
+                f"got {len(exprs)}"
+            )
+
+        variables = dict(self._variables)
+        for (k, v), expr in zip(dask_variables, exprs, strict=True):
+            variables[k] = v._replace(data=new_collection(expr))
+
+        return type(self)._construct_direct(
+            variables,
+            self._coord_names,
+            self._dims,
+            self._attrs,
+            self._indexes,
+            self._encoding,
+            self._close,
+        )
+
     def __dask_layers__(self):
         import dask
 

xarray/core/parallel.py

@@ -440,6 +440,7 @@ def _wrapper(
         dataarray_to_dataset(arg) if isinstance(arg, DataArray) else arg
         for arg in aligned
     )
+
     # rechunk any numpy variables appropriately
     xarray_objs = tuple(arg.chunk(arg.chunksizes) for arg in xarray_objs)