Repository: dask-contrib/dask-expr
Pull request #938 (2024-03-05 15:24:49)

Avoid fusing from_pandas ops to avoid duplicating data
fusing from_pandas expressions is problematic, in broadcasting expressions we fuse them into every partition, duplicating the data times the number of partitions. This created a 1.4GB task graph in our benchmarks compared to 67MB previously.

This gets us back to that number


Patch:
diff --git a/dask_expr/_expr.py b/dask_expr/_expr.py
index daccff0..5324eb5 100644
--- a/dask_expr/_expr.py
+++ b/dask_expr/_expr.py
@@ -2925,6 +2925,10 @@ def are_co_aligned(*exprs):
 ## Utilites for Expr fusion
 
 
+def is_valid_blockwise_op(expr):
+    return isinstance(expr, Blockwise) and not isinstance(expr, (FromPandas, FromArray))
+
+
 def optimize_blockwise_fusion(expr):
     """Traverse the expression graph and apply fusion"""
 
@@ -2943,7 +2947,7 @@ def optimize_blockwise_fusion(expr):
                 continue
             seen.add(next._name)
 
-            if isinstance(next, Blockwise):
+            if is_valid_blockwise_op(next):
                 dependencies[next._name] = set()
                 if next._name not in dependents:
                     dependents[next._name] = set()
@@ -2952,7 +2956,7 @@ def optimize_blockwise_fusion(expr):
             for operand in next.operands:
                 if isinstance(operand, Expr):
                     stack.append(operand)
-                    if isinstance(operand, Blockwise):
+                    if is_valid_blockwise_op(operand):
                         if next._name in dependencies:
                             dependencies[next._name].add(operand._name)
                         dependents[operand._name].add(next._name)
@@ -2966,7 +2970,7 @@ def optimize_blockwise_fusion(expr):
             expr_mapping[k]
             for k, v in dependents.items()
             if v == set()
-            or all(not isinstance(expr_mapping[_expr], Blockwise) for _expr in v)
+            or all(not is_valid_blockwise_op(expr_mapping[_expr]) for _expr in v)
         ]
         while roots:
             root = roots.pop()
@@ -3856,4 +3860,4 @@ from dask_expr._reductions import (
     TreeReduce,
     Var,
 )
-from dask_expr.io import IO, BlockwiseIO
+from dask_expr.io import IO, BlockwiseIO, FromArray, FromPandas


Test patch:
diff --git a/dask_expr/io/tests/test_io.py b/dask_expr/io/tests/test_io.py
index ccf0cd0..459bb4b 100644
--- a/dask_expr/io/tests/test_io.py
+++ b/dask_expr/io/tests/test_io.py
@@ -185,7 +185,8 @@ def test_io_culling(tmpdir, fmt):
 
     # All tasks should be fused for the single output partition
     assert df2.npartitions == 1
-    assert len(df2.dask) == df2.npartitions
+    # from pandas is not fused together
+    assert len(df2.dask) == df2.npartitions + (1 if fmt == "pandas" else 0)
     expected = pdf.iloc[5:][["a", "b"]] + 1
     assert_eq(df2, expected, check_index=False)
 
diff --git a/dask_expr/tests/test_collection.py b/dask_expr/tests/test_collection.py
index dcd41ae..9506384 100644
--- a/dask_expr/tests/test_collection.py
+++ b/dask_expr/tests/test_collection.py
@@ -119,6 +119,14 @@ def test_info(df, verbose, buf, memory_usage):
         raise NotImplementedError(f"Case not covered for kwargs: {kwargs}")
 
 
+def test_dont_fuse_from_pandas_ops(df, pdf):
+    df2 = from_pandas(pdf, npartitions=1)
+    df = df.astype("float64")
+    result = df.merge(df2)
+    dsk = result.optimize().dask
+    assert (df2._name, 0) in dsk
+
+
 def test_column_projection_modify_list(df, pdf):
     cols = ["x"]
     result = df[cols]
diff --git a/dask_expr/tests/test_fusion.py b/dask_expr/tests/test_fusion.py
index 65eeda0..6546c43 100644
--- a/dask_expr/tests/test_fusion.py
+++ b/dask_expr/tests/test_fusion.py
@@ -25,7 +25,8 @@ def test_simple(df):
     fused = optimize(out, fuse=True)
 
     # Should only get one task per partition
-    assert len(fused.dask) == df.npartitions
+    # from_pandas is not fused together
+    assert len(fused.dask) == df.npartitions + 10
     assert_eq(fused, unfused)
 
 
@@ -63,9 +64,9 @@ def test_optimize_fusion_many():
     result = (series_a + series_b) + 1
     fused = optimize(result, fuse=True)
     unfused = optimize(result, fuse=False)
-    unfused.pprint()
     assert fused.npartitions == a.npartitions
-    assert len(fused.dask) == fused.npartitions
+    # from_pandas is not fused together
+    assert len(fused.dask) == fused.npartitions + 20
     assert_eq(fused, unfused)
 
 
@@ -83,7 +84,8 @@ def test_optimize_fusion_repeat(df):
     # repeatedly call optimize after doing new fusable things
     fused = optimize(optimize(optimize(df) + 2).x)
 
-    assert len(fused.dask) == fused.npartitions == original.npartitions
+    # from_pandas is not fused together
+    assert len(fused.dask) - 10 == fused.npartitions == original.npartitions
     assert_eq(fused, df.x + 2)
 
 
@@ -115,13 +117,14 @@ def test_fuse_broadcast_deps():
     df3 = from_pandas(pdf3, npartitions=2)
 
     query = df.merge(df2).merge(df3)
-    assert len(query.optimize().__dask_graph__()) == 2
+    # from_pandas is not fused together
+    assert len(query.optimize().__dask_graph__()) == 2 + 4
     assert_eq(query, pdf.merge(pdf2).merge(pdf3))
 
 
 def test_name(df):
     out = (df["x"] + df["y"]) - 1
     fused = optimize(out, fuse=True)
-    assert "pandas" in str(fused.expr)
+    assert "getitem" in str(fused.expr)
     assert "sub" in str(fused.expr)
     assert str(fused.expr) == str(fused.expr).lower()
diff --git a/dask_expr/tests/test_merge.py b/dask_expr/tests/test_merge.py
index d228df3..13ae9c9 100644
--- a/dask_expr/tests/test_merge.py
+++ b/dask_expr/tests/test_merge.py
@@ -233,7 +233,7 @@ def test_merge_combine_similar(npartitions_left, npartitions_right):
     query["new"] = query.b + query.c
     query = query.groupby(["a", "e", "x"]).new.sum()
     assert (
-        len(query.optimize().__dask_graph__()) <= 25
+        len(query.optimize().__dask_graph__()) <= 30
     )  # 45 is the non-combined version
 
     expected = pdf.merge(pdf2)