Repository: dask/dask
Pull request #11796 (2025-02-27 16:00:37)

Fix missing imports in array-expr
This slipped in when merging https://github.com/dask/dask/pull/11783 I'm surprised CI didn't see this (or we might've missed the failure)


Patch:
diff --git a/dask/_expr.py b/dask/_expr.py
index 6170691ee..0be7bd77b 100644
--- a/dask/_expr.py
+++ b/dask/_expr.py
@@ -46,7 +46,9 @@ class Expr:
 
     operands: list
 
-    def __new__(cls, *args, **kwargs):
+    _determ_token: str | None
+
+    def __new__(cls, *args, _determ_token=None, **kwargs):
         operands = list(args)
         for parameter in cls._parameters[len(operands) :]:
             try:
@@ -55,6 +57,8 @@ class Expr:
                 operands.append(cls._defaults[parameter])
         assert not kwargs, kwargs
         inst = object.__new__(cls)
+
+        inst._determ_token = _determ_token
         inst.operands = [_unpack_collections(o) for o in operands]
         _name = inst._name
         if _name in Expr._instances:
@@ -127,10 +131,17 @@ class Expr:
     def __dask_tokenize__(self):
         return self._name
 
+    @staticmethod
+    def _reconstruct(*args):
+        typ, *operands, token = args
+        return typ(*operands, _determ_token=token)
+
     def __reduce__(self):
         if dask.config.get("dask-expr-no-serialize", False):
             raise RuntimeError(f"Serializing a {type(self)} object")
