Repository: dask/dask
Pull request #11762 (2025-02-19 12:53:35)

Ensure divisions not diverging from npartitions in Merge
Stumbled over this while investigating https://github.com/dask/dask/pull/11736

The problems I see on the branch over there are not reproducible on main but it may be related to this. Either way, I think these definitions shouldn't diverge since some optimizations might take the wrong turn


Patch:
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index 5506e5553..45e6c4667 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -58,7 +58,7 @@ import itertools
 import sys
 from collections import defaultdict
 from collections.abc import Callable, Container, Iterable, Mapping, MutableMapping
-from functools import partial
+from functools import lru_cache, partial
 from typing import Any, TypeVar, cast
 
 from dask.sizeof import sizeof
@@ -402,9 +402,10 @@ class GraphNode:
         return False
 
     def __sizeof__(self) -> int:
-        return sum(
-            sizeof(getattr(self, sl)) for sl in type(self).__slots__
-        ) + sys.getsizeof(type(self))
+        all_slots = self.get_all_slots()
+        return sum(sizeof(getattr(self, sl)) for sl in all_slots) + sys.getsizeof(
+            type(self)
+        )
 
     def substitute(
         self, subs: dict[KeyType, KeyType | GraphNode], key: KeyType | None = None
@@ -475,6 +476,16 @@ class GraphNode:
             _data_producer=any(t.data_producer for t in tasks),
         )
 
+    @classmethod
+    @lru_cache
+    def get_all_slots(cls):
+        slots = list()
+        for c in cls.mro():
+            slots.extend(getattr(c, "__slots__", ()))
+        # Interestingly, sorting this causes the nested containers to pickle
+        # more efficiently
+        return sorted(set(slots))
+
 
 _no_deps: frozenset = frozenset()
 
@@ -498,6 +509,9 @@ class Alias(GraphNode):
         self.target = target
         self._dependencies = frozenset((self.target,))
 
+    def __reduce__(self):
+        return Alias, (self.key, self.target)
+
     def copy(self):
         return Alias(self.key, self.target)
 
@@ -740,6 +754,15 @@ class Task(GraphNode):
             return self.func(*new_argspec, **kwargs)
         return self.func(*new_argspec)
 
+    def __setstate__(self, state):
+        slots = self.__class__.get_all_slots()
+        for sl, val in zip(slots, state):
+            setattr(self, sl, val)
+
+    def __getstate__(self):
+        slots = self.__class__.get_all_slots()
+        return tuple(getattr(self, sl) for sl in slots)
+
     @property
     def is_coro(self):
         if self._is_coro is None:
@@ -810,6 +833,23 @@ class NestedContainer(Task):
             **kwargs,
         )
 
+    def __getstate__(self):
+        state = super().__getstate__()
+        state = list(state)
+        slots = self.__class__.get_all_slots()
+        ix = slots.index("kwargs")
+        # The constructor as a kwarg is redundant since this is encoded in the
+        # class itself. Serializing the builtin types is not trivial
+        # This saves about 15% of overhead
+        state[ix] = state[ix].copy()
+        state[ix].pop("constructor", None)
+        return state
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        self.kwargs["constructor"] = self.__class__.constructor
+        return self
+
     def __repr__(self):
         return f"{type(self).__name__}({self.args})"
 
diff --git a/dask/dataframe/dask_expr/_merge.py b/dask/dataframe/dask_expr/_merge.py
index d6fecac69..79a0cff73 100644
--- a/dask/dataframe/dask_expr/_merge.py
+++ b/dask/dataframe/dask_expr/_merge.py
@@ -205,7 +205,7 @@ class Merge(Expr):
     def _npartitions(self):
         if self.operand("_npartitions") is not None:
             return self.operand("_npartitions")
-        return max(self.left.npartitions, self.right.npartitions)
+        return len(self._divisions()) - 1
 
     @property
     def _bcast_left(self):
@@ -273,8 +273,10 @@ class Merge(Expr):
             else:
                 _npartitions = max(self.left.npartitions, self.right.npartitions)
 
+        elif self.operand("_npartitions") is not None:
+            _npartitions = self.operand("_npartitions")
         else:
-            _npartitions = self._npartitions
+            _npartitions = max(self.left.npartitions, self.right.npartitions)
 
         return (None,) * (_npartitions + 1)
 


Test patch:
diff --git a/dask/dataframe/tests/test_multi.py b/dask/dataframe/tests/test_multi.py
index 32ba46074..bcc969968 100644
--- a/dask/dataframe/tests/test_multi.py
+++ b/dask/dataframe/tests/test_multi.py
@@ -2431,3 +2431,22 @@ def test_pairwise_merge_results_in_identical_output_df(
     ddf_pairwise = ddf_pairwise.join(dfs_to_merge, how=how)
 
     assert_eq(ddf_pairwise, ddf_loop)
+
+
+@pytest.mark.parametrize("npartitions_left", [1, 2, 3])
+@pytest.mark.parametrize("npartitions_right", [1, 2, 3])
+@pytest.mark.parametrize("how", ["left", "right", "outer"])
+def test_ensure_npartitions_properly_set(how, npartitions_left, npartitions_right):
+    ddf_right = dd.from_pandas(
+        pd.DataFrame(
+            {"A": [5, 6, 7, 8], "B": [4, 3, 2, 1]},
+            index=[0, 1, 2, 3],
+        ),
+        npartitions_right,
+    )
+    ddf_left = dd.from_pandas(pd.DataFrame(index=[0, 1, 3]), npartitions_left)
+    res = ddf_left.join(ddf_right, how=how)
+    assert res.expr._npartitions == res.npartitions
+    assert res.npartitions == len(res.divisions) - 1
+    assert len(res) == len(res.compute())
+    assert len(res) == sum(res.map_partitions(len).compute())
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index 1536a2b3f..c1711d537 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -8,6 +8,7 @@ from packaging.version import Version
 
 import dask
 from dask import delayed
+from dask._compatibility import WINDOWS
 from dask.base import collections_to_dsk, key_split, visualize_dsk
 from dask.core import get_deps
 from dask.order import _connecting_to_roots, diagnostics, ndependencies, order
@@ -254,6 +255,7 @@ def _rechunk_merge_graph():
     }
 
 
