Repository: dask/dask
Pull request #11855 (2025-03-31 15:51:56)

Ensure repartition does not trigger memory size computation during lowering (i.e. on the scheduler)
This _should_ fix https://github.com/coiled/benchmarks/issues/1686

I noticed that repartitioning is only collecting the memory size computation during materialization, i.e. the scheduler launched a "threaded" scheduler and started to compute the memory size... that's not intended! 

This change restores the stdlib behavior of how to deal with cached properties. Historically, I introduced dedicated properties to store this cache but the functools property cache may be a more ergonomic option. This, however, makes it harder to opt out of caching a property. 

fjetter: Note: I tried the copilot review. This is manual opt-in, i.e. doens't happen automatically
github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

      9 files  ±0        9 suites  ±0   3h 29m 51s ⏱️ +3s
 17 820 tests +1   16 602 ✅  -  2   1 215 💤 ±0   3 ❌ + 3 
159 454 runs  +9  147 319 ✅  - 12  12 115 💤 +1  20 ❌ +20 

For more details on these failures, see [this check](https://github.com/dask/dask/runs/39712989352).

Results for commit 5e92a3c7. ± Comparison against base commit 4399aef0.

[test-results]:data:application/gzip;base64,H4sIAD3E6mcC/1XMTQ7CIBCG4as0rF0UmAHHyxhKISH2x1BYNd5dSlqMy++ZzLszHya3sUdHt45tOaQ2xhxNCutSJhdIvFA5puPM9V30135u2doDlerFD1/hXT85NvMmTMVkAxfjGoscrZiXmkYChHO3NGjJqWFL13a1My2u0l/arvMcUhkMHQkjrYbBOPKorTQWNYID8EqSQm69GQfBPl8Pfku1GQEAAA==


Patch:
diff --git a/dask/_expr.py b/dask/_expr.py
index 4e3d7ed5d..67e61b807 100644
--- a/dask/_expr.py
+++ b/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 @@ class Expr:
 
     @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):
diff --git a/dask/dataframe/dask_expr/_repartition.py b/dask/dataframe/dask_expr/_repartition.py
index d5a9dd5d7..07eb67a59 100644
--- a/dask/dataframe/dask_expr/_repartition.py
+++ b/dask/dataframe/dask_expr/_repartition.py
@@ -415,6 +415,7 @@ class RepartitionFreq(Repartition):
 
 
 class RepartitionSize(Repartition):
+
     @functools.cached_property
     def _size(self):
         size = self.operand("partition_size")
@@ -422,14 +423,18 @@ class RepartitionSize(Repartition):
             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 @@ class RepartitionSize(Repartition):
             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] = {}
diff --git a/dask/dataframe/dask_expr/_rolling.py b/dask/dataframe/dask_expr/_rolling.py
index fd72ba447..2bf4f9cf5 100644
--- a/dask/dataframe/dask_expr/_rolling.py
+++ b/dask/dataframe/dask_expr/_rolling.py
@@ -287,56 +287,56 @@ class Rolling:
         )
 
     @derived_from(pd_Rolling)
