bind implementation

crusaderky · crusaderky · commit c68a459e5c3d · 2021-01-25T21:57:43.000Z
diff --git a/dask/graph_manipulation.py b/dask/graph_manipulation.py
@@ -199,32 +199,43 @@ def bind(
 
     omit, _ = unpack_collections(omit)
     if assume_layers:
-        bind_one = _bind_one_by_layers
-        omit_set = {layer for coll in omit for layer in coll.__dask_layers__()}
+        omit_layers = {layer for coll in omit for layer in coll.__dask_layers__()}
+        omit_keys = set()
     else:
         # One or more collections do not use HighLevelGraph; downgrade all graphs to
         # plain dicts
-        bind_one = _bind_one_by_keys
-        omit_set = {key for coll in omit for key in ensure_dict(coll.__dask_graph__())}
+        omit_layers = set()
+        omit_keys = {key for coll in omit for key in ensure_dict(coll.__dask_graph__())}
 
     return repack(
-        [bind_one(child, blocker, omit_set, seed) for child in unpacked_children]
+        [
+            _bind_one(child, blocker, omit_layers, omit_keys, seed)
+            for child in unpacked_children
+        ]
     )
 
 
-def _bind_one_by_layers(
-    child: T, blocker: Optional[Delayed], omit_layers: Set[str], seed: Hashable
+def _bind_one(
+    child: T,
+    blocker: Optional[Delayed],
+    omit_layers: Set[str],
+    omit_keys: Set[str],
+    seed: Hashable,
 ) -> T:
     try:
-        old_name = get_collection_name(child)
+        name = get_collection_name(child)
     except KeyError:
         return child  # Collection with no keys; e.g. Array of size 0
 
     dsk = child.__dask_graph__()
-    assert isinstance(dsk, HighLevelGraph)
+    if isinstance(dsk, HighLevelGraph):
+        new_layers = dict(dsk.layers)
+        new_deps = dict(dsk.dependencies)
+    else:
+        # Squash layers, if any
+        new_layers = {name: ensure_dict(dsk)}
+        new_deps = {name: set()}
 
-    new_layers = dict(dsk.layers)
-    new_deps = dict(dsk.dependencies)
     if blocker is not None:
         blocker_key = blocker.key
         blocker_dsk = blocker.__dask_graph__()
@@ -234,7 +245,10 @@ def _bind_one_by_layers(
     else:
         blocker_key = None
 
-    to_visit = set(child.__dask_layers__())
+    try:
+        to_visit = set(child.__dask_layers__())
+    except AttributeError:
+        to_visit = {name}
 
     while to_visit:
         name = to_visit.pop()
@@ -246,33 +260,29 @@ def _bind_one_by_layers(
         deps_to_visit = deps - omit_layers
         to_visit |= deps_to_visit
 
-        bind_before_keys = {
+        # If assume_layers=True (the default), the stop_before_keys set will be
+        # non-empty only in the layers immediately above the top layers of the
+        # collections passed through the 'omit' parameter.
+        # If assume_layers=False, it will instead always be populated with the output
+        # of __dask_keys__ of the collections passed through the 'omit' parameter.
+        stop_before_keys = {
             key for dep in deps_to_omit for key in new_layers[dep].get_external_keys()
-        }
-        new_layers[new_name] = layer.bind(
-            to_key=blocker_key, before_keys=bind_before_keys, seed=seed
+        } | omit_keys
+        new_layers[new_name] = layer.clone(
+            stop_before=stop_before_keys, bind_to=blocker_key, seed=seed
         )
         new_deps[new_name] = {
             clone_key(dep, seed=seed) for dep in deps_to_visit
         } | deps_to_omit
+        if deps_to_omit and blocker_key:
+            new_deps[new_name].add(blocker_key)
 
     rebuild, args = child.__dask_postpersist__()
     return rebuild(
-        HighLevelGraph(new_layers, new_deps), *args, name=clone_key(old_name, seed)
+        HighLevelGraph(new_layers, new_deps), *args, name=clone_key(name, seed)
     )
 
 
-def _bind_one_by_keys(
-    child: T, blocker: Optional[Delayed], omit_keys: set, seed: Hashable
-) -> T:
-    dsk = ensure_dict(child.__dask_graph__())
-
-    # TODO
-
-    rebuild, args = child.__dask_postpersist__()
-    return rebuild(dsk, *args, name="bind-" + tokenize(child, seed))
-
-
 def clone(*collections, omit=None, seed: Hashable = None, assume_layers: bool = True):
     """Clone dask collections, returning equivalent collections that are generated from
     independent calculations.
