Ensure repartition does not trigger memory size computation during lowering (i.e. on the scheduler) (dask#11855)

Author: fjetter

dask/_expr.py

@@ -45,6 +45,7 @@ class Expr:
     _parameters: list[str] = []
     _defaults: dict[str, Any] = {}
     _instances: weakref.WeakValueDictionary[str, Expr] = weakref.WeakValueDictionary()
+    _pickle_functools_cache: bool = True
 
     operands: list
 
@@ -146,14 +147,23 @@ def __dask_tokenize__(self):
 
     @staticmethod
     def _reconstruct(*args):
-        typ, *operands, token = args
-        return typ(*operands, _determ_token=token)
+        typ, *operands, token, cache = args
+        inst = typ(*operands, _determ_token=token)
+        for k, v in cache.items():
+            inst.__dict__[k] = v
+        return inst
 
     def __reduce__(self):
         if dask.config.get("dask-expr-no-serialize", False):
             raise RuntimeError(f"Serializing a {type(self)} object")
+        cache = {}
+        if type(self)._pickle_functools_cache:
+            for k, v in type(self).__dict__.items():
+                if isinstance(v, functools.cached_property) and k in self.__dict__:
+                    cache[k] = getattr(self, k)
+
         return Expr._reconstruct, tuple(
-            [type(self), *self.operands, self.deterministic_token]
+            [type(self), *self.operands, self.deterministic_token, cache]
         )
 
     def _depth(self, cache=None):

dask/dataframe/dask_expr/_repartition.py

@@ -415,21 +415,26 @@ def _lower(self):
 
 
 class RepartitionSize(Repartition):
+
     @functools.cached_property
     def _size(self):
         size = self.operand("partition_size")
         if isinstance(size, str):
             size = parse_bytes(size)
         return int(size)
 
+    @functools.cached_property
+    def _mem_usage(self):
+        return _get_mem_usages(self.frame)
+
     @functools.cached_property
     def _nsplits(self):
-        return 1 + _get_mem_usages(self.frame) // self._size
+        return 1 + self._mem_usage // self._size
 
     @functools.cached_property
     def _partition_boundaries(self):
         nsplits = self._nsplits
-        mem_usages = _get_mem_usages(self.frame)
+        mem_usages = self._mem_usage
 
         if np.any(nsplits > 1):
             split_mem_usages = []
@@ -449,6 +454,11 @@ def _divisions(self):
             return (None,) * len(self._partition_boundaries)
         return (self.frame.divisions[i] for i in self._partition_boundaries)
 
+    def _lower(self):
+        # populate cache
+        self._mem_usage  # noqa
+        return super()._lower()
+
     def _layer(self) -> dict:
         df = self.frame
         dsk: dict[tuple, Any] = {}

dask/dataframe/dask_expr/io/parquet.py

@@ -716,6 +716,7 @@ def default_types_mapper(pyarrow_dtype):
 
 
 class ReadParquet(PartitionsFiltered, BlockwiseIO):
+    _pickle_functools_cache = False
     _absorb_projections = True
     _filter_passthrough = False
 

dask/dataframe/dask_expr/tests/test_collection.py

@@ -1370,7 +1370,7 @@ def test_column_getattr(df):
 def test_serialization(pdf, df):
     before = pickle.dumps(df)
 
-    assert len(before) < 200 + len(pickle.dumps(pdf))
+    assert len(before) < 350 + len(pickle.dumps(pdf))
 
     part = df.partitions[0].compute()
     assert (
@@ -1380,8 +1380,6 @@ def test_serialization(pdf, df):
 
     after = pickle.dumps(df)
 
-    assert before == after  # caching doesn't affect serialization
-
     assert pickle.loads(before)._name == pickle.loads(after)._name
     assert_eq(pickle.loads(before), pickle.loads(after))
 

dask/tests/test_expr.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+import functools
+import pickle
+
 import pytest
 
 from dask._expr import Expr
@@ -28,3 +31,60 @@ def test_setattr2():
     assert e.bar == 3
     with pytest.raises(AttributeError):
         e.baz = 4
+
+
+class MyExprCachedProperty(Expr):
+    called_cached_property = False
+    _parameters = ["foo", "bar"]
+
+    @property
+    def baz(self):
+        return self.foo + self.bar
+
+    @functools.cached_property
+    def cached_property(self):
+        if MyExprCachedProperty.called_cached_property:
+            raise RuntimeError("No!")
+        MyExprCachedProperty.called_cached_property = True
+        return self.foo + self.bar
+
+
+def test_pickle_cached_properties():
+    pytest.importorskip("distributed")
+    from distributed import Nanny
+    from distributed.utils_test import gen_cluster
+
+    @gen_cluster(client=True, Worker=Nanny, nthreads=[("", 1)])
+    async def test(c, s, a):
+
+        expr = MyExprCachedProperty(foo=1, bar=2)
+        for _ in range(10):
+            assert expr.baz == 3
+            assert expr.cached_property == 3
+
+        assert MyExprCachedProperty.called_cached_property is True
+
+        rt = pickle.loads(pickle.dumps(expr))
+        assert rt.cached_property == 3
+        assert MyExprCachedProperty.called_cached_property is True
+
+        # Expressions are singletons, i.e. this doesn't crash
+        expr2 = MyExprCachedProperty(foo=1, bar=2)
+        assert expr2.cached_property == 3
+
+        # But this does
+        expr3 = MyExprCachedProperty(foo=1, bar=3)
+        with pytest.raises(RuntimeError):
+            expr3.cached_property
+
+        def f(expr):
+            # We want the cache to be part of the pickle, i.e. this is a
+            # different process such that the type is reset and the property can
+            # be accessed without side effects
+            assert MyExprCachedProperty.called_cached_property is False
+            assert expr.cached_property == 3
+            assert MyExprCachedProperty.called_cached_property is False
+
+        await c.submit(f, expr)
+
+    test()