Follow up to expression tokenization caching (dask#11837)

Author: fjetter

dask/array/_array_expr/_collection.py

@@ -67,6 +67,9 @@ def __dask_keys__(self):
         out = self.expr.lower_completely()
         return out.__dask_keys__()
 
+    def __dask_tokenize__(self):
+        return "Array", self.expr._name
+
     def compute(self, **kwargs):
         return DaskMethodsMixin.compute(self.optimize(), **kwargs)
 

dask/array/_array_expr/_expr.py

@@ -16,7 +16,6 @@
 from dask.array.core import T_IntOrNaN, common_blockdim, unknown_chunk_message
 from dask.blockwise import broadcast_dimensions
 from dask.layers import ArrayBlockwiseDep
-from dask.tokenize import _tokenize_deterministic
 from dask.utils import cached_cumsum
 
 
@@ -96,6 +95,9 @@ def keys(*args):
         self._cached_keys = result = keys()
         return result
 
+    def __dask_tokenize__(self):
+        return self._name
+
     def __hash__(self):
         return hash(self._name)
 
@@ -190,7 +192,7 @@ def chunks(self):
 
     @functools.cached_property
     def _name(self):
-        return "stack-" + _tokenize_deterministic(*self.operands)
+        return "stack-" + self.deterministic_token
 
     def _layer(self) -> dict:
         keys = list(product([self._name], *[range(len(bd)) for bd in self.chunks]))
@@ -239,7 +241,7 @@ def chunks(self):
 
     @functools.cached_property
     def _name(self):
-        return "stack-" + _tokenize_deterministic(*self.operands)
+        return "stack-" + self.deterministic_token
 
     def _layer(self) -> dict:
         axis = self.axis

dask/array/_array_expr/_shuffle.py

@@ -12,7 +12,6 @@
 from dask.array.chunk import getitem
 from dask.array.dispatch import concatenate_lookup, take_lookup
 from dask.base import tokenize
-from dask.tokenize import _tokenize_deterministic
 
 
 def _validate_indexer(chunks, indexer, axis):
@@ -56,10 +55,6 @@ def _meta(self):
     def _name(self):
         return f"{self.operand('name')}-{self._token}"
 
-    @functools.cached_property
-    def _token(self):
-        return _tokenize_deterministic(*self.operands)
-
     @functools.cached_property
     def chunks(self):
         output_chunks = []
@@ -118,7 +113,7 @@ def _layer(self) -> dict:
         intermediates: dict = dict()
         merges: dict = dict()
         dtype = np.min_scalar_type(max(max(chunks[axis]), self._chunk_size_limit))
-        split_name = f"shuffle-split-{self._token}"
+        split_name = f"shuffle-split-{self.deterministic_token}"
         slices = [slice(None)] * len(chunks)
         split_name_suffixes = count()
         sorter_name = "shuffle-sorter-"

dask/array/_array_expr/_slicing.py

@@ -22,7 +22,7 @@
 )
 from dask.array.utils import meta_from_array
 from dask.layers import ArrayBlockwiseDep
-from dask.tokenize import _tokenize_deterministic, tokenize
+from dask.tokenize import tokenize
 from dask.utils import cached_cumsum, is_arraylike
 
 
@@ -369,7 +369,7 @@ def take(x, index, axis=0):
 class Slice(ArrayExpr):
     @functools.cached_property
     def _name(self):
-        return f"getitem-{_tokenize_deterministic(*self.operands)}"
+        return f"getitem-{self.deterministic_token}"
 
     @functools.cached_property
     def _meta(self):

dask/dataframe/dask_expr/_expr.py

@@ -3045,7 +3045,7 @@ class _DelayedExpr(Expr):
     # TODO
     _parameters = ["obj"]
 
-    def __init__(self, obj):
+    def __init__(self, obj, _determ_token=None):
         self.obj = obj
         self.operands = [obj]
 
@@ -3073,7 +3073,7 @@ def ndim(self):
 class DelayedsExpr(Expr):
     _parameters = []
 
-    def __init__(self, *delayed_objects):
+    def __init__(self, *delayed_objects, _determ_token=None):
         self.operands = delayed_objects
 
     def __str__(self):