Merge branch 'dask:main' into main

relativistic · web-flow · commit 4d45cdfa70ae · 2025-01-09T08:40:07.000-07:00
diff --git a/.github/workflows/conda.yml b/.github/workflows/conda.yml
@@ -50,11 +50,9 @@ jobs:
 
           # conda search for the latest dask-core pre-release
           dask_core_arr=($(conda search --override-channels -c dask/label/dev dask-core | tail -n 1))
-          dask_expr_arr=($(conda search --override-channels -c dask/label/dev dask-expr | tail -n 1))
 
-          # extract dask-core & dask-expr pre-release versions
+          # extract dask-core pre-release version
           export DASK_CORE_VERSION=${dask_core_arr[1]}
-          export DASK_EXPR_VERSION=${dask_expr_arr[1]}
 
           # distributed pre-release build
           conda mambabuild continuous_integration/recipes/distributed \
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
@@ -73,12 +73,12 @@ jobs:
           - os: ubuntu-latest
             environment: mindeps
             label: pandas
-            extra_packages: [numpy=1.24, pandas=2.0, pyarrow=14.0.1, dask-expr]
+            extra_packages: [numpy=1.24, pandas=2.0, pyarrow=14.0.1]
             partition: "ci1"
           - os: ubuntu-latest
             environment: mindeps
             label: pandas
-            extra_packages: [numpy=1.24, pandas=2.0, pyarrow=14.0.1, dask-expr]
+            extra_packages: [numpy=1.24, pandas=2.0, pyarrow=14.0.1]
             partition: "not ci1"
 
           - os: ubuntu-latest
diff --git a/continuous_integration/environment-3.10.yaml b/continuous_integration/environment-3.10.yaml
@@ -57,7 +57,6 @@ dependencies:
   - httpx<0.28.0
   - pip:
       - git+https://github.com/dask/dask
-      - git+https://github.com/dask/dask-expr
       - git+https://github.com/dask/zict
       - git+https://github.com/dask/crick  # Only tested here
       # Revert after https://github.com/dask/distributed/issues/8614 is fixed
diff --git a/continuous_integration/environment-3.11.yaml b/continuous_integration/environment-3.11.yaml
@@ -51,7 +51,6 @@ dependencies:
   - httpx<0.28.0
   - pip:
       - git+https://github.com/dask/dask
-      - git+https://github.com/dask/dask-expr
       - git+https://github.com/dask/zict
       # Revert after https://github.com/dask/distributed/issues/8614 is fixed
       # - git+https://github.com/dask/s3fs
diff --git a/continuous_integration/environment-3.12.yaml b/continuous_integration/environment-3.12.yaml
@@ -51,7 +51,6 @@ dependencies:
   - httpx<0.28.0
   - pip:
       - git+https://github.com/dask/dask
-      - git+https://github.com/dask/dask-expr
       - git+https://github.com/dask/zict
       # Revert after https://github.com/dask/distributed/issues/8614 is fixed
       # - git+https://github.com/dask/s3fs
diff --git a/continuous_integration/environment-3.13.yaml b/continuous_integration/environment-3.13.yaml
@@ -50,7 +50,6 @@ dependencies:
   - httpx<0.28.0
   - pip:
       - git+https://github.com/dask/dask
-      - git+https://github.com/dask/dask-expr
       - git+https://github.com/dask/zict
       # Revert after https://github.com/dask/distributed/issues/8614 is fixed
       # - git+https://github.com/dask/s3fs
diff --git a/continuous_integration/recipes/dask/meta.yaml b/continuous_integration/recipes/dask/meta.yaml
@@ -2,7 +2,6 @@
 {% set new_patch = major_minor_patch[2] | int + 1 %}
 {% set version = (major_minor_patch[:2] + [new_patch]) | join('.') + environ.get('VERSION_SUFFIX', '') %}
 {% set dask_version = environ.get('DASK_CORE_VERSION', '0.0.0.dev') %}
-{% set dask_expr_version = environ.get('DASK_EXPR_VERSION', '0.0.0.dev') %}
 
 
 package:
@@ -21,12 +20,10 @@ requirements:
   host:
     - python >=3.10
     - dask-core {{ dask_version }}
-    - dask-expr {{ dask_expr_version }}
     - distributed {{ version }}
   run:
     - python >=3.10
     - {{ pin_compatible('dask-core', max_pin='x.x.x.x') }}
-    - {{ pin_compatible('dask-expr', max_pin='x.x.x.x') }}
     - {{ pin_compatible('distributed', exact=True) }}
     - cytoolz >=0.11.2
     - lz4 >=4.3.2
diff --git a/continuous_integration/scripts/test_report.py b/continuous_integration/scripts/test_report.py
@@ -142,7 +142,7 @@ def get_jobs(run, session, repo):
     # this to the JXML results.
 
     if repo.endswith("/dask"):