+@pytest.mark.skipif(WINDOWS, reason="Broken pipe on Windows")
 def test_visualize_int_overflow():
     pytest.importorskip("graphviz")
     pytest.importorskip("matplotlib")
diff --git a/dask/tests/test_task_spec.py b/dask/tests/test_task_spec.py
index 6cd2d215c..1b7667638 100644
--- a/dask/tests/test_task_spec.py
+++ b/dask/tests/test_task_spec.py
@@ -120,24 +120,6 @@ def use_kwargs(a, kwarg=None):
     return a + kwarg
 
 
-def test_unpickled_repr():
-    t = pickle.loads(pickle.dumps(Task("key", func, "a", "b")))
-    assert repr(t) == "<Task 'key' func('a', 'b')>"
-
-    t = pickle.loads(pickle.dumps(Task("nested", func2, t, t.ref())))
-    assert repr(t) == "<Task 'nested' func2(<Task 'key' func('a', 'b')>, Alias('key'))>"
-
-    t = pickle.loads(
-        pickle.dumps(Task("long", long_function_name_longer_even_longer, t, t.ref()))
-    )
-    assert repr(t) == "<Task 'long' long_function_name_longer_even_longer(...)>"
-
-    t = pickle.loads(
-        pickle.dumps(Task("kwarg", use_kwargs, "foo", kwarg="kwarg_value"))
-    )
-    assert repr(t) == "<Task 'kwarg' use_kwargs('foo', kwarg='kwarg_value')>"
-
-
 def test_convert_legacy_dsk():
     def func(*args):
         return "-".join(args)
@@ -265,16 +247,87 @@ class SerializeOnlyOnce:
 
 def test_pickle():
 
+    def assert_slots_equal(a, b):
+        def get_all_slots(obj):
+            slots = set()
+            for cls in obj.__class__.mro():
+                slots.update(getattr(cls, "__slots__", ()))
+            return slots
+
+        all_slots = get_all_slots(a) | get_all_slots(a)
+        assert all_slots == get_all_slots(a) == get_all_slots(a)
+        assert all(getattr(a, slot) == getattr(b, slot) for slot in all_slots)
+        assert not hasattr(a, "__dict__")
+        assert not hasattr(b, "__dict__")
+
     t1 = Task("key-1", func, "a", "b")
     t2 = Task("key-2", func, "c", "d")
 
     rtt1 = pickle.loads(pickle.dumps(t1))
+    assert repr(rtt1) == repr(t1)
     rtt2 = pickle.loads(pickle.dumps(t2))
+
+    assert_slots_equal(t1, rtt1)
+    assert_slots_equal(t2, rtt2)
     assert t1 == rtt1
     assert t1.func == rtt1.func
     assert t1.func is rtt1.func
     assert t1.func is rtt2.func
 
+    l = Tuple(t1, t2)
+    rtl = pickle.loads(pickle.dumps(l))
+    assert l == rtl
+    assert l() == rtl()
+    d = Dict(key=t1)
+    rtd = pickle.loads(pickle.dumps(d))
+    assert d == rtd
+    assert d() == rtd()
+
+
+def test_pickle_size():
+    # We will serialize many of these objects which drives both memory usage and
+    # serialization runtime performance.
+    # Reducing pickle size is beneficial but the numbers below are determined
+    # empirically
+    # Analyzing the output with pickletools.dis is useful to debug memoization
+    # and serialization by value
+
+    a = Alias("a", "b")
+    # We cannot shrink it to nothing
+    assert len(pickle.dumps(a)) < 55
+    b = Alias("b", "c")
+    # But most of it should be overhead that is memoized
+    assert len(pickle.dumps((a, b))) <= 70
+
+    # Pickle should be able to memoize this. On py3.10 that's 2 additional bytes
+    assert len(pickle.dumps((a, b, b))) <= len(pickle.dumps((a, b))) + 10
+
+    t1 = Task("key-1", func, "a", "b")
+    assert len(pickle.dumps(t1)) < 120
+
+    t2 = Task("key-2", func, TaskRef("key-1"), "c")
+    assert len(pickle.dumps(t2)) < 140
+
+    assert len(pickle.dumps((t1, t2))) < 170
+
+    l = List(t1, t2)
+    assert len(pickle.dumps(l)) <= 272
+
+    sizes = []
+    growth = []
+    inner = List(t1, t2)
+    for depth in range(20):
+        inner = List(inner, t1)
+        size = len(pickle.dumps(inner))
+        if len(sizes) > 0:
+            growth.append(size - sizes[-1][1])
+        sizes.append((depth, size))
+    growth = set(growth)
+    # If this breaks, something cannot be memoized. That's very concerning
+    assert len(growth) == 1
+    # If this goes up, that's not great but not a disaster
+    assert growth.pop() <= 32
+
 
 def test_tokenize():
     t = Task("key-1", func, "a", "b")