Repository: dask/dask
Pull request #11737 (2025-02-12 17:06:53)

Remove culling from delayed optimisation
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`


culling here has very limited use, we are culling scheduler side anyway and this can be catastrophic for performance in XGBoost applications


Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index 6a7ba4f9f..afe68e933 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -1938,30 +1938,43 @@ class Array(DaskMethodsMixin):
         if value is np.ma.masked:
             value = np.ma.masked_all((), dtype=self.dtype)
 
-        if (
-            not is_dask_collection(value)
-            and issubclass(self.dtype.type, Integral)
-            and np.isnan(value).any()
+        if not is_dask_collection(value) and self.dtype.kind in "iu":
+            if np.isnan(value).any():
+                raise ValueError("cannot convert float NaN to integer")
+            if np.isinf(value).any():
+                raise ValueError("cannot convert float infinity to integer")
+
+        # Suppress dtype broadcasting; __setitem__ can't change the dtype.
+        # Use asanyarray to retain e.g. np.ma objects.
+        value = asanyarray(value, dtype=self.dtype, like=self)
+
+        if isinstance(key, Array) and (
+            key.dtype.kind in "iu"
+            or (key.dtype == bool and key.ndim == 1 and self.ndim > 1)
         ):
-            raise ValueError("cannot convert float NaN to integer")
+            key = (key,)
 
-        ## Use the "where" method for cases when key is an Array
+        ## Use the "where" method for cases when key is an Array of bools
         if isinstance(key, Array):
             from dask.array.routines import where
 
-            if isinstance(value, Array) and value.ndim > 1:
-                raise ValueError("boolean index array should have 1 dimension")
-            try:
-                y = where(key, value, self)
-            except ValueError as e:
-                raise ValueError(
-                    "Boolean index assignment in Dask "
-                    "expects equally shaped arrays.\nExample: da1[da2] = da3 "
-                    "where da1.shape == (4,), da2.shape == (4,) "
-                    "and da3.shape == (4,).\n"
-                    "Alternatively, you can use the extended API that supports"
-                    "indexing with tuples.\nExample: da1[(da2,)] = da3."
-                ) from e
+            if key.shape != self.shape:
+                raise IndexError(
+                    f"boolean index shape {key.shape} must match indexed array's "
+                    f"{self.shape}."
+                )
+
+            # If value has ndim > 0, they must be broadcastable to self.shape[idx].
+            # This raises when the bool mask causes the size to become unknown,
+            # e.g. this is valid in numpy but raises here:
+            # x = da.array([1,2,3])
+            # x[da.array([True, True, False])] = [4, 5]
+            if value.ndim:
+                value = broadcast_to(value, self[key].shape)
+
+            y = where(key, value, self)
+            # FIXME does any backend allow mixed ops vs. numpy?
+            # If yes, is it wise to let them change the meta?
             self._meta = y._meta
             self.dask = y.dask
             self._name = y.name
@@ -1973,7 +1986,6 @@ class Array(DaskMethodsMixin):
 
         # Still here? Then apply the assignment to other type of
         # indices via the `setitem_array` function.
-        value = asanyarray(value)
 
         out = "setitem-" + tokenize(self, key, value)
         dsk = setitem_array(out, self, key, value)
diff --git a/dask/dataframe/dask_expr/io/io.py b/dask/dataframe/dask_expr/io/io.py
index 0753a6479..e0edf672d 100644
--- a/dask/dataframe/dask_expr/io/io.py
+++ b/dask/dataframe/dask_expr/io/io.py
@@ -7,6 +7,7 @@ import operator
 import numpy as np
 import pyarrow as pa
 
+from dask import delayed
 from dask._task_spec import DataNode, List, Task
 from dask.dataframe import methods
 from dask.dataframe._pyarrow import to_pyarrow_string
@@ -256,10 +257,11 @@ class FromMap(PartitionsFiltered, BlockwiseIO):
     def _meta(self):
         if self.operand("user_meta") is not no_default:
             meta = self.operand("user_meta")
+            return make_meta(meta)
         else:
             vals = [v[0] for v in self.iterables]
-            meta = self.func(*vals, *self.args, **self.kwargs)
-        return make_meta(meta)
+            meta = delayed(self.func)(*vals, *self.args, **self.kwargs)
+            return delayed(make_meta)(meta).compute()
 
     def _divisions(self):
         if self.operand("user_divisions"):
diff --git a/dask/delayed.py b/dask/delayed.py
index 1718ca88a..e02b60699 100644
--- a/dask/delayed.py
+++ b/dask/delayed.py
@@ -20,7 +20,7 @@ from dask.base import (
 )
 from dask.base import tokenize as _tokenize
 from dask.context import globalmethod
-from dask.core import flatten, quote
+from dask.core import quote
 from dask.highlevelgraph import HighLevelGraph
 from dask.optimization import fuse
 from dask.typing import Graph, NestedKeys
@@ -539,7 +539,6 @@ def optimize(dsk, keys, **kwargs):
 
     if not isinstance(dsk, HighLevelGraph):
         dsk = HighLevelGraph.from_collections(id(dsk), dsk, dependencies=())
-    dsk = dsk.cull(set(flatten(keys)))
     return dsk
 
 


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index b7ffb5783..d319c8b1d 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -4184,16 +4184,14 @@ def test_setitem_1d():
 
     x[x > 6] = -1
     x[x % 2 == 0] = -2
+    x[[2, 3]] = -3
 
     dx[dx > 6] = -1
     dx[dx % 2 == 0] = -2
+    dx[da.asarray([2, 3])] = -3
 
     assert_eq(x, dx)
 
-    index = da.arange(3)
-    with pytest.raises(ValueError, match="Boolean index assignment in Dask"):
-        dx[index] = 1
-
 
 def test_setitem_masked():
     # Test np.ma.masked assignment to object-type arrays
@@ -4467,7 +4465,6 @@ def test_setitem_extended_API_2d_rhs_func_of_lhs():
 def test_setitem_extended_API_2d_mask(index, value):
     x = np.ma.arange(60).reshape((6, 10))
     dx = da.from_array(x.data, chunks=(2, 3))
-    dx[index] = value
     # See https://github.com/numpy/numpy/issues/23000 for the `RuntimeWarning`
     with warnings.catch_warnings():
         warnings.filterwarnings(
@@ -4476,7 +4473,8 @@ def test_setitem_extended_API_2d_mask(index, value):
             message="invalid value encountered in cast",
         )
         x[index] = value
-        dx = dx.persist()
+        dx[index] = value
+    dx = dx.persist()
     assert_eq(x, dx.compute())
     assert_eq(x.mask, da.ma.getmaskarray(dx).compute())
 
@@ -4535,13 +4533,6 @@ def test_setitem_errs():
     with pytest.raises(IndexError):
         x[[[True, True, False, False]], 0] = 5
 
-    # Too many/not enough booleans
-    with pytest.raises(IndexError):
-        x[[True, True, False]] = 5
-
-    with pytest.raises(IndexError):
-        x[[False, True, True, True, False]] = 5
-
     # 2-d indexing a single dimension
     with pytest.raises(IndexError):
         x[[[1, 2, 3]], 0] = 5
@@ -4586,6 +4577,45 @@ def test_setitem_errs():
     x = da.ones((3, 3), dtype=int)
     with pytest.raises(ValueError, match="cannot convert float NaN to integer"):
         x[:, 1] = np.nan
+    with pytest.raises(ValueError, match="cannot convert float infinity to integer"):
+        x[:, 1] = np.inf
+    with pytest.raises(ValueError, match="cannot convert float infinity to integer"):
+        x[:, 1] = -np.inf
+
+
+@pytest.mark.parametrize("idx_namespace", [np, da])
+def test_setitem_bool_index_errs(idx_namespace):
+    x = da.ones((3, 4), chunks=(2, 2))
+    y = da.ones(4, chunks=2)
+    array = idx_namespace.array
+
+    # Shape mismatch
+    with pytest.raises(ValueError):
+        y[array([True, True, True, False])] = [2, 3]
+
+    with pytest.raises(ValueError):
+        # A naive where(idx, val, x) would produce a result
+        y[array([True, True, True, False])] = [1, 2, 3, 4]
+
+    with pytest.raises(ValueError):
+        y[array([True, True, True, False])] = [1, 2, 3, 4, 5]
+
+    with pytest.raises(ValueError):
+        y[array([True, False, False, True])] = [1, 2, 3, 4, 5]
+
+    # Too many/not enough booleans
+    with pytest.raises(IndexError):
+        y[array([True, False, True])] = 5
+
+    with pytest.raises(IndexError):
+        y[array([False, True, True, True, False])] = 5
+
+    # Situations where a naive da.where(idx, val, x) would produce a result
+    with pytest.raises(IndexError):
+        x[array([True, False, False, True])] = 1
+
+    with pytest.raises(IndexError):
+        y[array([[True], [False]])] = 1  # da.where would broadcast to ndim=2
 
 
 def test_zero_slice_dtypes():
@@ -5788,3 +5818,16 @@ def test_scalar_setitem():
     y[()] = 2
     assert_eq(y, 2.0)
     assert isinstance(y.compute(), np.ndarray)
+
+
+@pytest.mark.parametrize(
+    "idx", [[0], [True, False], da.array([0]), da.array([True, False])]
+)
+@pytest.mark.parametrize(
+    "val",
+    [3.3, np.float64(3.3), np.int64(3), da.array(3.3), da.array(3, dtype=np.int64)],
+)
+def test_setitem_no_dtype_broadcast(idx, val):
+    x = da.array([1, 2], dtype=np.int32)
+    x[idx] = val
+    assert_eq(x, da.array([3, 2], dtype=np.int32))
diff --git a/dask/array/tests/test_cupy_core.py b/dask/array/tests/test_cupy_core.py
index 71c0a3b25..750d2c65b 100644
--- a/dask/array/tests/test_cupy_core.py
+++ b/dask/array/tests/test_cupy_core.py
@@ -273,15 +273,21 @@ def test_store_kwargs():
     assert called[0]
 
 
-def test_setitem_1d():
+@pytest.mark.parametrize(
+    "chunks", (-1, pytest.param(5, marks=pytest.mark.xfail(reason="dask/dask#11730")))
+)
+def test_setitem_1d(chunks):
     x = cupy.arange(10)
-    dx = da.from_array(x.copy(), chunks=(5,))
+    dx = da.from_array(x.copy(), chunks=chunks)
 
     x[x > 6] = -1
     x[x % 2 == 0] = -2
+    idx = cupy.asarray([2, 3])
+    x[idx,] = -3
 
     dx[dx > 6] = -1
     dx[dx % 2 == 0] = -2
+    dx[da.from_array(idx)] = -3
 
     assert_eq(x, dx)
 
diff --git a/dask/dataframe/dask_expr/io/tests/test_delayed.py b/dask/dataframe/dask_expr/io/tests/test_delayed.py
index 42508dcd3..bc63534d9 100644
--- a/dask/dataframe/dask_expr/io/tests/test_delayed.py
+++ b/dask/dataframe/dask_expr/io/tests/test_delayed.py
@@ -14,7 +14,6 @@ pd = _backend_library()
 def test_from_delayed_optimizing():
     parts = from_dict({"a": np.arange(300)}, npartitions=30).to_delayed()
     result = from_delayed(parts[0], meta=pd.DataFrame({"a": pd.Series(dtype=np.int64)}))
-    assert len(result.optimize().dask) == 2
     assert_eq(result, pd.DataFrame({"a": pd.Series(np.arange(10))}))
 
 
diff --git a/dask/tests/test_delayed.py b/dask/tests/test_delayed.py
index 4ce1dbc8e..124c53fda 100644
--- a/dask/tests/test_delayed.py
+++ b/dask/tests/test_delayed.py
@@ -331,8 +331,8 @@ def test_common_subexpressions():
 def test_delayed_optimize():
     x = Delayed("b", {"a": 1, "b": (inc, "a"), "c": (inc, "b")})
     (x2,) = dask.optimize(x)
-    # Delayed's __dask_optimize__ culls out 'c'
-    assert sorted(x2.dask.keys()) == ["a", "b"]
+    # Delayed's __dask_optimize__ does not cull out 'c'
+    assert sorted(x2.dask.keys()) == ["a", "b", "c"]
     assert x2._layer != x2._key
     # Optimize generates its own layer name, which doesn't match the key.
     # `Delayed._rebuild` handles this.
@@ -843,7 +843,7 @@ def test_annotations_survive_optimization():
     (d_opt,) = dask.optimize(d)
     assert type(d_opt.dask) is HighLevelGraph
     assert len(d_opt.dask.layers) == 1
-    assert len(d_opt.dask.layers["b"]) == 2  # c is culled
+    assert len(d_opt.dask.layers["b"]) == 3  # c is not culled
     assert d_opt.dask.layers["b"].annotations == {"foo": "bar"}