Repository: dask/dask
Pull request #11943 (2025-05-10 14:30:19)

Fix Array.setitem with nan shapes on both sides
Closes https://github.com/dask/dask/issues/11753, specifically https://github.com/dask/dask/issues/11753#issuecomment-2868902284


github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

      9 files  ± 0        9 suites  ±0   3h 23m 1s ⏱️ + 5m 12s
 18 032 tests + 4   16 816 ✅ + 2   1 216 💤 + 2  0 ❌ ±0 
161 362 runs  +36  149 229 ✅ +19  12 133 💤 +17  0 ❌ ±0 

Results for commit 4e6eba7c. ± Comparison against base commit 7df5674f.

[test-results]:data:application/gzip;base64,H4sIAOdrH2gC/02OzQ6DIBAGX8Vw7sFdyAp9mQZxSUj9aRBOTd+9aJV6nPmyk30LH0Zexb0xt0asOaQKQ442hWUuCAgaiipj2mbQrcSTH2t2bpOkgf7yGV6/y7/zNozFtVVwjEs8TMzzniaQhAfXtDKIpsqalvJ0l/TO17JbpimkAkIxcW87h8xOq/Ia8YBEUnkLnXdtbzQpzyA+X2tlQxYYAQAA


Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index 5123c4ce3..4f89e3dc0 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -1934,7 +1934,14 @@ class Array(DaskMethodsMixin):
         if isinstance(key, Array):
             from dask.array.routines import where
 
