Repository: dask/dask
Pull request #11849 (2025-03-25 09:57:12)

Ensure HLGExpr tokenize uniquely
Closes https://github.com/pydata/xarray/issues/10171

I've seen these kind of failures during development. I don't fully understand why dask CI was fine with the current code but it might've been related to serialization. This fixes the xarray issue

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 26m 59s ⏱️ - 2m 18s
 17 820 tests +1   16 433 ✅  -   171   1 215 💤 ± 0    172 ❌ +  172 
159 454 runs  +9  145 838 ✅  - 1 493  12 102 💤  - 12  1 514 ❌ +1 514 

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

Results for commit d8e07abc. ± Comparison against base commit e0877d0a.

[test-results]:data:application/gzip;base64,H4sIAPGI4mcC/02MOw7DIBAFr2JRp2AxC0suE2GwJRR/IgyVlbsHrJiknHl6c7ApzOPO7p25dWzPITXwOdoUtrUgCAlVlTHVGTQJfvFjz85VqWTf/+QzvM4nYHOTDfP5Fk2NMW6xuFqLeT3jaCTKL7e4ROqpyRbn4nJXHKF9/+NuW5aQCjBPI9d2cEpzUjRo44y3ElCT9B6EJa+0RRDs/QGXZjf7HQEAAA==


Patch:
diff --git a/dask/_expr.py b/dask/_expr.py
index dfb1d9093..4e3d7ed5d 100644
--- a/dask/_expr.py
+++ b/dask/_expr.py
@@ -136,7 +136,13 @@ class Expr:
         return hash(self._name)
 
     def __dask_tokenize__(self):
-        return self._name
+        if not self._determ_token:
+            # If the subclass does not implement a __dask_tokenize__ we'll want
+            # to tokenize all operands.
+            # Note how this differs to the implementation of
+            # Expr.deterministic_token
+            self._determ_token = _tokenize_deterministic(type(self), *self.operands)
+        return self._determ_token
 
     @staticmethod
     def _reconstruct(*args):
@@ -494,7 +500,9 @@ class Expr:
     @property
     def deterministic_token(self):
         if not self._determ_token:
-            self._determ_token = _tokenize_deterministic(*self.operands)
+            # Just tokenize self to fall back on __dask_tokenize__
+            # Note how this differs to the implementation of __dask_tokenize__
+            self._determ_token = self.__dask_tokenize__()
         return self._determ_token
 
     @functools.cached_property
@@ -1074,6 +1082,11 @@ class _ExprSequence(Expr):
             all_keys.append(op.__dask_keys__())
         return all_keys
 
