Avoid sorting large stacks in order (#3298)

mrocklin · web-flow · commit f5779130fc2a · 2018-03-19T16:42:42.000-04:00
When performning task ordering we sort tasks based on the number of dependents/dependencies they have. This is critical to low-memory processing. However, sometimes individual tasks have millions of dependencies, for which an n*log(n) sort adds significant overhead. In these cases we give up on sorting, and just hope that the tasks are well ordered naturally (such as is often the case in Python 3.6+ due to sorted dicts and the natural ordering that exists when constructing common graphs) See pangeo-data/pangeo#150 (comment) for a real-world case
diff --git a/dask/order.py b/dask/order.py
@@ -109,7 +109,10 @@ def dependents_key(x):
     i = 0
 
     stack = [k for k, v in dependents.items() if not v]
-    stack = sorted(stack, key=dependencies_key)
+    if len(stack) < 10000:
+        stack = sorted(stack, key=dependencies_key)
+    else:
+        stack = stack[::-1]
 
     while stack:
         item = stack.pop()
@@ -120,7 +123,9 @@ def dependents_key(x):
         deps = waiting[item]
         if deps:
             stack.append(item)
-            stack.extend(sorted(deps, key=dependencies_key))
+            if len(deps) < 1000:
+                deps = sorted(deps, key=dependencies_key)
+            stack.extend(deps)
             continue
 
         result[item] = i
@@ -130,7 +135,9 @@ def dependents_key(x):
             waiting[dep].discard(item)
 
         deps = [d for d in dependents[item] if d not in result]
-        stack.extend(sorted(deps, key=dependents_key, reverse=True))
+        if len(deps) < 1000:
+            deps = sorted(deps, key=dependents_key, reverse=True)
+        stack.extend(deps)
 
     return result
 
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
@@ -30,6 +30,7 @@ Core
 
 - Fixed bug when using unexpected but hashable types for keys (:pr:`3238`) `Daniel Collins`_
 - Fix bug in task ordering so that we break ties consistently with the key name (:pr:`3271`) `Matthew Rocklin`_
+- Avoid sorting tasks in order when the number of tasks is very large (:pr:`3298`) `Matthew Rocklin`_
 
 
 0.17.1 / 2018-02-22
