Repository: dask/dask
Pull request #11836 (2025-03-21 09:24:15)

Expr setattr
I believe this is an important first step towards making the Expr classes more traditional classes. The concept with _parameters, etc. has been causing issues in the past in the context of inheritance (not to mention the confusion around recursion errors and other things)

This change also further highlights a couple of anti patterns where some implementations set attributes on the classes without them being registered as parameters which poses challenges to serialization and tokenization

This needs https://github.com/dask/dask/pull/11835 since otherwise I cannot properly subclass `Expr` for the test


Patch:
diff --git a/dask/_expr.py b/dask/_expr.py
index 2865fd247..640f4b679 100644
--- a/dask/_expr.py
+++ b/dask/_expr.py
@@ -164,6 +164,19 @@ class Expr:
                     cache[expr._name] = result[-1]
             return max(result)
 
+    def __setattr__(self, name: str, value: Any) -> None:
+        if name in ["operands", "_determ_token"]:
+            object.__setattr__(self, name, value)
+            return
+        try:
+            params = object.__getattribute__(type(self), "_parameters")
+            operands = object.__getattribute__(self, "operands")
+            operands[params.index(name)] = value
+        except ValueError:
+            raise AttributeError(
+                f"{type(self).__name__} object has no attribute {name}"
+            )
+
     def operand(self, key):
         # Access an operand unambiguously
         # (e.g. if the key is reserved by a method/property)
diff --git a/dask/array/_array_expr/_collection.py b/dask/array/_array_expr/_collection.py
index 3655070bf..99e1e476f 100644
--- a/dask/array/_array_expr/_collection.py
+++ b/dask/array/_array_expr/_collection.py
@@ -67,6 +67,9 @@ class Array(DaskMethodsMixin):
         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)
 
diff --git a/dask/array/_array_expr/_expr.py b/dask/array/_array_expr/_expr.py
index 5c1de318c..52810f06a 100644
--- a/dask/array/_array_expr/_expr.py
+++ b/dask/array/_array_expr/_expr.py
@@ -11,17 +11,15 @@ import toolz
 from tlz import accumulate
 
 from dask._expr import Expr
-from dask._task_spec import Task, TaskRef
+from dask._task_spec import List, 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
 from dask.layers import ArrayBlockwiseDep
-from dask.tokenize import _tokenize_deterministic
 from dask.utils import cached_cumsum
 
 
 class ArrayExpr(Expr):
-    _cached_keys = None
 
     def _operands_for_repr(self):
         return []
@@ -76,25 +74,38 @@ class ArrayExpr(Expr):
             raise ValueError(msg)
         return int(sum(self.chunks[0]))
 
-    def __dask_keys__(self):
+    @functools.cached_property
+    def _cached_keys(self):
         out = self.lower_completely()
-        if self._cached_keys is not None:
-            return self._cached_keys
 
         name, chunks, numblocks = out.name, out.chunks, out.numblocks
 
         def keys(*args):
             if not chunks:
-                return [(name,)]
+                return List(TaskRef((name,)))
             ind = len(args)
             if ind + 1 == len(numblocks):
-                result = [(name,) + args + (i,) for i in range(numblocks[ind])]
+                result = List(
+                    *(TaskRef((name,) + args + (i,)) for i in range(numblocks[ind]))
+                )
             else:
-                result = [keys(*(args + (i,))) for i in range(numblocks[ind])]
+                result = List(*(keys(*(args + (i,))) for i in range(numblocks[ind])))
             return result
 
-        self._cached_keys = result = keys()
-        return result
+        return keys()
+
+    def __dask_keys__(self):
+        key_refs = self._cached_keys
+
+        def unwrap(task):
+            if isinstance(task, List):
+                return [unwrap(t) for t in task.args]
+            return task.key
+
+        return unwrap(key_refs)
+
+    def __dask_tokenize__(self):
+        return self._name
 
     def __hash__(self):
         return hash(self._name)
@@ -190,7 +201,7 @@ class Stack(ArrayExpr):
 
     @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 +250,7 @@ class Concatenate(ArrayExpr):
 
     @functools.cached_property
     def _name(self):
-        return "stack-" + _tokenize_deterministic(*self.operands)
+        return "stack-" + self.deterministic_token
 
     def _layer(self) -> dict:
         axis = self.axis
diff --git a/dask/array/_array_expr/_shuffle.py b/dask/array/_array_expr/_shuffle.py
index 6c7bc4d29..1a137194c 100644
--- a/dask/array/_array_expr/_shuffle.py
+++ b/dask/array/_array_expr/_shuffle.py
@@ -12,7 +12,6 @@ from dask.array._array_expr._expr import ArrayExpr
 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 @@ class Shuffle(ArrayExpr):
     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 @@ class Shuffle(ArrayExpr):
         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-"
diff --git a/dask/array/_array_expr/_slicing.py b/dask/array/_array_expr/_slicing.py
index 86afb3a04..22cefd793 100644
--- a/dask/array/_array_expr/_slicing.py
+++ b/dask/array/_array_expr/_slicing.py
@@ -22,7 +22,7 @@ from dask.array.slicing import (
 )
 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):
diff --git a/dask/dataframe/dask_expr/_expr.py b/dask/dataframe/dask_expr/_expr.py
index 846cb9219..dd67b09d9 100644
--- a/dask/dataframe/dask_expr/_expr.py
+++ b/dask/dataframe/dask_expr/_expr.py
@@ -438,8 +438,7 @@ class Expr(core.Expr):
     @property
     def npartitions(self):
         if "npartitions" in self._parameters:
-            idx = self._parameters.index("npartitions")
-            return self.operands[idx]
+            return self.operand("npartitions")
         else:
             return len(self.divisions) - 1
 
