Repository: coiled/benchmarks (language: python)

Issue:
Toggle dask-expr for all benchmarks
dask/dask CI looks good and we're preparing to switch over to dask-expr as the dataframe backend soon. This removes the need to watch "legacy dask" very closely. I don't expect anything major to slow down dask/dask (that wouldn't break dask-expr in a similar way).

We should set the dask-expr toggle for our benchmark suite to reflect this (possibly run one A/B test before we switch to verify our baseline)


Discussion / hints:

Fix patch:
diff --git a/AB_environments/AB_baseline.dask.yaml b/AB_environments/AB_baseline.dask.yaml
index 4edfd74109..5afc7be70c 100644
--- a/AB_environments/AB_baseline.dask.yaml
+++ b/AB_environments/AB_baseline.dask.yaml
@@ -2,3 +2,6 @@
 # (overriding the built-in config) for the baseline environment.
 # Change contents, but do not rename.
 # Leave empty if you don't want to override anything.
+dask:
+  dataframe:
+    query-planning: True
diff --git a/AB_environments/AB_sample.dask.yaml b/AB_environments/AB_sample.dask.yaml
index baa367a5f5..9ae43cb8fb 100644
--- a/AB_environments/AB_sample.dask.yaml
+++ b/AB_environments/AB_sample.dask.yaml
@@ -11,3 +11,6 @@
 # distributed:
 #   scheduler:
 #     worker-saturation: 1.2
+dask:
+  dataframe:
+    query-planning: True


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 8b85a1c421..2b65a86c25 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -141,6 +141,7 @@ jobs:
           COILED_RUNTIME_VERSION: ${{ matrix.runtime-version }}
           DB_NAME: ${{ matrix.name_prefix }}-${{ matrix.os }}-py${{ matrix.python_version }}.db
           CLUSTER_DUMP: always
+          DASK_DATAFRAME__QUERY_PLANNING: True
         run: |
           pytest --benchmark -n 4 --dist loadscope ${{ env.PYTEST_MARKERS }} ${{ matrix.pytest_args }}
 
diff --git a/tests/tpch/dask_queries.py b/tests/tpch/dask_queries.py
index 67d772d881..a1973c0025 100644
--- a/tests/tpch/dask_queries.py
+++ b/tests/tpch/dask_queries.py
@@ -1,6 +1,6 @@
 from datetime import datetime, timedelta
 
-import dask_expr as dd
+import dask.dataframe as dd
 
 
 def query_1(dataset_path, fs):
@@ -360,6 +360,7 @@ def query_8(dataset_path, fs):
     final = mkt_total.merge(
         mkt_brazil, left_on="o_year", right_on="o_year", suffixes=("_mkt", "_brazil")
     )
+
     final["mkt_share"] = final.volume_brazil / final.volume_mkt
     return final.sort_values(by=["o_year"], ascending=[True])[["o_year", "mkt_share"]]
 
@@ -494,6 +495,7 @@ def query_10(dataset_path, fs):
     )
 
     # TODO: ideally the filters are pushed up before the merge during optimization
+
     # query = query[
     #     (query.o_orderdate >= orderdate_from)
     #     & (query.o_orderdate < orderdate_to)
diff --git a/tests/tpch/test_dask.py b/tests/tpch/test_dask.py
index 48d2f7d8ba..160033d495 100644
--- a/tests/tpch/test_dask.py
+++ b/tests/tpch/test_dask.py
@@ -4,7 +4,7 @@
 
 pytestmark = pytest.mark.tpch_dask
 
-dd = pytest.importorskip("dask_expr")
+dd = pytest.importorskip("dask.dataframe")
 
 
 def test_query_1(client, dataset_path, fs):