Expr setattr (dask#11836)

Author: fjetter

dask/_expr.py

@@ -164,6 +164,19 @@ def _depth(self, cache=None):
                     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)

dask/array/_array_expr/_expr.py

@@ -11,7 +11,7 @@
 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
@@ -20,7 +20,6 @@
 
 
 class ArrayExpr(Expr):
-    _cached_keys = None
 
     def _operands_for_repr(self):
         return []
@@ -75,25 +74,35 @@ def __len__(self):
             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

dask/dataframe/dask_expr/_expr.py

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

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

dask/dataframe/dask_expr/io/parquet.py

@@ -717,7 +717,6 @@ def default_types_mapper(pyarrow_dtype):
 
 
 class ReadParquet(PartitionsFiltered, BlockwiseIO):
-    _pq_length_stats = None
     _absorb_projections = True
     _filter_passthrough = False
 
@@ -1074,9 +1073,7 @@ def _dataset_info(self):
 
         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 @@ def _dataset_info(self):
         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 @@ def _get_lengths(self) -> tuple | None:
         """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))
 
 
 #

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