@@ -3045,7 +3044,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 +3072,7 @@ class _DelayedExpr(Expr):
 class DelayedsExpr(Expr):
     _parameters = []
 
-    def __init__(self, *delayed_objects):
+    def __init__(self, *delayed_objects, _determ_token=None):
         self.operands = delayed_objects
 
     def __str__(self):
diff --git a/dask/dataframe/dask_expr/io/io.py b/dask/dataframe/dask_expr/io/io.py
index 7874346fc..0578dd9a9 100644
--- a/dask/dataframe/dask_expr/io/io.py
+++ b/dask/dataframe/dask_expr/io/io.py
@@ -398,6 +398,7 @@ class FromPandas(PartitionsFiltered, BlockwiseIO):
         "pyarrow_strings_enabled",
         "_partitions",
         "_series",
+        "_pd_length_stats",
     ]
     _defaults = {
         "npartitions": None,
@@ -407,8 +408,9 @@ class FromPandas(PartitionsFiltered, BlockwiseIO):
         "_series": False,
         "chunksize": None,
         "pyarrow_strings_enabled": True,
+        "_pd_length_stats": None,
     }
-    _pd_length_stats = None
+    _pd_length_stats: tuple | None
     _absorb_projections = True
 
     @functools.cached_property
@@ -538,8 +540,14 @@ class FromPandasDivisions(FromPandas):
         "pyarrow_strings_enabled",
         "_partitions",
         "_series",
+        "_pd_length_stats",
     ]
-    _defaults = {"columns": None, "_partitions": None, "_series": False}
+    _defaults = {
+        "columns": None,
+        "_partitions": None,
+        "_series": False,
+        "_pd_length_stats": None,
+    }
     sort = True
 
     @functools.cached_property
diff --git a/dask/dataframe/dask_expr/io/parquet.py b/dask/dataframe/dask_expr/io/parquet.py
index 3a98ebe52..74ba74f80 100644
--- a/dask/dataframe/dask_expr/io/parquet.py
+++ b/dask/dataframe/dask_expr/io/parquet.py
@@ -717,7 +717,6 @@ def _determine_type_mapper(
 
 
 class ReadParquet(PartitionsFiltered, BlockwiseIO):
-    _pq_length_stats = None
     _absorb_projections = True
     _filter_passthrough = False
 
@@ -1074,9 +1073,7 @@ class ReadParquetPyarrowFS(ReadParquet):
 
         dataset_info["schema"] = dataset.schema
         dataset_info["base_meta"] = dataset.schema.empty_table().to_pandas()
-        self.operands[type(self)._parameters.index("_dataset_info_cache")] = (
-            dataset_info
-        )
+        self._dataset_info_cache = dataset_info
         return dataset_info
 
     @cached_property
@@ -1279,6 +1276,7 @@ class ReadParquetFSSpec(ReadParquet):
         "_partitions",
         "_series",
         "_dataset_info_cache",
+        "_pq_length_stats",
     ]
     _defaults = {
         "columns": None,
@@ -1299,6 +1297,7 @@ class ReadParquetFSSpec(ReadParquet):
         "_partitions": None,
         "_series": False,
         "_dataset_info_cache": None,
+        "_pq_length_stats": None,
     }
 
     @property
@@ -1410,9 +1409,7 @@ class ReadParquetFSSpec(ReadParquet):
         dataset_info["all_columns"] = all_columns
         dataset_info["calculate_divisions"] = self.calculate_divisions
 
-        self.operands[type(self)._parameters.index("_dataset_info_cache")] = (
-            dataset_info
-        )
+        self._dataset_info_cache - dataset_info
         return dataset_info
 
     def _filtered_task(self, name: Key, index: int) -> Task:
@@ -1480,29 +1477,27 @@ class ReadParquetFSSpec(ReadParquet):
         """Return known partition lengths using parquet statistics"""
         if not self.filters:
             self._update_length_statistics()
-            return tuple(  # type: ignore
+            return tuple(
                 length
-                for i, length in enumerate(self._pq_length_stats)  # type: ignore
+                for i, length in enumerate(self._pq_length_stats)
                 if not self._filtered or i in self._partitions
             )
         return None
 
-    def _update_length_statistics(self):
+    @cached_property
+    def _pq_length_stats(self):
         """Ensure that partition-length statistics are up to date"""
 
-        if not self._pq_length_stats:
-            if self._plan["statistics"]:
-                # Already have statistics from original API call
-                self._pq_length_stats = tuple(
-                    stat["num-rows"]
-                    for i, stat in enumerate(self._plan["statistics"])
-                    if not self._filtered or i in self._partitions
-                )
-            else:
-                # Need to go back and collect statistics
-                self._pq_length_stats = tuple(
-                    stat["num-rows"] for stat in _collect_pq_statistics(self)
-                )
+        if self._plan["statistics"]:
+            # Already have statistics from original API call
+            return tuple(
+                stat["num-rows"]
+                for i, stat in enumerate(self._plan["statistics"])
+                if not self._filtered or i in self._partitions
+            )
+        else:
+            # Need to go back and collect statistics
+            return tuple(stat["num-rows"] for stat in _collect_pq_statistics(self))
 
 
 #


Test patch:
diff --git a/dask/tests/test_expr.py b/dask/tests/test_expr.py
new file mode 100644
index 000000000..f7f146915
--- /dev/null
+++ b/dask/tests/test_expr.py
@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+import pytest
+
+from dask._expr import Expr
+
+
+def test_setattr():
+    class MyExpr(Expr):
+        _parameters = ["foo", "bar"]
+
+    e = MyExpr(foo=1, bar=2)
+    e.bar = 3
+    assert e.bar == 3
+    with pytest.raises(AttributeError):
+        e.baz = 4