Ensure HLGExpr tokenize uniquely (dask#11849)

Author: fjetter

dask/_expr.py

@@ -136,7 +136,13 @@ def __hash__(self):
         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 @@ def _funcname(self) -> str:
     @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 @@ def __dask_keys__(self) -> list:
             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),

dask/array/_array_expr/_blockwise.py

@@ -116,8 +116,7 @@ def chunks(self):
     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(

dask/array/_array_expr/_expr.py

@@ -104,9 +104,6 @@ def unwrap(task):
 
         return unwrap(key_refs)
 
-    def __dask_tokenize__(self):
-        return self._name
-
     def __hash__(self):
         return hash(self._name)
 

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(

dask/dataframe/dask_expr/_expr.py

@@ -49,7 +49,6 @@
     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 @@ def ndim(self):
         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

dask/dataframe/dask_expr/io/parquet.py

@@ -776,8 +776,7 @@ def columns(self):
     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(

dask/tests/test_base.py

@@ -19,7 +19,6 @@
 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))

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)