Repository: dask/dask

Issue:
Extract long lists passed as arguments into separate tasks
In #4727 it was identified that a long list included in many tasks led to slowdowns for 3 reasons:

- Longer optimization time (need to traverse the list looking for tasks multiple times)
- Longer execution time (scheduler traverses lists in tasks during execution)
- Longer serialization time (including large data in each task is a general no-no).

We already pull out arguments larger than 1e6 bytes into a separate task (wrapped in delayed), but for lists which are interpreted by the optimization/scheduling passes this threshold should likely be lower. If it's not already done, this should also be done for the other collections.


Discussion:

Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -299,6 +299,8 @@ def normalize_arg(x):
         return x
     elif isinstance(x, str) and re.match(r'_\d+', x):
         return delayed(x)
+    elif isinstance(x, list) and len(x) >= 10:
+        return delayed(x)
     elif sizeof(x) > 1e6:
         return delayed(x)
     else:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -32,7 +32,6 @@
 from ..array.core import Array, normalize_arg
 from ..blockwise import blockwise, Blockwise
 from ..base import DaskMethodsMixin, tokenize, dont_optimize, is_dask_collection
-from ..sizeof import sizeof
 from ..delayed import delayed, Delayed, unpack_collections
 from ..highlevelgraph import HighLevelGraph
 
@@ -3772,15 +3771,12 @@ def map_partitions(func, *args, **kwargs):
     name = kwargs.pop('token', None)
     transform_divisions = kwargs.pop('transform_divisions', True)
 
-    # Normalize keyword arguments
-    kwargs2 = {k: normalize_arg(v) for k, v in kwargs.items()}
-
     assert callable(func)
     if name is not None:
-        token = tokenize(meta, *args, **kwargs2)
+        token = tokenize(meta, *args, **kwargs)
     else:
         name = funcname(func)
-        token = tokenize(func, meta, *args, **kwargs2)
+        token = tokenize(func, meta, *args, **kwargs)
     name = '{0}-{1}'.format(name, token)
 
     from .multi import _maybe_align_partitions
@@ -3816,8 +3812,7 @@ def map_partitions(func, *args, **kwargs):
             args2.append(arg)
             dependencies.append(arg)
             continue
-        if not is_dask_collection(arg) and sizeof(arg) > 1e6:
-            arg = delayed(arg)
+        arg = normalize_arg(arg)
         arg2, collections = unpack_collections(arg)
         if collections:
             args2.append(arg2)
@@ -3826,7 +3821,8 @@ def map_partitions(func, *args, **kwargs):
             args2.append(arg)
 
     kwargs3 = {}
-    for k, v in kwargs2.items():
+    for k, v in kwargs.items():
+        v = normalize_arg(v)
         v, collections = unpack_collections(v)
         dependencies.extend(collections)
         kwargs3[k] = v


Test patch:
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -3544,3 +3544,15 @@ def test_replace():
     assert_eq(df.replace({1: 10, 2: 20}), ddf.replace({1: 10, 2: 20}))
     assert_eq(df.x.replace(1, 10), ddf.x.replace(1, 10))
     assert_eq(df.x.replace({1: 10, 2: 20}), ddf.x.replace({1: 10, 2: 20}))
+
+
+def test_map_partitions_delays_lists():
+    df = pd.DataFrame({'x': [1, 2, 3, 4, 5]})
+    ddf = dd.from_pandas(df, npartitions=2)
+
+    L = list(range(100))
+    out = ddf.map_partitions(lambda x, y: x + sum(y), y=L)
+    assert any(str(L) == str(v) for v in out.__dask_graph__().values())
+
+    out = ddf.map_partitions(lambda x, y: x + sum(y), L)
+    assert any(str(L) == str(v) for v in out.__dask_graph__().values())