Use HighLevelGraph.merge in collections_to_dsk (#6836)

dask · Nov 13, 2020 · 1e0a9a1 · 1e0a9a1
1 parent aeb7815
commit 1e0a9a1
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/dask/base.py b/dask/base.py
@@ -263,6 +263,8 @@ def collections_to_dsk(collections, optimize_graph=True, **kwargs):
     """
     Convert many collections into a single dask graph, after optimization
     """
+    from .highlevelgraph import HighLevelGraph
+
     optimizations = kwargs.pop("optimizations", None) or config.get("optimizations", [])
 
     if optimize_graph:
@@ -284,12 +286,11 @@ def collections_to_dsk(collections, optimize_graph=True, **kwargs):
                 _opt_list.append(_opt)
             groups = group
 
-        dsk = merge(
-            *map(
-                ensure_dict,
-                _opt_list,
-            )
-        )
+        # Merge all graphs
+        if any(isinstance(graph, HighLevelGraph) for graph in _opt_list):
+            dsk = HighLevelGraph.merge(*_opt_list)
+        else:
+            dsk = merge(*map(ensure_dict, _opt_list))
     else:
         dsk, _ = _extract_graph_and_keys(collections)