Integrate dask-expr and make CI happy (#980)

Author: milesgranger

.github/workflows/lint.yaml

@@ -5,6 +5,8 @@ jobs:
   lint:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v3
-      - uses: pre-commit/[email protected]
+      - uses: actions/[email protected]
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.9'
+      - uses: pre-commit/[email protected]

.github/workflows/tests.yaml

@@ -10,12 +10,14 @@ jobs:
       matrix:
         # os: ["windows-latest", "ubuntu-latest", "macos-latest"]
         os: ["ubuntu-latest"]
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.9", "3.10", "3.11"]
+        query-planning: [true, false]
 
     env:
       PYTHON_VERSION: ${{ matrix.python-version }}
       PARALLEL: "true"
       COVERAGE: "true"
+      DASK_DATAFRAME__QUERY_PLANNING: ${{ matrix.query-planning }}
 
     steps:
       - name: Checkout source

.pre-commit-config.yaml

@@ -1,16 +1,19 @@
 repos:
--   repo: https://github.com/python/black
-    rev: 22.3.0
+-   repo: https://github.com/psf/black
+    rev: 23.12.1
     hooks:
     - id: black
       language_version: python3
+      args:
+          - --target-version=py39
 -   repo: https://github.com/pycqa/flake8
-    rev: 3.7.9
+    rev: 7.0.0
     hooks:
     - id: flake8
       language_version: python3
--   repo: https://github.com/timothycrosley/isort
-    rev: 4.3.21
+      args: ["--ignore=E501,W503,E203,E741,E731"]
+-   repo: https://github.com/pycqa/isort
+    rev: 5.13.2
     hooks:
     - id: isort
       language_version: python3

ci/environment-3.10.yaml

@@ -3,7 +3,6 @@ channels:
   - conda-forge
   - defaults
 dependencies:
-  - dask
   - dask-glm
   - multipledispatch >=0.4.9
   - mypy
@@ -21,3 +20,7 @@ dependencies:
   - scipy
   - sparse
   - toolz
+  - pip
+  - pip:
+      - git+https://github.com/dask-contrib/dask-expr
+      - git+https://github.com/dask/dask

ci/environment-3.8.yaml renamed to ci/environment-3.11.yaml

@@ -1,9 +1,8 @@
-name: dask-ml-3.8
+name: dask-ml-3.11
 channels:
   - conda-forge
   - defaults
 dependencies:
-  - dask
   - dask-glm
   - multipledispatch >=0.4.9
   - mypy
@@ -16,8 +15,12 @@ dependencies:
   - pytest
   - pytest-cov
   - pytest-mock
-  - python=3.8.*
+  - python=3.11.*
   - scikit-learn >=1.2.0
   - scipy
   - sparse
   - toolz
+  - pip
+  - pip:
+      - git+https://github.com/dask-contrib/dask-expr
+      - git+https://github.com/dask/dask

ci/environment-3.9.yaml

@@ -3,7 +3,6 @@ channels:
   - conda-forge
   - defaults
 dependencies:
-  - dask
   - dask-glm
   - multipledispatch >=0.4.9
   - mypy
@@ -21,3 +20,7 @@ dependencies:
   - scipy
   - sparse
   - toolz
+  - pip
+  - pip:
+      - git+https://github.com/dask-contrib/dask-expr
+      - git+https://github.com/dask/dask

ci/environment-docs.yaml

@@ -5,28 +5,21 @@ channels:
 dependencies:
   - black
   - coverage
-  - graphviz
   - heapdict
   - ipykernel
   - ipython
   - multipledispatch
   - mypy
-  - nbsphinx
   - nomkl
   - nose
   - numba
   - numpy
-  - numpydoc
-  - pandas
   - psutil
   - python=3.10
   - sortedcontainers
   - scikit-learn >=1.2.0
   - scipy
   - sparse
-  - sphinx
-  - sphinx_rtd_theme
-  - sphinx-gallery
   - tornado
   - toolz
   - zict
@@ -35,5 +28,30 @@ dependencies:
   - dask-glm
   - dask-xgboost
   - pip:
-    - dask-sphinx-theme >=3.0.0
     - graphviz
+    - numpydoc
+    - sphinx>=4.0.0,<5.0.0
+    - dask-sphinx-theme>=3.0.0
+    - sphinx-click
+    - sphinx-copybutton
+    - sphinx-remove-toctrees
+    - sphinx_autosummary_accessors
+    - sphinx-tabs
+    - sphinx-design
+    - jupyter_sphinx
+    # FIXME: `sphinxcontrib-*` pins are a workaround until we have sphinx>=5.
+    #        See https://github.com/dask/dask-sphinx-theme/issues/68.
+    - sphinxcontrib-applehelp>=1.0.0,<1.0.7
+    - sphinxcontrib-devhelp>=1.0.0,<1.0.6
+    - sphinxcontrib-htmlhelp>=2.0.0,<2.0.5
+    - sphinxcontrib-serializinghtml>=1.1.0,<1.1.10
+    - sphinxcontrib-qthelp>=1.0.0,<1.0.7
+    - toolz
+    - cloudpickle>=1.5.0
+    - pandas>=1.4.0
+    - dask-expr
+    - fsspec
+    - scipy
+    - pytest
+    - pytest-check-links
+    - requests-cache

dask_ml/_partial.py

@@ -29,7 +29,7 @@ def fit(
     shuffle_blocks=True,
     random_state=None,
     assume_equal_chunks=False,
-    **kwargs
+    **kwargs,
 ):
     """Fit scikit learn model against dask arrays
 

dask_ml/ensemble/_blockwise.py

@@ -1,12 +1,11 @@
 import dask
 import dask.array as da
-import dask.dataframe as dd
 import numpy as np
 import sklearn.base
 from sklearn.utils.validation import check_is_fitted
 
 from ..base import ClassifierMixin, RegressorMixin
-from ..utils import check_array
+from ..utils import check_array, is_frame_base
 
 
 class BlockwiseBase(sklearn.base.BaseEstimator):
@@ -62,7 +61,7 @@ def _predict(self, X):
                 dtype=np.dtype(dtype),
                 chunks=chunks,
             )
-        elif isinstance(X, dd._Frame):
+        elif is_frame_base(X):
             meta = np.empty((0, len(self.classes_)), dtype=dtype)
             combined = X.map_partitions(
                 _predict_stack, estimators=self.estimators_, meta=meta
@@ -184,7 +183,7 @@ def _collect_probas(self, X):
                 chunks=chunks,
                 meta=meta,
             )
-        elif isinstance(X, dd._Frame):
+        elif is_frame_base(X):
             # TODO: replace with a _predict_proba_stack version.
             # This current raises; dask.dataframe doesn't like map_partitions that
             # return new axes.

dask_ml/linear_model/utils.py

@@ -5,29 +5,60 @@
 import numpy as np
 from multipledispatch import dispatch
 
+if getattr(dd, "_dask_expr_enabled", lambda: False)():
+    import dask_expr
 
-@dispatch(dd._Frame)
-def exp(A):
-    return da.exp(A)
+    @dispatch(dask_expr.FrameBase)
+    def exp(A):
+        return da.exp(A)
 
+    @dispatch(dask_expr.FrameBase)
+    def absolute(A):
+        return da.absolute(A)
 
-@dispatch(dd._Frame)
-def absolute(A):
-    return da.absolute(A)
+    @dispatch(dask_expr.FrameBase)
+    def sign(A):
+        return da.sign(A)
 
+    @dispatch(dask_expr.FrameBase)
+    def log1p(A):
+        return da.log1p(A)
 
-@dispatch(dd._Frame)
-def sign(A):
-    return da.sign(A)
+    @dispatch(dask_expr.FrameBase)  # noqa: F811
+    def add_intercept(X):  # noqa: F811
+        columns = X.columns
+        if "intercept" in columns:
+            raise ValueError("'intercept' column already in 'X'")
+        return X.assign(intercept=1)[["intercept"] + list(columns)]
 
+else:
 
-@dispatch(dd._Frame)
-def log1p(A):
-    return da.log1p(A)
+    @dispatch(dd._Frame)
+    def exp(A):
+        return da.exp(A)
 
+    @dispatch(dd._Frame)
+    def absolute(A):
+        return da.absolute(A)
 
-@dispatch(np.ndarray)
-def add_intercept(X):
+    @dispatch(dd._Frame)
+    def sign(A):
+        return da.sign(A)
+
+    @dispatch(dd._Frame)
+    def log1p(A):
+        return da.log1p(A)
+
+    @dispatch(dd._Frame)  # noqa: F811
+    def add_intercept(X):  # noqa: F811
+        columns = X.columns
+        if "intercept" in columns:
+            raise ValueError("'intercept' column already in 'X'")
+        return X.assign(intercept=1)[["intercept"] + list(columns)]
+
+
+@dispatch(np.ndarray)  # noqa: F811
+def add_intercept(X):  # noqa: F811
     return _add_intercept(X)
 
 
@@ -53,14 +84,6 @@ def add_intercept(X):  # noqa: F811
     return X.map_blocks(_add_intercept, dtype=X.dtype, chunks=chunks)
 
 
-@dispatch(dd.DataFrame)  # noqa: F811
-def add_intercept(X):  # noqa: F811
-    columns = X.columns
-    if "intercept" in columns:
-        raise ValueError("'intercept' column already in 'X'")
-    return X.assign(intercept=1)[["intercept"] + list(columns)]
-
-
 @dispatch(np.ndarray)  # noqa: F811
 def lr_prob_stack(prob):  # noqa: F811
     return np.vstack([1 - prob, prob]).T