-        # example name: "test (windows-latest, 3.9)" or "test (ubuntu-latest, 3.12, dask-expr)"
+        # example name: "test (windows-latest, 3.9)" or "test (ubuntu-latest, 3.12)"
         df_jobs = df_jobs[~df_jobs.name.str.contains("Event File")]
 
         def format(row):
@@ -389,10 +389,6 @@ def download_and_parse_artifacts(
                 if a["name"].startswith("test-results") and repo.endswith("/dask"):
                     continue
 
-                # NOTE: Temporarily ignore reporting dask-expr related test cases
-                if "dask-expr" in a["name"]:
-                    continue
-
                 # Note: we assign a column with the workflow run timestamp rather
                 # than the artifact timestamp so that artifacts triggered under
                 # the same workflow run can be aligned according to the same trigger
diff --git a/distributed/client.py b/distributed/client.py
@@ -45,7 +45,6 @@
 from dask.core import flatten, validate_key
 from dask.highlevelgraph import HighLevelGraph
 from dask.layers import Layer
-from dask.optimization import SubgraphCallable
 from dask.tokenize import tokenize
 from dask.typing import Key, NestedKeys, NoDefault, no_default
 from dask.utils import (
@@ -6120,8 +6119,6 @@ def futures_of(o, client=None):
             stack.extend(x)
         elif type(x) is dict:
             stack.extend(x.values())
-        elif type(x) is SubgraphCallable:
-            stack.extend(x.dsk.values())
         elif isinstance(x, TaskRef):
             if x not in seen:
                 seen.add(x)
diff --git a/distributed/shuffle/tests/test_merge.py b/distributed/shuffle/tests/test_merge.py
@@ -57,9 +57,9 @@ async def test_basic_merge(c, s, a, b, how):
     joined = a.merge(b, left_on="y", right_on="y", how=how)
 
     # Ensure we're using a hash join
-    from dask_expr._merge import HashJoinP2P
-
-    assert any(isinstance(expr, HashJoinP2P) for expr in joined.optimize()._expr.walk())
+    assert any(
+        expr._name.startswith("hashjoinp2p") for expr in joined.optimize()._expr.walk()
+    )
 
     expected = pd.merge(A, B, how, "y")
     await list_eq(joined, expected)
diff --git a/distributed/shuffle/tests/test_shuffle.py b/distributed/shuffle/tests/test_shuffle.py
@@ -1637,33 +1637,6 @@ async def test_multi(c, s, a, b):
     await assert_scheduler_cleanup(s)
 
 
-@pytest.mark.skipif(reason="worker restrictions are not supported in dask-expr")
-@gen_cluster(client=True)
-async def test_restrictions(c, s, a, b):
-    df = dask.datasets.timeseries(
-        start="2000-01-01",
-        end="2000-01-10",
-        dtypes={"x": float, "y": float},
-        freq="10 s",
-    ).persist(workers=a.address)
-    await df
-    assert a.data
-    assert not b.data
-
-    with dask.config.set({"dataframe.shuffle.method": "p2p"}):
-        x = df.shuffle("x")
-        y = df.shuffle("y")
-
-    x = x.persist(workers=b.address)
-    y = y.persist(workers=a.address)
-
-    await x
-    assert all(key in b.data for key in x.__dask_keys__())
-
-    await y
-    assert all(key in a.data for key in y.__dask_keys__())
-
-
 @gen_cluster(client=True)
 async def test_delete_some_results(c, s, a, b):
     df = dask.datasets.timeseries(
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
@@ -42,7 +42,6 @@
 import dask
 import dask.bag as db
 from dask import delayed
-from dask.optimization import SubgraphCallable
 from dask.tokenize import tokenize
 from dask.utils import get_default_shuffle_method, parse_timedelta, tmpfile
 
@@ -2626,13 +2625,6 @@ async def test_futures_of_get(c, s, a, b):
     b = db.Bag({("b", i): f for i, f in enumerate([x, y, z])}, "b", 3)
     assert set(futures_of(b)) == {x, y, z}
 
-    sg = SubgraphCallable(
-        {"x": x, "y": y, "z": z, "out": (add, (add, (add, x, y), z), "in")},
-        "out",
-        ("in",),
-    )
-    assert set(futures_of(sg)) == {x, y, z}
-
 
 def test_futures_of_class():
     pytest.importorskip("numpy")
@@ -6191,43 +6183,6 @@ async def test_profile_bokeh(c, s, a, b):
         assert os.path.exists(fn)
 
 
-@gen_cluster(client=True, nthreads=[("", 1)])
-async def test_get_mix_futures_and_SubgraphCallable(c, s, a):
-    future = c.submit(add, 1, 2)
-
-    subgraph = SubgraphCallable(
-        {"_2": (add, "_0", "_1"), "_3": (add, future, "_2")},
-        "_3",
-        ("_0", "_1"),
-    )
-    dsk = {
-        "a": 1,
-        "b": 2,
-        "c": (subgraph, "a", "b"),
-        "d": (subgraph, "c", "b"),
-    }
-
-    future2 = c.get(dsk, "d", sync=False)
-    result = await future2
-    assert result == 11
-
-    # Nested subgraphs
-    subgraph2 = SubgraphCallable(
-        {
-            "_2": (subgraph, "_0", "_1"),
-            "_3": (subgraph, "_2", "_1"),
-            "_4": (add, "_3", future2),
-        },
-        "_4",
-        ("_0", "_1"),
-    )
-
-    dsk2 = {"e": 1, "f": 2, "g": (subgraph2, "e", "f")}
-
-    result = await c.get(dsk2, "g", sync=False)
-    assert result == 22
-
-
 @gen_cluster(client=True)
 async def test_get_mix_futures_and_SubgraphCallable_dask_dataframe(c, s, a, b):
     pd = pytest.importorskip("pandas")
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
@@ -2829,9 +2829,7 @@ async def test_default_task_duration_splits(c, s, a, b):
     await wait(fut)
 
     split_prefix = [pre for pre in s.task_prefixes.keys() if "split" in pre]
-    # dask-expr enabled: ['split-taskshuffle', 'split-stage']
-    # dask-expr disabled: ['split-shuffle']
-    assert split_prefix
+    assert split_prefix == ["split-taskshuffle", "split-stage"]
     default_times = dask.config.get("distributed.scheduler.default-task-durations")
     for p in split_prefix:
         default_time = parse_timedelta(default_times[p])
diff --git a/distributed/tests/test_spans.py b/distributed/tests/test_spans.py
@@ -803,7 +803,7 @@ def f():
 
 @gen_cluster(client=True)
 async def test_span_on_persist(c, s, a, b):
-    """As a workaround to lack of annotations support in dask-expr and loss of
+    """As a workaround to lack of annotations support in dask.dataframe and loss of
     annotations due to low level optimization in dask.array, you can use span() to wrap
     calls to persist() and compute()
     """
diff --git a/distributed/tests/test_utils_comm.py b/distributed/tests/test_utils_comm.py
@@ -7,7 +7,6 @@
 import pytest
 
 from dask._task_spec import TaskRef
-from dask.optimization import SubgraphCallable
 
 from distributed import wait
 from distributed.compatibility import asyncio_run
@@ -246,18 +245,3 @@ def assert_eq(keys1: set[TaskRef], keys2: set[TaskRef]) -> None:
     res, keys = unpack_remotedata(TaskRef("mykey"))
     assert res == "mykey"
     assert_eq(keys, {TaskRef("mykey")})
-
-    # Check unpack of SC that contains a wrapped key
-    sc = SubgraphCallable({"key": (TaskRef("data"),)}, outkey="key", inkeys=["arg1"])
-    dsk = (sc, "arg1")
-    res, keys = unpack_remotedata(dsk)
-    assert res[0] != sc  # Notice, the first item (the SC) has been changed
-    assert res[1:] == ("arg1", "data")
-    assert_eq(keys, {TaskRef("data")})
-
-    # Check unpack of SC when it takes a wrapped key as argument
-    sc = SubgraphCallable({"key": ("arg1",)}, outkey="key", inkeys=[TaskRef("arg1")])
-    dsk = (sc, "arg1")
-    res, keys = unpack_remotedata(dsk)
-    assert res == (sc, "arg1")  # Notice, the first item (the SC) has NOT been changed
-    assert_eq(keys, set())
diff --git a/distributed/utils_comm.py b/distributed/utils_comm.py
@@ -13,7 +13,6 @@
 
 import dask.config
 from dask._task_spec import TaskRef
-from dask.optimization import SubgraphCallable
 from dask.typing import Key
 from dask.utils import is_namedtuple_instance, parse_timedelta
 
@@ -197,30 +196,6 @@ def _unpack_remotedata_inner(
     if typ is tuple:
         if not o:
             return o
-        if type(o[0]) is SubgraphCallable:
-            # Unpack futures within the arguments of the subgraph callable
-            futures: set[TaskRef] = set()
-            args = tuple(_unpack_remotedata_inner(i, byte_keys, futures) for i in o[1:])
-            found_futures.update(futures)
-
-            # Unpack futures within the subgraph callable itself
-            sc: SubgraphCallable = o[0]
-            futures = set()
-            dsk = {
-                k: _unpack_remotedata_inner(v, byte_keys, futures)
-                for k, v in sc.dsk.items()
-            }
-            future_keys: tuple = ()
-            if futures:  # If no futures is in the subgraph, we just use `sc` as-is
-                found_futures.update(futures)
-                future_keys = (
-                    tuple(f.key for f in futures)
-                    if byte_keys
-                    else tuple(f.key for f in futures)
-                )
-                inkeys = tuple(sc.inkeys) + future_keys
-                sc = SubgraphCallable(dsk, sc.outkey, inkeys, sc.name)
-            return (sc,) + args + future_keys
         else:
             return tuple(
                 _unpack_remotedata_inner(item, byte_keys, found_futures) for item in o
