Blocklist-based implementation

dask/base.py

@@ -210,28 +210,25 @@ def is_dask_collection(x) -> bool:
     implementation of the protocol.
 
     """
-    try:
-        if "xarray" in type(x).__module__:
-            import xarray
-
-            if isinstance(x, xarray.Dataset):
-                return any(is_dask_collection(v._data) for _, v in x.variables.items())
-            elif isinstance(x, xarray.DataArray):
-                return is_dask_collection(x.variable._data)
-            elif isinstance(x, xarray.Variable):
-                return is_dask_collection(x._data)
-            else:
-                raise TypeError("Unfamiliar with xarray type", type(x))
-        else:
-            return (
-                hasattr(x, "__dask_graph__")
-                and callable(x.__dask_graph__)
-                and not isinstance(x, type)
-            )
-
-    except (AttributeError, TypeError):
+    if (
+        isinstance(x, type)
+        or not hasattr(x, "__dask_graph__")
+        or not callable(x.__dask_graph__)
+    ):
         return False
 
+    pkg_name = getattr(type(x), "__module__", "").split(".")[0]
+    if pkg_name == "dask_expr":
+        # Temporary hack to avoid graph materialization. Note that this won't work with
+        # dask_expr.array objects wrapped by xarray or pint. By the time dask_expr.array
+        # is published, we hope to be able to rewrite this method completely.
+        # Read: https://github.com/dask/dask/pull/10676
+        return True
+
+    # In all known dask collections other than dask-expr,
+    # calling __dask_graph__ is cheap
+    return x.__dask_graph__() is not None
+
 
 class DaskMethodsMixin:
     """A mixin adding standard dask collection methods"""

dask/tests/test_base.py

@@ -677,10 +677,39 @@ def __dask_graph__(self):
 
     x = delayed(1) + 2
     assert is_dask_collection(x)
+    assert not is_dask_collection(2)
     assert is_dask_collection(DummyCollection({}))
     assert not is_dask_collection(DummyCollection)
 
 
+def test_is_dask_collection_dask_expr():
+    pd = pytest.importorskip("pandas")
+    dx = pytest.importorskip("dask_expr")
+
+    df = pd.Series([1, 2, 3])
+    dxf = dx.from_pandas(df)
+    assert not is_dask_collection(df)
+    assert is_dask_collection(dxf)
+
+
+def test_is_dask_collection_dask_expr_does_not_materialize():
+    dx = pytest.importorskip("dask_expr")
+
+    class DoNotMaterialize(dx.Expr):
+        @property
+        def _meta(self):
+            return 0
+
+        def __dask_graph__(self):
+            assert False, "must not reach"
+
+    dyf = dx.new_collection(DoNotMaterialize())
+    assert is_dask_collection(dyf)
+
+    with pytest.raises(AssertionError, match="must not reach"):
+        dyf.__dask_graph__()
+
+
 def test_unpack_collections():
     class ANamedTuple(NamedTuple):
         a: int  # type: ignore[annotation-unchecked]