Repository: dask/dask
Pull request #11688 (2025-01-21 13:58:14)

Disentangle core expression class from dataframe specific code
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`

This mostly moves stuff around so that the array class can inherit from the core expression implementation

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._

     15 files  ±0       15 suites  ±0   4h 24m 24s ⏱️ -56s
 17 141 tests ±0   15 960 ✅  -  1   1 180 💤 ±0   1 ❌ + 1 
211 212 runs  ±0  194 178 ✅  - 13  17 022 💤 +1  12 ❌ +12 

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

Results for commit 15f20001. ± Comparison against base commit f5bd5590.

[test-results]:data:application/gzip;base64,H4sIAL+xj2cC/02Myw6DIBBFf8Ww7mJmBIH+TGNHSUh9NAirpv9eMIpdnnNvzkc4P42buDeobo3Yko8XDSn00a/LzqaTWeU17rtGiSc/tsS8n2wHl3z5d5FoLud6PxVXxRjCGrIpl5CWkiZEQjq4pq1Ebao80hqoHs90Ff9pXufZxwwClSMAQAXGWktWPvuW0XGHirklo/XgDLSj+P4A34BSwRsBAAA=

phofl: merging to unblock follow up stuff

Patch:
diff --git a/dask/dataframe/dask_expr/_core.py b/dask/_expr.py
similarity index 87%
rename from dask/dataframe/dask_expr/_core.py
rename to dask/_expr.py
index 33e86c90a..9265bc1b6 100644
--- a/dask/dataframe/dask_expr/_core.py
+++ b/dask/_expr.py
@@ -7,15 +7,13 @@ from collections import defaultdict
 from collections.abc import Generator
 from typing import TYPE_CHECKING, Literal
 
-import pandas as pd
 import toolz
 
 import dask
 from dask._task_spec import Task
-from dask.dataframe.core import is_dataframe_like, is_index_like, is_series_like
-from dask.dataframe.dask_expr._util import _BackendData, _tokenize_deterministic
+from dask.tokenize import _tokenize_deterministic
 from dask.typing import Key
-from dask.utils import funcname, import_required, is_arraylike
+from dask.utils import funcname, import_required
 
 if TYPE_CHECKING:
     # TODO import from typing (requires Python >=3.10)
@@ -70,14 +68,7 @@ class Expr:
         return None
 
     def _operands_for_repr(self):
-        to_include = []
-        for param, operand in zip(self._parameters, self.operands):
-            if isinstance(operand, Expr) or (
-                not isinstance(operand, (pd.Series, pd.DataFrame))
-                and operand != self._defaults.get(param)
-            ):
-                to_include.append(f"{param}={operand!r}")
-        return to_include
+        raise NotImplementedError("Subclasses should implement this method")
 
     def __str__(self):
         s = ", ".join(self._operands_for_repr())
@@ -102,33 +93,7 @@ class Expr:
         return header
 
     def _tree_repr_lines(self, indent=0, recursive=True):
-        header = funcname(type(self)) + ":"
-        lines = []
-        for i, op in enumerate(self.operands):
-            if isinstance(op, Expr):
-                if recursive:
-                    lines.extend(op._tree_repr_lines(2))
-            else:
-                if isinstance(op, _BackendData):
-                    op = op._data
-
-                # TODO: this stuff is pandas-specific
-                if isinstance(op, pd.core.base.PandasObject):
-                    op = "<pandas>"
-                elif is_dataframe_like(op):
-                    op = "<dataframe>"
-                elif is_index_like(op):
-                    op = "<index>"
-                elif is_series_like(op):
-                    op = "<series>"
-                elif is_arraylike(op):
-                    op = "<array>"
-                header = self._tree_repr_argument_construction(i, op, header)
-
-        lines = [header] + lines
-        lines = [" " * indent + line for line in lines]
-
-        return lines
+        raise NotImplementedError("Subclasses should implement this method")
 
     def tree_repr(self):
         return os.linesep.join(self._tree_repr_lines())
@@ -248,7 +213,7 @@ class Expr:
 
         return {
             (self._name, i): self._task((self._name, i), i)
-            for i in range(self.npartitions)
+            for i in range(self.npartitions)  # type: ignore
         }
 
     def rewrite(self, kind: str, rewritten):
@@ -301,7 +266,7 @@ class Expr:
             # Rewrite all of the children
             new_operands = []
             changed = False
-            for operand in expr.operands:
+            for operand in expr.operands:  # type: ignore
                 if isinstance(operand, Expr):
                     new = operand.rewrite(kind=kind, rewritten=rewritten)
                     rewritten[operand._name] = new
@@ -368,7 +333,7 @@ class Expr:
             # Rewrite all of the children
             new_operands = []
             changed = False
-            for operand in expr.operands:
+            for operand in expr.operands:  # type: ignore
                 if isinstance(operand, Expr):
                     # Bandaid for now, waiting for Singleton
                     dependents[operand._name].append(weakref.ref(expr))
@@ -431,7 +396,7 @@ class Expr:
         # Lower all children
         new_operands = []
         changed = False
-        for operand in out.operands:
+        for operand in out.operands:  # type: ignore
             if isinstance(operand, Expr):
                 new = operand.lower_once(lowered)
                 if new._name != operand._name:
@@ -488,35 +453,6 @@ class Expr:
     def _meta(self):
         raise NotImplementedError()
 
-    def __getattr__(self, key):
-        try:
-            return object.__getattribute__(self, key)
-        except AttributeError as err:
-            if key.startswith("_meta"):
-                # Avoid a recursive loop if/when `self._meta*`
-                # produces an `AttributeError`
-                raise RuntimeError(
-                    f"Failed to generate metadata for {self}. "
-                    "This operation may not be supported by the current backend."
-                )
-
-            # Allow operands to be accessed as attributes
-            # as long as the keys are not already reserved
-            # by existing methods/properties
-            _parameters = type(self)._parameters
-            if key in _parameters:
-                idx = _parameters.index(key)
-                return self.operands[idx]
-            if is_dataframe_like(self._meta) and key in self._meta.columns:
-                return self[key]
-
-            link = "https://github.com/dask-contrib/dask-expr/blob/main/README.md#api-coverage"
-            raise AttributeError(
-                f"{err}\n\n"
-                "This often means that you are attempting to use an unsupported "
-                f"API function. Current API coverage is documented here: {link}."
-            )
-
     def __dask_graph__(self):
         """Traverse expression tree, collect layers"""
         stack = [self]
@@ -627,7 +563,7 @@ class Expr:
 
         changed = False
         new_operands = []
-        for i, operand in enumerate(self.operands):
+        for i, operand in enumerate(self.operands):  # type: ignore
             if i < len(self._parameters) and self._parameters[i] in substitutions:
                 new_operands.append(substitutions[self._parameters[i]])
                 changed = True
diff --git a/dask/dataframe/dask_expr/_collection.py b/dask/dataframe/dask_expr/_collection.py
index 8622e20c0..f0479c2b9 100644
--- a/dask/dataframe/dask_expr/_collection.py
+++ b/dask/dataframe/dask_expr/_collection.py
@@ -26,6 +26,7 @@ import dask.array as da
 import dask.dataframe.dask_expr._backends  # noqa: F401
 import dask.dataframe.methods as methods
 from dask import compute, get_annotations
+from dask._expr import OptimizerStage
 from dask.array import Array
 from dask.base import DaskMethodsMixin, is_dask_collection, named_schedulers
 from dask.core import flatten
@@ -50,7 +51,6 @@ from dask.dataframe.dask_expr._categorical import (
     GetCategories,
 )
 from dask.dataframe.dask_expr._concat import Concat
-from dask.dataframe.dask_expr._core import OptimizerStage
 from dask.dataframe.dask_expr._datetime import DatetimeAccessor
 from dask.dataframe.dask_expr._describe import DescribeNonNumeric, DescribeNumeric
 from dask.dataframe.dask_expr._dispatch import get_collection_type
@@ -104,7 +104,6 @@ from dask.dataframe.dask_expr._util import (
     _is_any_real_numeric_dtype,
     _maybe_from_pandas,
     _raise_if_object_series,
-    _tokenize_deterministic,
     _validate_axis,
     get_specified_shuffle,
     is_scalar,
@@ -127,6 +126,7 @@ from dask.dataframe.utils import (
     pyarrow_strings_enabled,
 )
 from dask.delayed import Delayed, delayed
+from dask.tokenize import _tokenize_deterministic
 from dask.utils import (
     IndexCallable,
     M,
diff --git a/dask/dataframe/dask_expr/_expr.py b/dask/dataframe/dask_expr/_expr.py
index dac7d0615..fe7390bdb 100644
--- a/dask/dataframe/dask_expr/_expr.py
+++ b/dask/dataframe/dask_expr/_expr.py
@@ -14,6 +14,7 @@ import pandas as pd
 from pandas.errors import PerformanceWarning
 from tlz import merge_sorted, partition, unique
 
+from dask import _expr as core
 from dask._task_spec import Alias, DataNode, Task, TaskRef, execute_graph
 from dask.array import Array
 from dask.core import flatten
@@ -31,11 +32,10 @@ from dask.dataframe.core import (
     safe_head,
     total_mem_usage,
 )
-from dask.dataframe.dask_expr import _core as core
 from dask.dataframe.dask_expr._util import (
+    _BackendData,
     _calc_maybe_new_divisions,
     _convert_to_list,
-    _tokenize_deterministic,
     _tokenize_partial,
     is_scalar,
 )
@@ -48,7 +48,7 @@ from dask.dataframe.utils import (
     raise_on_meta_error,
     valid_divisions,
 )
-from dask.tokenize import normalize_token
+from dask.tokenize import _tokenize_deterministic, normalize_token
 from dask.typing import Key, no_default
 from dask.utils import (
     M,
@@ -106,6 +106,74 @@ class Expr(core.Expr):
     def nbytes(self):
         return NBytes(self)
 
+    def _tree_repr_lines(self, indent=0, recursive=True):
+        header = funcname(type(self)) + ":"
+        lines = []
+        for i, op in enumerate(self.operands):
+            if isinstance(op, Expr):
+                if recursive:
+                    lines.extend(op._tree_repr_lines(2))
+            else:
+                if isinstance(op, _BackendData):
+                    op = op._data
+
+                # TODO: this stuff is pandas-specific
+                if isinstance(op, pd.core.base.PandasObject):
+                    op = "<pandas>"
+                elif is_dataframe_like(op):
+                    op = "<dataframe>"
+                elif is_index_like(op):
+                    op = "<index>"
+                elif is_series_like(op):
+                    op = "<series>"
+                elif is_arraylike(op):
+                    op = "<array>"
+                header = self._tree_repr_argument_construction(i, op, header)
+
+        lines = [header] + lines
+        lines = [" " * indent + line for line in lines]
+
+        return lines
+
+    def _operands_for_repr(self):
+        to_include = []
+        for param, operand in zip(self._parameters, self.operands):
+            if isinstance(operand, Expr) or (
+                not isinstance(operand, (pd.Series, pd.DataFrame))
+                and operand != self._defaults.get(param)
+            ):
+                to_include.append(f"{param}={operand!r}")
+        return to_include
+
+    def __getattr__(self, key):
+        try:
+            return object.__getattribute__(self, key)
+        except AttributeError as err:
+            if key.startswith("_meta"):
+                # Avoid a recursive loop if/when `self._meta*`
+                # produces an `AttributeError`
+                raise RuntimeError(
+                    f"Failed to generate metadata for {self}. "
+                    "This operation may not be supported by the current backend."
+                )
+
+            # Allow operands to be accessed as attributes
+            # as long as the keys are not already reserved
+            # by existing methods/properties
+            _parameters = type(self)._parameters
+            if key in _parameters:
+                idx = _parameters.index(key)
+                return self.operands[idx]
+            if is_dataframe_like(self._meta) and key in self._meta.columns:
+                return self[key]
+
+            link = "https://github.com/dask-contrib/dask-expr/blob/main/README.md#api-coverage"
+            raise AttributeError(
+                f"{err}\n\n"
+                "This often means that you are attempting to use an unsupported "
+                f"API function. Current API coverage is documented here: {link}."
+            )
+
     def __getitem__(self, other):
         if isinstance(other, Expr):
             if not are_co_aligned(self, other):
@@ -590,7 +658,7 @@ class MapPartitions(Blockwise):
         "token",
         "kwargs",
     ]
-    _defaults = {
+    _defaults: dict = {
         "kwargs": None,
         "align_dataframes": True,
         "parent_meta": None,
@@ -842,7 +910,7 @@ class MapOverlap(MapPartitions):
         "token",
         "kwargs",
     ]
-    _defaults = {
+    _defaults: dict = {
         "meta": None,
         "enfore_metadata": True,
         "transform_divisions": True,
@@ -1420,7 +1488,7 @@ class ToTimestamp(Elemwise):
 
 class CombineSeries(Elemwise):
     _parameters = ["frame", "other", "func", "fill_value"]
-    _defaults = {"fill_value": None}
+    _defaults: dict = {"fill_value": None}
     operation = M.combine
 
     @functools.cached_property
diff --git a/dask/dataframe/dask_expr/_merge.py b/dask/dataframe/dask_expr/_merge.py
index 424936594..190cde1e2 100644
--- a/dask/dataframe/dask_expr/_merge.py
+++ b/dask/dataframe/dask_expr/_merge.py
@@ -33,11 +33,7 @@ from dask.dataframe.dask_expr._shuffle import (
     _is_numeric_cast_type,
     _select_columns_or_index,
 )
-from dask.dataframe.dask_expr._util import (
-    _convert_to_list,
-    _tokenize_deterministic,
-    is_scalar,
-)
+from dask.dataframe.dask_expr._util import _convert_to_list, is_scalar
 from dask.dataframe.dispatch import make_meta, meta_nonempty
 from dask.dataframe.multi import (
     _concat_wrapper,
@@ -46,6 +42,7 @@ from dask.dataframe.multi import (
     merge_chunk,
 )
 from dask.dataframe.shuffle import partitioning_index
+from dask.tokenize import _tokenize_deterministic
 from dask.typing import Key
 from dask.utils import apply, get_default_shuffle_method
 
diff --git a/dask/dataframe/dask_expr/_reductions.py b/dask/dataframe/dask_expr/_reductions.py
index e95f5323f..5b5f248dd 100644
--- a/dask/dataframe/dask_expr/_reductions.py
+++ b/dask/dataframe/dask_expr/_reductions.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import functools
 from collections.abc import Callable
+from typing import Any as AnyT
 
 import numpy as np
 import pandas as pd
@@ -37,8 +38,9 @@ from dask.dataframe.dask_expr._expr import (
     determine_column_projection,
     plain_column_projection,
 )
-from dask.dataframe.dask_expr._util import _tokenize_deterministic, is_scalar
+from dask.dataframe.dask_expr._util import is_scalar
 from dask.dataframe.dispatch import make_meta, meta_nonempty
+from dask.tokenize import _tokenize_deterministic
 from dask.typing import no_default
 from dask.utils import M, apply, funcname
 
@@ -783,7 +785,7 @@ class Reduction(ApplyConcatApply):
     methods.
     """
 
-    _defaults = {
+    _defaults: AnyT = {
         "skipna": True,
         "numeric_only": False,
         "min_count": 0,
diff --git a/dask/dataframe/dask_expr/_util.py b/dask/dataframe/dask_expr/_util.py
index 6869f2465..aef4d8e90 100644
--- a/dask/dataframe/dask_expr/_util.py
+++ b/dask/dataframe/dask_expr/_util.py
@@ -12,7 +12,7 @@ from packaging.version import Version
 from dask import config, is_dask_collection
 from dask.dataframe._compat import is_string_dtype
 from dask.dataframe.core import is_dataframe_like, is_series_like
-from dask.tokenize import normalize_token, tokenize
+from dask.tokenize import _tokenize_deterministic, normalize_token
 from dask.utils import get_default_shuffle_method
 
 K = TypeVar("K", bound=Hashable)
@@ -95,11 +95,6 @@ def is_scalar(x):
     return not isinstance(x, Expr)
 
 
-def _tokenize_deterministic(*args, **kwargs) -> str:
-    # Utility to be strict about deterministic tokens
-    return tokenize(*args, ensure_deterministic=True, **kwargs)
-
-
 def _tokenize_partial(expr, ignore: list | None = None) -> str:
     # Helper function to "tokenize" the operands
     # that are not in the `ignore` list
@@ -145,7 +140,7 @@ class _BackendData:
 
     @functools.cached_property
     def _token(self):
-        from dask.dataframe.dask_expr._util import _tokenize_deterministic
+        from dask.tokenize import _tokenize_deterministic
 
         return _tokenize_deterministic(self._data)
 
diff --git a/dask/dataframe/dask_expr/diagnostics/_analyze.py b/dask/dataframe/dask_expr/diagnostics/_analyze.py
index 283b0ff58..29d2cadc4 100644
--- a/dask/dataframe/dask_expr/diagnostics/_analyze.py
+++ b/dask/dataframe/dask_expr/diagnostics/_analyze.py
@@ -7,13 +7,14 @@ import pandas as pd
 
 from dask.base import DaskMethodsMixin
 from dask.dataframe.dask_expr._expr import Blockwise, Expr
-from dask.dataframe.dask_expr._util import _tokenize_deterministic, is_scalar
+from dask.dataframe.dask_expr._util import is_scalar
 from dask.dataframe.dask_expr.diagnostics._explain import (
     _add_graphviz_edges,
     _explain_info,
 )
 from dask.dataframe.dask_expr.io.io import FusedIO
 from dask.sizeof import sizeof
+from dask.tokenize import _tokenize_deterministic
 from dask.utils import format_bytes, import_required
 
 if TYPE_CHECKING:
diff --git a/dask/dataframe/dask_expr/diagnostics/_explain.py b/dask/dataframe/dask_expr/diagnostics/_explain.py
index 2e60c2ec0..1d5beebc3 100644
--- a/dask/dataframe/dask_expr/diagnostics/_explain.py
+++ b/dask/dataframe/dask_expr/diagnostics/_explain.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from dask.dataframe.dask_expr._core import OptimizerStage
+from dask._expr import OptimizerStage
 from dask.dataframe.dask_expr._expr import Expr, Projection, optimize_until
 from dask.dataframe.dask_expr._merge import Merge
 from dask.dataframe.dask_expr.io.parquet import ReadParquet
diff --git a/dask/dataframe/dask_expr/io/_delayed.py b/dask/dataframe/dask_expr/io/_delayed.py
index f2fb8eeac..879191bb3 100644
--- a/dask/dataframe/dask_expr/io/_delayed.py
+++ b/dask/dataframe/dask_expr/io/_delayed.py
@@ -12,11 +12,11 @@ from dask.dataframe.dask_expr._expr import (
     DelayedsExpr,
     PartitionsFiltered,
 )
-from dask.dataframe.dask_expr._util import _tokenize_deterministic
 from dask.dataframe.dask_expr.io import BlockwiseIO
 from dask.dataframe.dispatch import make_meta
 from dask.dataframe.utils import check_meta, pyarrow_strings_enabled
 from dask.delayed import Delayed, delayed
+from dask.tokenize import _tokenize_deterministic
 from dask.typing import Key
 
 if TYPE_CHECKING:
diff --git a/dask/dataframe/dask_expr/io/io.py b/dask/dataframe/dask_expr/io/io.py
index 5c85c31e9..cec887496 100644
--- a/dask/dataframe/dask_expr/io/io.py
+++ b/dask/dataframe/dask_expr/io/io.py
@@ -22,13 +22,10 @@ from dask.dataframe.dask_expr._expr import (
     no_default,
 )
 from dask.dataframe.dask_expr._reductions import Len
-from dask.dataframe.dask_expr._util import (
-    _BackendData,
-    _convert_to_list,
-    _tokenize_deterministic,
-)
+from dask.dataframe.dask_expr._util import _BackendData, _convert_to_list
 from dask.dataframe.dispatch import make_meta
 from dask.dataframe.io.io import _meta_from_array, sorted_division_locations
+from dask.tokenize import _tokenize_deterministic
 from dask.typing import Key
 from dask.utils import funcname, is_series_like
 
diff --git a/dask/dataframe/dask_expr/io/parquet.py b/dask/dataframe/dask_expr/io/parquet.py
index a56d80de0..733d27fb6 100644
--- a/dask/dataframe/dask_expr/io/parquet.py
+++ b/dask/dataframe/dask_expr/io/parquet.py
@@ -45,7 +45,7 @@ from dask.dataframe.dask_expr._expr import (
     determine_column_projection,
 )
 from dask.dataframe.dask_expr._reductions import Len
-from dask.dataframe.dask_expr._util import _convert_to_list, _tokenize_deterministic
+from dask.dataframe.dask_expr._util import _convert_to_list
 from dask.dataframe.dask_expr.io import BlockwiseIO, PartitionsFiltered
 from dask.dataframe.dask_expr.io.io import FusedParquetIO
 from dask.dataframe.io.parquet.core import (
@@ -60,7 +60,7 @@ from dask.dataframe.io.parquet.core import (
 from dask.dataframe.io.parquet.utils import _split_user_options
 from dask.dataframe.io.utils import _is_local_fs
 from dask.delayed import delayed
-from dask.tokenize import normalize_token, tokenize
+from dask.tokenize import _tokenize_deterministic, normalize_token, tokenize
 from dask.typing import Key
 from dask.utils import apply, funcname, natural_sort_key, parse_bytes, typename
 
diff --git a/dask/tokenize.py b/dask/tokenize.py
index 1f0c6ff32..d6601778c 100644
--- a/dask/tokenize.py
+++ b/dask/tokenize.py
@@ -442,3 +442,8 @@ def register_numpy():
     @normalize_token.register(np.dtype)
     def normalize_dtype(dtype):
         return dtype.str
+
+
+def _tokenize_deterministic(*args, **kwargs) -> str:
+    # Utility to be strict about deterministic tokens
+    return tokenize(*args, ensure_deterministic=True, **kwargs)


Test patch:
diff --git a/dask/dataframe/dask_expr/tests/test_core.py b/dask/dataframe/dask_expr/tests/test_core.py
index d1c75650a..a039f369a 100644
--- a/dask/dataframe/dask_expr/tests/test_core.py
+++ b/dask/dataframe/dask_expr/tests/test_core.py
@@ -2,18 +2,24 @@ from __future__ import annotations
 
 import pytest
 
-from dask.dataframe.dask_expr._core import Expr
+from dask._expr import Expr
 
 
 class ExprB(Expr):
     def _simplify_down(self):
         return ExprA()
 
+    def _operands_for_repr(self):
+        return []
+
 
 class ExprA(Expr):
     def _simplify_down(self):
         return ExprB()
 
+    def _operands_for_repr(self):
+        return []
+
 
 def test_endless_simplify():
     expr = ExprA()
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
index 0331200f4..0a07ce438 100644
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -80,7 +80,7 @@ def test_is_dask_collection_dask_expr_does_not_materialize():
     pytest.importorskip("pandas")
     dx = pytest.importorskip("dask.dataframe.dask_expr")
 
-    class DoNotMaterialize(dx._core.Expr):
+    class DoNotMaterialize(dask._expr.Expr):
         @property
         def _meta(self):
             return 0