-        return type(self), tuple(self.operands)
+        return Expr._reconstruct, tuple(
+            [type(self)] + self.operands + [self.deterministic_token]
+        )
 
     def _depth(self, cache=None):
         """Depth of the expression tree
@@ -460,9 +471,15 @@ class Expr:
     def _funcname(self) -> str:
         return funcname(type(self)).lower()
 
+    @property
+    def deterministic_token(self):
+        if not self._determ_token:
+            self._determ_token = _tokenize_deterministic(*self.operands)
+        return self._determ_token
+
     @functools.cached_property
     def _name(self) -> str:
-        return self._funcname + "-" + _tokenize_deterministic(*self.operands)
+        return self._funcname + "-" + self.deterministic_token
 
     @property
     def _meta(self):
diff --git a/dask/array/_array_expr/_blockwise.py b/dask/array/_array_expr/_blockwise.py
index 5def4f8c3..394bccfb6 100644
--- a/dask/array/_array_expr/_blockwise.py
+++ b/dask/array/_array_expr/_blockwise.py
@@ -116,6 +116,15 @@ class Blockwise(ArrayExpr):
     def dtype(self):
         return self.operand("dtype")
 
+    @property
+    def deterministic_token(self):
+        if not self._determ_token:
+            # TODO: Is there an actual need to overwrite this?
+            self._determ_token = _tokenize_deterministic(
+                self.func, self.out_ind, self.dtype, *self.args, **self.kwargs
+            )
+        return self._determ_token
+
     @cached_property
     def _name(self):
         if "name" in self._parameters and self.operand("name"):
@@ -123,9 +132,7 @@ class Blockwise(ArrayExpr):
         else:
             return (
                 f"{self.token or funcname(self.func).strip('_')}-"
-                + _tokenize_deterministic(
-                    self.func, self.out_ind, self.dtype, *self.args, **self.kwargs
-                )
+                + self.deterministic_token
             )
 
     def _layer(self):
diff --git a/dask/array/_array_expr/_creation.py b/dask/array/_array_expr/_creation.py
index 4368299be..4202bc43b 100644
--- a/dask/array/_array_expr/_creation.py
+++ b/dask/array/_array_expr/_creation.py
@@ -6,6 +6,7 @@ from functools import partial
 import numpy as np
 
 from dask._collections import new_collection
+from dask._task_spec import Task
 from dask.array._array_expr._expr import ArrayExpr
 from dask.array.chunk import arange as _arange
 from dask.array.chunk import linspace as _linspace
diff --git a/dask/array/_array_expr/_expr.py b/dask/array/_array_expr/_expr.py
index 7fc9f5e91..ea13e444b 100644
--- a/dask/array/_array_expr/_expr.py
+++ b/dask/array/_array_expr/_expr.py
@@ -11,6 +11,7 @@ import toolz
 from tlz import accumulate
 
 from dask._expr import Expr
+from dask._task_spec import Task, TaskRef
 from dask.array.chunk import getitem
 from dask.array.core import T_IntOrNaN, common_blockdim, unknown_chunk_message
 from dask.blockwise import broadcast_dimensions
diff --git a/dask/array/_array_expr/_io.py b/dask/array/_array_expr/_io.py
index 0c7995d42..186fc0692 100644
--- a/dask/array/_array_expr/_io.py
+++ b/dask/array/_array_expr/_io.py
@@ -15,7 +15,6 @@ from dask.array.core import (
     slices_from_chunks,
 )
 from dask.array.utils import meta_from_array
-from dask.tokenize import _tokenize_deterministic
 from dask.utils import SerializableLock
 
 
@@ -36,9 +35,7 @@ class FromGraph(IO):
 
     @functools.cached_property
     def _name(self):
-        return (
-            self.operand("name_prefix") + "-" + _tokenize_deterministic(*self.operands)
-        )
+        return self.operand("name_prefix") + "-" + self.deterministic_token
 
     def _layer(self):
         dsk = dict(self.operand("layer"))
diff --git a/dask/array/_array_expr/_reductions.py b/dask/array/_array_expr/_reductions.py
index 2b65f5050..faf6827ba 100644
--- a/dask/array/_array_expr/_reductions.py
+++ b/dask/array/_array_expr/_reductions.py
@@ -276,14 +276,21 @@ class PartialReduce(ArrayExpr):
         "reduced_meta": None,
     }
 
+    @property
+    def deterministic_token(self):
+        if not self._determ_token:
+            # TODO: Is there an actual need to overwrite this?
+            self._determ_token = _tokenize_deterministic(
+                self.func, self.array, self.split_every, self.keepdims, self.dtype
+            )
+        return self._determ_token
+
     @cached_property
     def _name(self):
         return (
             (self.operand("name") or funcname(self.func))
             + "-"
-            + _tokenize_deterministic(
-                self.func, self.array, self.split_every, self.keepdims, self.dtype
-            )
+            + self.deterministic_token
         )
 
     @cached_property
diff --git a/dask/dataframe/dask_expr/_expr.py b/dask/dataframe/dask_expr/_expr.py
index 02a57624e..e586973c0 100644
--- a/dask/dataframe/dask_expr/_expr.py
+++ b/dask/dataframe/dask_expr/_expr.py
@@ -48,7 +48,7 @@ from dask.dataframe.utils import (
     raise_on_meta_error,
     valid_divisions,
 )
-from dask.tokenize import _tokenize_deterministic, normalize_token
+from dask.tokenize import normalize_token
 from dask.typing import Key, no_default
 from dask.utils import (
     M,
@@ -607,7 +607,7 @@ class Blockwise(Expr):
             head = funcname(self.operation)
         else:
             head = funcname(type(self)).lower()
-        return head + "-" + _tokenize_deterministic(*self.operands)
+        return head + "-" + self.deterministic_token
 
     def _blockwise_arg(self, arg, i):
         """Return a Blockwise-task argument"""
@@ -685,7 +685,7 @@ class MapPartitions(Blockwise):
             head = self.token
         else:
             head = funcname(self.func).lower()
-        return head + "-" + _tokenize_deterministic(*self.operands)
+        return head + "-" + self.deterministic_token
 
     def _broadcast_dep(self, dep: Expr):
         # Always broadcast single-partition dependencies in MapPartitions
@@ -3101,7 +3101,7 @@ class DelayedsExpr(Expr):
 
     @functools.cached_property
     def _name(self):
-        return "delayed-container-" + _tokenize_deterministic(*self.operands)
+        return "delayed-container-" + self.deterministic_token
 
     def _layer(self) -> dict:
         dask = {}
@@ -3782,7 +3782,7 @@ class Fused(Blockwise):
 
     @functools.cached_property
     def _name(self):
-        return f"{str(self)}-{_tokenize_deterministic(*self.operands)}"
+        return f"{str(self)}-{self.deterministic_token}"
 
     def _divisions(self):
         return self.exprs[0]._divisions()
diff --git a/dask/dataframe/dask_expr/_reductions.py b/dask/dataframe/dask_expr/_reductions.py
index da9fb6a2a..53d549d32 100644
--- a/dask/dataframe/dask_expr/_reductions.py
+++ b/dask/dataframe/dask_expr/_reductions.py
@@ -40,7 +40,6 @@ from dask.dataframe.dask_expr._expr import (
 )
 from dask.dataframe.dispatch import make_meta, meta_nonempty
 from dask.dataframe.utils import is_scalar
-from dask.tokenize import _tokenize_deterministic
 from dask.typing import no_default
 from dask.utils import M, apply, funcname
 
@@ -326,7 +325,7 @@ class TreeReduce(Expr):
             name = funcname(self.combine.__self__).lower() + "-tree"
         else:
             name = funcname(self.combine)
-        return name + "-" + _tokenize_deterministic(*self.operands)
+        return name + "-" + self.deterministic_token
 
     def __dask_postcompute__(self):
         return toolz.first, ()
@@ -858,7 +857,7 @@ class CustomReduction(Reduction):
     @functools.cached_property
     def _name(self):
         name = self.operand("token") or funcname(type(self)).lower()
-        return name + "-" + _tokenize_deterministic(*self.operands)
+        return name + "-" + self.deterministic_token
 
     @classmethod
     def chunk(cls, df, **kwargs):
diff --git a/dask/dataframe/dask_expr/diagnostics/_analyze.py b/dask/dataframe/dask_expr/diagnostics/_analyze.py
index 0dfd822ca..3b315c715 100644
--- a/dask/dataframe/dask_expr/diagnostics/_analyze.py
+++ b/dask/dataframe/dask_expr/diagnostics/_analyze.py
@@ -14,7 +14,6 @@ from dask.dataframe.dask_expr.diagnostics._explain import (
 from dask.dataframe.dask_expr.io.io import FusedIO
 from dask.dataframe.utils import is_scalar
 from dask.sizeof import sizeof
-from dask.tokenize import _tokenize_deterministic
 from dask.utils import format_bytes, import_required
 
 if TYPE_CHECKING:
@@ -201,4 +200,4 @@ class Analyze(Blockwise):
 
     @functools.cached_property
     def _name(self):
-        return "analyze-" + _tokenize_deterministic(*self.operands)
+        return "analyze-" + self.deterministic_token
diff --git a/dask/dataframe/dask_expr/io/_delayed.py b/dask/dataframe/dask_expr/io/_delayed.py
index 879191bb3..e948a4a20 100644
--- a/dask/dataframe/dask_expr/io/_delayed.py
+++ b/dask/dataframe/dask_expr/io/_delayed.py
@@ -16,7 +16,6 @@ from dask.dataframe.dask_expr.io import BlockwiseIO
 from dask.dataframe.dispatch import make_meta
 from dask.dataframe.utils import check_meta, pyarrow_strings_enabled
 from dask.delayed import Delayed, delayed
-from dask.tokenize import _tokenize_deterministic
 from dask.typing import Key
 
 if TYPE_CHECKING:
@@ -44,7 +43,7 @@ class FromDelayed(PartitionsFiltered, BlockwiseIO):
     def _name(self):
         if self.prefix is None:
             return super()._name
-        return self.prefix + "-" + _tokenize_deterministic(*self.operands)
+        return self.prefix + "-" + self.deterministic_token
 
     @functools.cached_property
     def _meta(self):
diff --git a/dask/dataframe/dask_expr/io/io.py b/dask/dataframe/dask_expr/io/io.py
index ab416156f..7874346fc 100644
--- a/dask/dataframe/dask_expr/io/io.py
+++ b/dask/dataframe/dask_expr/io/io.py
@@ -107,11 +107,7 @@ class FusedIO(BlockwiseIO):
 
     @functools.cached_property
     def _name(self):
-        return (
-            self.operand("_expr")._funcname
-            + "-fused-"
-            + _tokenize_deterministic(*self.operands)
-        )
+        return self.operand("_expr")._funcname + "-fused-" + self.deterministic_token
 
     @functools.cached_property
     def _meta(self):
@@ -167,7 +163,7 @@ class FusedParquetIO(FusedIO):
         return (
             funcname(type(self.operand("_expr"))).lower()
             + "-fused-parq-"
-            + _tokenize_deterministic(*self.operands)
+            + self.deterministic_token
         )
 
     @staticmethod
@@ -245,13 +241,9 @@ class FromMap(PartitionsFiltered, BlockwiseIO):
     @functools.cached_property
     def _name(self):
         if self.label is None:
-            return (
-                funcname(self.func).lower()
-                + "-"
-                + _tokenize_deterministic(*self.operands)
-            )
+            return funcname(self.func).lower() + "-" + self.deterministic_token
         else:
-            return self.label + "-" + _tokenize_deterministic(*self.operands)
+            return self.label + "-" + self.deterministic_token
 
     @functools.cached_property
     def _meta(self):
@@ -552,7 +544,7 @@ class FromPandasDivisions(FromPandas):
 
     @functools.cached_property
     def _name(self):
-        return "from_pd_divs" + "-" + _tokenize_deterministic(*self.operands)
+        return "from_pd_divs" + "-" + self.deterministic_token
 
     @property
     def _divisions_and_locations(self):
diff --git a/dask/dataframe/dask_expr/io/parquet.py b/dask/dataframe/dask_expr/io/parquet.py
index 737ef69b6..fce3d5472 100644
--- a/dask/dataframe/dask_expr/io/parquet.py
+++ b/dask/dataframe/dask_expr/io/parquet.py
@@ -777,15 +777,18 @@ class ReadParquet(PartitionsFiltered, BlockwiseIO):
     def _funcname(self):
         return "read_parquet"
 
-    @cached_property
-    def _name(self):
-        return (
-            self._funcname
-            + "-"
-            + _tokenize_deterministic(
+    @property
+    def deterministic_token(self):
+        if not self._determ_token:
+            # TODO: Is there an actual need to overwrite this?
+            self._determ_token = _tokenize_deterministic(
                 funcname(type(self)), self.checksum, *self.operands[:-1]
             )
-        )
+        return self._determ_token
+
+    @cached_property
+    def _name(self):
+        return self._funcname + "-" + self.deterministic_token
 
     @property
     def checksum(self):


Test patch:
diff --git a/dask/dataframe/dask_expr/tests/test_core.py b/dask/dataframe/dask_expr/tests/test_core.py
index a039f369a..b564690b9 100644
--- a/dask/dataframe/dask_expr/tests/test_core.py
+++ b/dask/dataframe/dask_expr/tests/test_core.py
@@ -1,8 +1,12 @@
 from __future__ import annotations
 
+import pickle
+import uuid
+
 import pytest
 
 from dask._expr import Expr
+from dask.tokenize import tokenize
 
 
 class ExprB(Expr):
@@ -25,3 +29,21 @@ def test_endless_simplify():
     expr = ExprA()
     with pytest.raises(RuntimeError, match="converge"):
         expr.simplify()
+
+
+class NondeterminisitcToken:
+    def __dask_tokenize__(self):
+        return uuid.uuid4()
+
+
+def test_expr_nondeterministic_token_pickle_roundtrip():
+
+    foo = NondeterminisitcToken()
+    assert tokenize(foo) != tokenize(foo)
+    assert tokenize(Expr(foo)) != tokenize(Expr(foo))
+
+    inst = Expr(foo)
+    tok = tokenize(inst)
+    assert tokenize(inst) == tok
+    rt = pickle.loads(pickle.dumps(inst))
+    assert tokenize(rt) == tok
diff --git a/dask/dataframe/dask_expr/tests/test_distributed.py b/dask/dataframe/dask_expr/tests/test_distributed.py
index a229199de..4033e2012 100644
--- a/dask/dataframe/dask_expr/tests/test_distributed.py
+++ b/dask/dataframe/dask_expr/tests/test_distributed.py
@@ -283,7 +283,7 @@ async def test_merge_p2p_shuffle_projection_error(c, s, a, b):
 
 
 def test_sort_values():
-    with LocalCluster(processes=False, n_workers=2) as cluster:
+    with LocalCluster(processes=False, n_workers=2, dashboard_address=":0") as cluster:
         with Client(cluster) as client:  # noqa: F841
             pdf = pd.DataFrame({"a": [5] + list(range(100)), "b": 2})
             df = from_pandas(pdf, npartitions=10)
@@ -297,7 +297,7 @@ def test_sort_values():
 
 @pytest.mark.parametrize("add_repartition", [True, False])
 def test_merge_combine_similar_squash_merges(add_repartition):
-    with LocalCluster(processes=False, n_workers=2) as cluster:
+    with LocalCluster(processes=False, n_workers=2, dashboard_address=":0") as cluster:
         with Client(cluster) as client:  # noqa: F841
             pdf = pd.DataFrame(
                 {