bugfixes

crusaderky · crusaderky · commit 802ece799399 · 2021-01-28T16:22:47.000Z
diff --git a/dask/graph_manipulation.py b/dask/graph_manipulation.py
@@ -228,7 +228,7 @@ def _bind_one(
     seed: Hashable,
 ) -> T:
     try:
-        name = get_collection_name(child)
+        prev_coll_name = get_collection_name(child)
     except KeyError:
         return child  # Collection with no keys; e.g. Array of size 0
 
@@ -241,11 +241,11 @@ def _bind_one(
         all_keys = dsk.get_all_external_keys()
     else:
         dsk = ensure_dict(dsk)
-        prev_layers = {name: dsk}
-        prev_deps = {name: set()}
+        prev_layers = {prev_coll_name: dsk}
+        prev_deps = {prev_coll_name: set()}
         all_keys = dsk.keys()
 
-    replace_keys = all_keys - omit_keys
+    clone_keys = all_keys - omit_keys
 
     if blocker is not None:
         blocker_key = blocker.key
@@ -259,52 +259,53 @@ def _bind_one(
     try:
         deps_to_clone = set(child.__dask_layers__())
     except AttributeError:
-        deps_to_clone = {name}
+        deps_to_clone = {prev_coll_name}
     deps_to_copy_verbatim = set()
 
     while deps_to_clone:
-        name = deps_to_clone.pop()
-        new_name = clone_key(name, seed=seed)
-        if new_name in new_layers:
+        prev_layer_name = deps_to_clone.pop()
+        new_layer_name = clone_key(prev_layer_name, seed=seed)
+        if new_layer_name in new_layers:
             continue
 
-        layer = prev_layers[name]
-        layer_deps = prev_deps[name]
+        layer = prev_layers[prev_layer_name]
+        layer_deps = prev_deps[prev_layer_name]
         layer_deps_to_omit = layer_deps & omit_layers
         layer_deps_to_clone = layer_deps - omit_layers
         deps_to_clone |= layer_deps_to_clone
         deps_to_copy_verbatim |= layer_deps_to_omit
 
-        omit_keys_from_layers = {
-            key
-            for dep in layer_deps_to_omit
-            for key in new_layers[dep].get_output_keys()
-        }
-        new_layers[new_name], bound = layer.clone(
-            keys=replace_keys - omit_keys_from_layers,
+        layer_clone_keys = clone_keys.copy()
+        for dep in layer_deps_to_omit:
+            layer_clone_keys -= new_layers[dep].get_output_keys()
+
+        new_layers[new_layer_name], bound = layer.clone(
+            keys=layer_clone_keys,
             seed=seed,
             bind_to=blocker_key,
         )
-        new_deps[new_name] = {
+        new_deps[new_layer_name] = {
             clone_key(dep, seed=seed) for dep in layer_deps_to_clone
         } | layer_deps_to_omit
         if bound:
-            new_deps[new_name].add(blocker_key)
+            new_deps[new_layer_name].add(blocker_key)
 
     # Add the layers of the collections from omit from child.dsk. Note that, when
     # assume_layers=False, it would be unsafe to simply do HighLevelGraph.merge(dsk,
     # omit[i].dsk). Also, collections in omit may or may not be parents of this specific
     # child, or of any children at all.
     while deps_to_copy_verbatim:
-        name = deps_to_clone.pop()
-        layer_deps = prev_deps[name]
+        layer_name = deps_to_clone.pop()
+        layer_deps = prev_deps[layer_name]
         deps_to_copy_verbatim |= layer_deps
-        new_deps[name] = layer_deps
-        new_layers[name] = prev_layers[name]
+        new_deps[layer_name] = layer_deps
+        new_layers[layer_name] = prev_layers[layer_name]
 
     rebuild, args = child.__dask_postpersist__()
     return rebuild(
-        HighLevelGraph(new_layers, new_deps), *args, name=clone_key(name, seed)
+        HighLevelGraph(new_layers, new_deps),
+        *args,
+        name=clone_key(prev_coll_name, seed),
     )
 
 
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
@@ -246,11 +246,12 @@ def clone_value(o):
                 return {k: clone_value(v) for k, v in o.items()}
             else:
                 try:
-                    if o in keys:
-                        is_leaf = False
-                        return clone_key(o, seed)
+                    if o not in keys:
+                        return o
                 except TypeError:
                     return o
+                is_leaf = False
+                return clone_key(o, seed)
 
         dsk_new = {}
         bound = False