+    def __repr__(self):
+        return "ExprSequence(" + ", ".join(map(repr, self.operands)) + ")"
+
+    __str__ = __repr__
+
     def finalize_compute(self):
         return _ExprSequence(
             *(op.finalize_compute() for op in self.operands),
diff --git a/dask/array/_array_expr/_blockwise.py b/dask/array/_array_expr/_blockwise.py
index 394bccfb6..b123494cd 100644
--- a/dask/array/_array_expr/_blockwise.py
+++ b/dask/array/_array_expr/_blockwise.py
@@ -116,8 +116,7 @@ class Blockwise(ArrayExpr):
     def dtype(self):
         return self.operand("dtype")
 
-    @property
-    def deterministic_token(self):
+    def __dask_tokenize__(self):
         if not self._determ_token:
             # TODO: Is there an actual need to overwrite this?
             self._determ_token = _tokenize_deterministic(
diff --git a/dask/array/_array_expr/_expr.py b/dask/array/_array_expr/_expr.py
index b727e763b..1357ef0c8 100644
--- a/dask/array/_array_expr/_expr.py
+++ b/dask/array/_array_expr/_expr.py
@@ -104,9 +104,6 @@ class ArrayExpr(Expr):
 
         return unwrap(key_refs)
 
-    def __dask_tokenize__(self):
-        return self._name
-
     def __hash__(self):
         return hash(self._name)
 
diff --git a/dask/array/_array_expr/_reductions.py b/dask/array/_array_expr/_reductions.py
index faf6827ba..44ff47161 100644
--- a/dask/array/_array_expr/_reductions.py
+++ b/dask/array/_array_expr/_reductions.py
@@ -276,8 +276,7 @@ class PartialReduce(ArrayExpr):
         "reduced_meta": None,
     }
 
-    @property
-    def deterministic_token(self):
+    def __dask_tokenize__(self):
         if not self._determ_token:
             # TODO: Is there an actual need to overwrite this?
             self._determ_token = _tokenize_deterministic(
diff --git a/dask/dataframe/dask_expr/_expr.py b/dask/dataframe/dask_expr/_expr.py
index 408bab56b..6adf948f9 100644
--- a/dask/dataframe/dask_expr/_expr.py
+++ b/dask/dataframe/dask_expr/_expr.py
@@ -49,7 +49,6 @@ from dask.dataframe.utils import (
     raise_on_meta_error,
     valid_divisions,
 )
-from dask.tokenize import normalize_token
 from dask.typing import Key, no_default
 from dask.utils import (
     M,
@@ -3112,11 +3111,6 @@ class DelayedsExpr(Expr):
         return 0
 
 
-@normalize_token.register(Expr)
-def normalize_expression(expr):
-    return expr._name
-
-
 def is_broadcastable(dfs, s):
     """
     This Series is broadcastable against another dataframe in the sequence
diff --git a/dask/dataframe/dask_expr/io/parquet.py b/dask/dataframe/dask_expr/io/parquet.py
index 8e12e9abd..3f3c9ff5c 100644
--- a/dask/dataframe/dask_expr/io/parquet.py
+++ b/dask/dataframe/dask_expr/io/parquet.py
@@ -776,8 +776,7 @@ class ReadParquet(PartitionsFiltered, BlockwiseIO):
     def _funcname(self):
         return "read_parquet"
 
-    @property
-    def deterministic_token(self):
+    def __dask_tokenize__(self):
         if not self._determ_token:
             # TODO: Is there an actual need to overwrite this?
             self._determ_token = _tokenize_deterministic(


Test patch:
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
index 76df1b678..6d3ecced6 100644
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -19,7 +19,6 @@ import dask.bag as db
 from dask.base import (
     DaskMethodsMixin,
     clone_key,
-    collections_to_expr,
     compute,
     compute_as_if_collection,
     get_collection_names,
@@ -931,16 +930,6 @@ def test_num_workers_config(scheduler):
     assert len(workers) == num_workers
 
 
-def test_optimizations_ctd():
-    pytest.importorskip("numpy")
-    da = pytest.importorskip("dask.array")
-    x = da.arange(2, chunks=1)[:1]
-    dsk1 = collections_to_expr([x])
-    with dask.config.set({"optimizations": [lambda dsk, keys: dsk]}):
-        dsk2 = collections_to_expr([x])
-    assert dsk1 == dsk2
-
-
 def test_clone_key():
     for key, seed in [("x", 123), (("x", 1), 456), (("sum-1-2-3", h1, 1), 123)]:
         validate_key(clone_key(key, seed))
diff --git a/dask/tests/test_hlgexpr.py b/dask/tests/test_hlgexpr.py
new file mode 100644
index 000000000..391d83d83
--- /dev/null
+++ b/dask/tests/test_hlgexpr.py
@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import pickle
+
+from dask._expr import HLGExpr
+from dask.tokenize import tokenize
+
+
+def test_tokenize():
+    # Ensure tokens are different for different high-level graphs The current
+    # implementation actually ensures that no HLGExpr are tokenizing equally.
+    # Technically, we do not need such a strong guarantee. but tokenizing a full
+    # HLG reliably is tricky and we do not require the reproducibility for
+    # HLGExpr since they do not undergo the same kind of optimization as the
+    # rest of the graph.
+    from dask.highlevelgraph import HighLevelGraph
+
+    dsk = HighLevelGraph.from_collections("x", {"foo": None})
+    dsk2 = HighLevelGraph.from_collections("x", {"bar": None})
+    dsk3 = HighLevelGraph.from_collections("y", {"foo": None})
+    assert tokenize(HLGExpr(dsk)) != tokenize(HLGExpr(dsk2))
+    assert tokenize(HLGExpr(dsk)) != tokenize(HLGExpr(dsk3))
+    assert tokenize(HLGExpr(dsk2)) != tokenize(HLGExpr(dsk3))
+
+    # Roundtrip preserves the tokens
+    for expr in [HLGExpr(dsk), HLGExpr(dsk2), HLGExpr(dsk3)]:
+        assert tokenize(pickle.loads(pickle.dumps(expr))) == tokenize(expr)