-            if key.shape != self.shape:
+            left_shape = np.array(key.shape)
+            right_shape = np.array(self.shape)
+
+            # We want to treat unknown shape on *either* sides as a match
+            match = left_shape == right_shape
+            match |= np.isnan(left_shape) | np.isnan(right_shape)
+
+            if not match.all():
                 raise IndexError(
                     f"boolean index shape {key.shape} must match indexed array's "
                     f"{self.shape}."
diff --git a/dask/local.py b/dask/local.py
index 66b72f5ea..0c276aecb 100644
--- a/dask/local.py
+++ b/dask/local.py
@@ -187,6 +187,12 @@ def start_state_from_dask(dsk, cache=None, sortkey=None, keys=None):
         dependencies[key]
         task = dsk.get(key, None)
         if task is None:
+            if dependents[key] and not cache.get(key, None):
+                raise ValueError(
+                    "Missing dependency {} for dependents {}".format(
+                        key, dependents[key]
+                    )
+                )
             continue
         elif isinstance(task, DataNode):
             cache[key] = task()
diff --git a/dask/order.py b/dask/order.py
index b6db4578b..e33489029 100644
--- a/dask/order.py
+++ b/dask/order.py
@@ -195,7 +195,7 @@ def order(
     runnable_hull = set()
     reachable_hull = set()
 
-    runnable = []
+    runnable: list[Key] = []
 
     known_runnable_paths: dict[Key, list[list[Key]]] = {}
     known_runnable_paths_pop = known_runnable_paths.pop
@@ -203,12 +203,8 @@ def order(
     crit_path_counter = 0
     scrit_path: set[Key] = set()
     _crit_path_counter_offset: int | float = 0
-
     _sort_keys_cache: dict[Key, tuple[int, int, int, int, str]] = {}
 
-    leafs_connected_to_loaded_roots: set[Key] = set()
-    processed_roots = set()
-
     def sort_key(x: Key) -> tuple[int, int, int, int, str]:
         try:
             return _sort_keys_cache[x]
@@ -231,6 +227,7 @@ def order(
         # recursion depth errors. This is the only reason for the while loop
         next_items = [item]
         nonlocal i
+        nonlocal min_leaf_degree  # type: ignore[misc]
         while next_items:
             item = next_items.pop()
             runnable_hull.discard(item)
@@ -248,7 +245,25 @@ def order(
             if item not in external_keys:
                 i += 1
             if item in root_nodes:
-                processed_roots.add(item)
+                for leaf in leafs_connected[item]:
+                    if leaf in leaf_nodes:
+                        degree = leafs_degree.pop(leaf)
+                        leafs_per_degree[degree].remove(leaf)
+
+                        new_degree = degree - 1
+                        if new_degree > 0:
+                            if new_degree < min_leaf_degree:
+                                min_leaf_degree = new_degree
+                            leafs_per_degree[new_degree].add(leaf)
+                            leafs_degree[leaf] = new_degree
+                        elif not leafs_per_degree[degree]:
+                            assert degree == min_leaf_degree
+                            while min_leaf_degree != max_leaf_degree and (
+                                min_leaf_degree not in leafs_per_degree
+                                or not leafs_per_degree[min_leaf_degree]
+                            ):
+                                min_leaf_degree += 1
+
             # Note: This is a `set` and therefore this introduces a certain
             # randomness. However, this randomness should not have any impact on
             # the final result since the `process_runnable` should produce
@@ -396,42 +411,38 @@ def order(
     # writing, the most expensive part of ordering is the prep work (mostly
     # connected roots + sort_key) which can be reused for multiple orderings.
 
-    def get_target(longest_path: bool = False) -> Key:
-        # Some topologies benefit if the node with the most dependencies
-        # is used as first choice, others benefit from the opposite.
-        candidates = leaf_nodes
-        skey: Callable = sort_key
-        if reachable_hull:
-            skey = lambda k: (num_needed[k], sort_key(k))
-
-        all_leafs_accessible = len(leafs_connected_to_loaded_roots) >= len(candidates)
-        if reachable_hull and not all_leafs_accessible:
-            # Avoid this branch if we can already access all leafs. Just pick
-            # one of them without the expensive selection process in this case.
-
-            candidates = reachable_hull & candidates
+    # Degree in this context is the number root nodes that have to be loaded for
+    # this leaf to become accessible. Zero means the leaf is already accessible
+    # in which case it _should_ either already be in result or be accessible vie
+    # process_runnables
+    # When picking a new target, we prefer the leafs with the least number of
+    # roots that need loading.
+    leafs_degree = {}
+    leafs_per_degree = defaultdict(set)
+    min_leaf_degree = len(dsk)
+    max_leaf_degree = len(dsk)
+    for leaf in leaf_nodes - root_nodes:
+        degree = len(roots_connected[leaf])
+        min_leaf_degree = min(min_leaf_degree, degree)
+        max_leaf_degree = max(max_leaf_degree, degree)
+        leafs_degree[leaf] = degree
+        leafs_per_degree[degree].add(leaf)
+
+    def get_target() -> Key:
+        # If we're already mid run and there is a runnable_hull we'll attempt to
+        # pick the next target in a way that minimizes the number of additional
+        # root nodes that are needed
+        all_leafs_accessible = min_leaf_degree == max_leaf_degree
+        is_trivial_lookup = not reachable_hull or all_leafs_accessible
+        if not is_trivial_lookup:
+            candidates = reachable_hull & leafs_per_degree[min_leaf_degree]
             if not candidates:
-                # We can't reach a leaf node directly, but we still have nodes
-                # with results in memory, these notes can inform our path towards
-                # a new preferred leaf node.
-                for r in processed_roots:
-                    leafs_connected_to_loaded_roots.update(leafs_connected[r])
-                processed_roots.clear()
-
-                leafs_connected_to_loaded_roots.intersection_update(leaf_nodes)
-                candidates = leafs_connected_to_loaded_roots
-            else:
-                leafs_connected_to_loaded_roots.update(candidates)
-
-        elif not reachable_hull:
-            # We reach a new and independent branch, so throw away previous branch
-            leafs_connected_to_loaded_roots.clear()
-
-        if longest_path and (not reachable_hull or all_leafs_accessible):
-            return leaf_nodes_sorted.pop()
+                candidates = leafs_per_degree[min_leaf_degree]
+            # Even without reachable hull overlap this should be relatively
+            # small so one full pass should be fine
+            return min(candidates, key=sort_key)
         else:
-            # FIXME: This can be very expensive
-            return min(candidates, key=skey)
+            return leaf_nodes_sorted.pop()
 
     def use_longest_path() -> bool:
         size = 0
@@ -450,10 +461,10 @@ def order(
 
         return True
 
+    # Some topologies benefit if the node with the most dependencies
+    # is used as first choice, others benefit from the opposite.
     longest_path = use_longest_path()
-    leaf_nodes_sorted = []
-    if longest_path:
-        leaf_nodes_sorted = sorted(leaf_nodes, key=sort_key, reverse=False)
+    leaf_nodes_sorted = sorted(leaf_nodes, key=sort_key, reverse=not longest_path)
 
     # *************************************************************************
     # CORE ALGORITHM STARTS HERE
@@ -553,7 +564,7 @@ def order(
         assert not scrit_path
 
         # A. Build the critical path
-        target = get_target(longest_path=longest_path)
+        target = get_target()
         next_deps = dependencies[target]
         path_append(target)
 
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index 849301c25..5380ef914 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -5,6 +5,16 @@ Changelog
 
     This is not exhaustive. For an exhaustive list of changes, see the git log.
 
+
+.. _v2025.5.0:
+
+2025.5.0
+--------
+
+- Fixed Array setitem when both the array and the indexer have unknown shape (:issue:`11753`) `Tom Augspurger`_
+
+
+
 .. _v2025.4.1:
 
 2025.4.1


Test patch:
diff --git a/dask/array/tests/test_slicing.py b/dask/array/tests/test_slicing.py
index 28eef7841..e08ee7dcd 100644
--- a/dask/array/tests/test_slicing.py
+++ b/dask/array/tests/test_slicing.py
@@ -6,6 +6,7 @@ import warnings
 import pytest
 
 from dask._task_spec import Alias, Task, TaskRef
+from dask.delayed import delayed
 
 np = pytest.importorskip("numpy")
 
@@ -1137,3 +1138,39 @@ def test_positional_indexer_newaxis():
     arr = da.array([0, 1, 2])
     new = arr[[True, True, False], np.newaxis]
     assert_eq(new, arr.compute()[[True, True, False], np.newaxis])
+
+
+@pytest.mark.parametrize(
+    "shapes",
+    [
+        (10, 10),
+        (np.nan, np.nan),
+        pytest.param(
+            (10, np.nan), marks=pytest.mark.xfail(reason="Not implemented", strict=True)
+        ),
+        pytest.param(
+            (np.nan, 10), marks=pytest.mark.xfail(reason="Not implemented", strict=True)
+        ),
+    ],
+)
+def test_boolean_mask_with_unknown_shape(
+    shapes: tuple[float | int, float | int],
+) -> None:
+    x_shape, mask_shape = shapes
+    arr = delayed(np.ones(10))
+    x = da.concatenate(
+        [
+            da.from_delayed(arr, shape=(x_shape,), dtype=float),
+            da.from_delayed(arr, shape=(x_shape,), dtype=float),
+        ]
+    )
+    mask = da.concatenate(
+        [
+            da.from_delayed(arr, shape=(mask_shape,), dtype=bool),
+            da.from_delayed(arr, shape=(mask_shape,), dtype=bool),
+        ]
+    )
+    x[mask] = 2
+
+    expected = np.full(20, 2.0)
+    assert_eq(x, expected)
diff --git a/dask/tests/test_local.py b/dask/tests/test_local.py
index 684e43949..016443820 100644
--- a/dask/tests/test_local.py
+++ b/dask/tests/test_local.py
@@ -3,7 +3,7 @@ from __future__ import annotations
 import pytest
 
 import dask
-from dask._task_spec import Task
+from dask._task_spec import Task, TaskRef
 from dask.local import finish_task, get_sync, sortkey, start_state_from_dask
 from dask.order import order
 from dask.utils_test import GetFunctionTestMixin, add, inc
@@ -213,3 +213,10 @@ def test_ensure_calculate_only_whats_needed():
     ]
     dsk = {t.key: t for t in tasks}
     assert get_sync(dsk, "y") == 2
+
+
+def test_detect_malformed_graph():
+    dsk = {"y": Task("y", lambda x: x, TaskRef("x"))}
+    with pytest.raises(ValueError, match="Missing dependency"):
+        get_sync(dsk, "y")
+    assert get_sync(dsk, "y", cache={"x": 1}) == 1
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index 4074c5b78..b93c95eb0 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -560,10 +560,10 @@ def test_run_smaller_sections(abcde):
     }
     o = order(dsk)
     assert_topological_sort(dsk, o)
-    assert max(diagnostics(dsk)[1]) <= 4  # optimum is 3
-    # This is a mildly ambiguous example
+    # The actual order is a bit ambiguous but the optimal pressure is well
+    # defined, see also
     # https://github.com/dask/dask/pull/10535/files#r1337528255
-    assert (o[aa] < o[a] and o[dd] < o[a]) or (o[b] < o[e] and o[b] < o[cc])
+    assert max(diagnostics(dsk)[1]) == 3
 
 
 def test_local_parents_of_reduction(abcde):