-    def cov(self):
-        return self._single_agg(RollingCov)
+    def cov(self, *args, **kwargs):
+        return self._single_agg(RollingCov, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
     def apply(self, func, *args, **kwargs):
         return self._single_agg(RollingApply, how_args=(func, *args), how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def count(self):
-        return self._single_agg(RollingCount)
+    def count(self, *args, **kwargs):
+        return self._single_agg(RollingCount, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def sum(self):
-        return self._single_agg(RollingSum)
+    def sum(self, *args, **kwargs):
+        return self._single_agg(RollingSum, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def mean(self):
-        return self._single_agg(RollingMean)
+    def mean(self, *args, **kwargs):
+        return self._single_agg(RollingMean, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def min(self):
-        return self._single_agg(RollingMin)
+    def min(self, *args, **kwargs):
+        return self._single_agg(RollingMin, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def max(self):
-        return self._single_agg(RollingMax)
+    def max(self, *args, **kwargs):
+        return self._single_agg(RollingMax, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def var(self):
-        return self._single_agg(RollingVar)
+    def var(self, *args, **kwargs):
+        return self._single_agg(RollingVar, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def std(self):
-        return self._single_agg(RollingStd)
+    def std(self, *args, **kwargs):
+        return self._single_agg(RollingStd, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def median(self):
-        return self._single_agg(RollingMedian)
+    def median(self, *args, **kwargs):
+        return self._single_agg(RollingMedian, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def quantile(self, q):
-        return self._single_agg(RollingQuantile, how_args=(q,))
+    def quantile(self, q, *args, **kwargs):
+        return self._single_agg(RollingQuantile, how_args=(q, *args), how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def skew(self):
-        return self._single_agg(RollingSkew)
+    def skew(self, *args, **kwargs):
+        return self._single_agg(RollingSkew, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
-    def kurt(self):
-        return self._single_agg(RollingKurt)
+    def kurt(self, *args, **kwargs):
+        return self._single_agg(RollingKurt, how_args=args, how_kwargs=kwargs)
 
     @derived_from(pd_Rolling)
     def agg(self, func, *args, **kwargs):
diff --git a/dask/dataframe/dask_expr/io/parquet.py b/dask/dataframe/dask_expr/io/parquet.py
index 3f3c9ff5c..bf7bfe372 100644
--- a/dask/dataframe/dask_expr/io/parquet.py
+++ b/dask/dataframe/dask_expr/io/parquet.py
@@ -716,6 +716,7 @@ def _determine_type_mapper(
 
 
 class ReadParquet(PartitionsFiltered, BlockwiseIO):
+    _pickle_functools_cache = False
     _absorb_projections = True
     _filter_passthrough = False
 


Test patch:
diff --git a/dask/dataframe/dask_expr/tests/test_collection.py b/dask/dataframe/dask_expr/tests/test_collection.py
index bbd7fb43f..1f1754199 100644
--- a/dask/dataframe/dask_expr/tests/test_collection.py
+++ b/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))
 
diff --git a/dask/dataframe/dask_expr/tests/test_rolling.py b/dask/dataframe/dask_expr/tests/test_rolling.py
index 35026d0d9..a1384ce61 100644
--- a/dask/dataframe/dask_expr/tests/test_rolling.py
+++ b/dask/dataframe/dask_expr/tests/test_rolling.py
@@ -14,7 +14,10 @@ pd = _backend_library()
 @pytest.fixture
 def pdf():
     idx = pd.date_range("2000-01-01", periods=12, freq="min")
-    pdf = pd.DataFrame({"foo": range(len(idx))}, index=idx)
+    pdf = pd.DataFrame(
+        {"foo": range(len(idx)), "bar": idx},
+        index=idx,
+    )
     pdf["bar"] = 1
     yield pdf
 
@@ -34,30 +37,44 @@ def df(pdf, request):
         ("min", ()),
         ("max", ()),
         ("var", ()),
+        ("var", (2,)),  # ddof
         ("std", ()),
+        ("std", (2,)),  # ddof
         ("median", ()),
         ("skew", ()),
         ("quantile", (0.5,)),
         ("kurt", ()),
     ],
 )
+@pytest.mark.parametrize("numeric_only", [True, False])
 @pytest.mark.parametrize("window,min_periods", ((1, None), (3, 2), (3, 3)))
 @pytest.mark.parametrize("center", (True, False))
 @pytest.mark.parametrize("df", (1, 2), indirect=True)
-def test_rolling_apis(df, pdf, window, api, how_args, min_periods, center):
+def test_rolling_apis(
+    df, pdf, window, api, how_args, min_periods, center, numeric_only
+):
     args = (window,)
-    kwargs = dict(min_periods=min_periods, center=center)
-
-    result = getattr(df.rolling(*args, **kwargs), api)(*how_args)
-    expected = getattr(pdf.rolling(*args, **kwargs), api)(*how_args)
+    kwargs = dict(
+        min_periods=min_periods,
+        center=center,
+    )
+    how_kwargs = dict(
+        numeric_only=numeric_only,
+    )
+    result = getattr(df.rolling(*args, **kwargs), api)(*how_args, **how_kwargs)
+    expected = getattr(pdf.rolling(*args, **kwargs), api)(*how_args, **how_kwargs)
     assert_eq(result, expected)
 
-    result = getattr(df.rolling(*args, **kwargs), api)(*how_args)["foo"]
-    expected = getattr(pdf.rolling(*args, **kwargs), api)(*how_args)["foo"]
+    result = getattr(df.rolling(*args, **kwargs), api)(*how_args, **how_kwargs)["foo"]
+    expected = getattr(pdf.rolling(*args, **kwargs), api)(*how_args, **how_kwargs)[
+        "foo"
+    ]
     assert_eq(result, expected)
 
     q = result.simplify()
-    eq = getattr(df["foo"].rolling(*args, **kwargs), api)(*how_args).simplify()
+    eq = getattr(df["foo"].rolling(*args, **kwargs), api)(
+        *how_args, **how_kwargs
+    ).simplify()
     assert q._name == eq._name
 
 
diff --git a/dask/tests/test_expr.py b/dask/tests/test_expr.py
index 2907e5e5a..c40a64513 100644
--- a/dask/tests/test_expr.py
+++ b/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()