Merge branch 'main' into io-tasks

dask · Jan 3, 2025 · 808eb1b · 808eb1b
2 parents 2ed6a46 + 8f1b241
commit 808eb1b
Show file tree

Hide file tree

Showing 15 changed files with 22 additions and 751 deletions.
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
@@ -57,16 +57,6 @@ jobs:
             label: no_queue
             partition: "not ci1"
 
-          # Set dataframe.query-planning: false
-          - os: ubuntu-latest
-            environment: "3.10"
-            label: no_expr
-            partition: "ci1"
-          - os: ubuntu-latest
-            environment: "3.10"
-            label: no_expr
-            partition: "not ci1"
-
           # dask.array P2P shuffle
           - os: ubuntu-latest
             environment: mindeps
@@ -83,12 +73,12 @@ jobs:
           - os: ubuntu-latest
             environment: mindeps
             label: pandas
-            extra_packages: [numpy=1.24, pandas=2.0, pyarrow=14.0.1]
+            extra_packages: [numpy=1.24, pandas=2.0, pyarrow=14.0.1, dask-expr]
             partition: "ci1"
           - os: ubuntu-latest
             environment: mindeps
             label: pandas
-            extra_packages: [numpy=1.24, pandas=2.0, pyarrow=14.0.1]
+            extra_packages: [numpy=1.24, pandas=2.0, pyarrow=14.0.1, dask-expr]
             partition: "not ci1"
 
           - os: ubuntu-latest
@@ -221,11 +211,6 @@ jobs:
         if: ${{ matrix.label == 'no_queue' }}
         run: echo "DASK_DISTRIBUTED__SCHEDULER__WORKER_SATURATION=inf" >> $GITHUB_ENV
 
-      - name: Set up dask env to disable dask-expr
-        shell: bash -l {0}
-        if: ${{ matrix.label == 'no_expr' }}
-        run: echo "DASK_DATAFRAME__QUERY_PLANNING=False" >> $GITHUB_ENV
-
       - name: Print host info
         # host_info.py imports numpy, which isn't a direct dependency of distributed
         if: matrix.environment != 'mindeps'

diff --git a/continuous_integration/environment-3.10.yaml b/continuous_integration/environment-3.10.yaml
@@ -19,6 +19,7 @@ dependencies:
   - ipykernel
   - ipywidgets
   - jinja2
+  - jupyter_events<0.11
   - jupyter-server-proxy
   - jupyterlab
   - locket

diff --git a/continuous_integration/environment-3.11.yaml b/continuous_integration/environment-3.11.yaml
@@ -18,6 +18,7 @@ dependencies:
   - ipykernel
   - ipywidgets
   - jinja2
+  - jupyter_events<0.11
   - jupyter-server-proxy
   - jupyterlab
   - locket

diff --git a/continuous_integration/environment-3.12.yaml b/continuous_integration/environment-3.12.yaml
@@ -18,6 +18,7 @@ dependencies:
   - ipykernel
   - ipywidgets
   - jinja2
+  - jupyter_events<0.11
   - jupyter-server-proxy
   - jupyterlab
   - locket

diff --git a/continuous_integration/environment-3.13.yaml b/continuous_integration/environment-3.13.yaml
@@ -17,6 +17,7 @@ dependencies:
   - ipykernel
   - ipywidgets
   - jinja2
+  - jupyter_events<0.11
   - jupyter-server-proxy
   - jupyterlab
   - locket

diff --git a/distributed/protocol/tests/test_highlevelgraph.py b/distributed/protocol/tests/test_highlevelgraph.py
@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-import contextlib
-
 import pytest
 
 np = pytest.importorskip("numpy")
@@ -175,11 +173,9 @@ async def test_dataframe_annotations(c, s, a, b):
     acol = df["a"]
     bcol = df["b"]
 
-    ctx = contextlib.nullcontext()
-    if dd._dask_expr_enabled():
-        ctx = pytest.warns(
-            UserWarning, match="Annotations will be ignored when using query-planning"
-        )
+    ctx = pytest.warns(
+        UserWarning, match="Annotations will be ignored when using query-planning"
+    )
 
     with dask.annotate(retries=retries), ctx:
         df = acol + bcol
@@ -189,7 +185,3 @@ async def test_dataframe_annotations(c, s, a, b):
 
     assert rdf.dtypes == np.float64
     assert (rdf == 10.0).all()
-
-    if not dd._dask_expr_enabled():
-        # There is an annotation match per partition (i.e. task)
-        assert plugin.retry_matches == df.npartitions
diff --git a/distributed/shuffle/__init__.py b/distributed/shuffle/__init__.py
@@ -1,18 +1,12 @@
 from __future__ import annotations
 
 from distributed.shuffle._arrow import check_minimal_arrow_version
-from distributed.shuffle._merge import HashJoinP2PLayer, hash_join_p2p
 from distributed.shuffle._rechunk import rechunk_p2p
 from distributed.shuffle._scheduler_plugin import ShuffleSchedulerPlugin
-from distributed.shuffle._shuffle import P2PShuffleLayer, rearrange_by_column_p2p
 from distributed.shuffle._worker_plugin import ShuffleWorkerPlugin
 
 __all__ = [
     "check_minimal_arrow_version",
-    "hash_join_p2p",
-    "HashJoinP2PLayer",
-    "P2PShuffleLayer",
-    "rearrange_by_column_p2p",
     "rechunk_p2p",
     "ShuffleSchedulerPlugin",
     "ShuffleWorkerPlugin",