Use Taskspec fuse implementation

dask_expr/_expr.py

@@ -3766,32 +3766,29 @@ def _broadcast_dep(self, dep: Expr):
         return dep.npartitions == 1
 
     def _task(self, name: Key, index: int) -> Task:
-        internal_tasks = []
-        seen_keys = set()
-        external_deps = set()
+        internal_tasks = {}
         for _expr in self.exprs:
             if self._broadcast_dep(_expr):
                 subname = (_expr._name, 0)
             else:
                 subname = (_expr._name, index)
             t = _expr._task(subname, subname[1])
+
             assert t.key == subname
-            internal_tasks.append(t)
-            seen_keys.add(subname)
-            external_deps.update(t.dependencies)
-        external_deps -= seen_keys
-        dependencies = {dep: TaskRef(dep) for dep in external_deps}
-        t = Task(
-            name,
-            Fused._execute_internal_graph,
-            # Wrap the actual subgraph as a data node such that the tasks are
-            # not erroneously parsed. The external task would otherwise carry
-            # the internal keys as dependencies which is not satisfiable
-            DataNode(None, internal_tasks),
-            dependencies,
-            (self.exprs[0]._name, index),
-        )
-        return t
+            internal_tasks[t.key] = t
+        # The above is ambiguous and we have to cull to get an unambiguous graph
+        outkey = (self.exprs[0]._name, index)
+        work = [outkey]
+        internal_tasks_culled = []
+        while work:
+            tkey = work.pop()
+            if tkey not in internal_tasks:
+                # External dependency
+                continue
+            t = internal_tasks[tkey]
+            internal_tasks_culled.append(t)
+            work.extend(t.dependencies)
+        return Task.fuse(*internal_tasks_culled, key=name)
 
     @staticmethod
     def _execute_internal_graph(internal_tasks, dependencies, outkey):