Remove recursion in task spec

.github/workflows/dask_test.yaml

@@ -47,10 +47,10 @@ jobs:
           cache-environment-key: environment-${{ steps.date.outputs.date }}-0
 
       - name: Install current main versions of dask
-        run: python -m pip install git+https://github.com/dask/dask
+        run: python -m pip install git+https://github.com/fjetter/dask@task_spec_remove_recursion
 
       - name: Install current main versions of distributed
-        run: python -m pip install git+https://github.com/dask/distributed
+        run: python -m pip install git+https://github.com/fjetter/distributed@task_spec_remove_recursion
 
       - name: Install dask-expr
         run: python -m pip install -e . --no-deps

.github/workflows/test.yaml

@@ -47,11 +47,11 @@ jobs:
           cache-environment-key: environment-${{ steps.date.outputs.date }}-1
 
       - name: Install current main versions of dask
-        run: python -m pip install git+https://github.com/dask/dask
+        run: python -m pip install git+https://github.com/fjetter/dask@task_spec_remove_recursion
         if: ${{ matrix.environment-file == 'ci/environment.yml' }}
 
       - name: Install current main versions of distributed
-        run: python -m pip install git+https://github.com/dask/distributed
+        run: python -m pip install git+https://github.com/fjetter/distributed@task_spec_remove_recursion
         if: ${{ matrix.environment-file == 'ci/environment.yml' }}
 
       - name: Install dask-expr

ci/environment.yml

@@ -12,5 +12,5 @@ dependencies:
   - sqlalchemy
   - xarray
   - pip:
-      - git+https://github.com/dask/distributed
-      - git+https://github.com/dask/dask
+      - git+https://github.com/fjetter/distributed@task_spec_remove_recursion
+      - git+https://github.com/fjetter/dask@task_spec_remove_recursion

dask_expr/_merge.py

@@ -680,7 +680,7 @@ def _layer(self) -> dict:
             _barrier_key_left,
             p2p_barrier,
             token_left,
-            transfer_keys_left,
+            *transfer_keys_left,
             spec=DataFrameShuffleSpec(
                 id=token_left,
                 npartitions=self.npartitions,
@@ -698,7 +698,7 @@ def _layer(self) -> dict:
             _barrier_key_right,
             p2p_barrier,
             token_right,
-            transfer_keys_right,
+            *transfer_keys_right,
             spec=DataFrameShuffleSpec(
                 id=token_right,
                 npartitions=self.npartitions,

dask_expr/_shuffle.py

@@ -592,7 +592,7 @@ def _layer(self):
             _barrier_key,
             p2p_barrier,
             token,
-            transfer_keys,
+            *transfer_keys,
             spec=DataFrameShuffleSpec(
                 id=shuffle_id,
                 npartitions=self.npartitions_out,

dask_expr/io/io.py

@@ -6,7 +6,7 @@
 
 import numpy as np
 import pyarrow as pa
-from dask._task_spec import Task
+from dask._task_spec import List, Task
 from dask.dataframe import methods
 from dask.dataframe._pyarrow import to_pyarrow_string
 from dask.dataframe.core import apply_and_enforce, is_dataframe_like, make_meta
@@ -135,7 +135,7 @@ def _task(self, name: Key, index: int) -> Task:
         bucket = self._fusion_buckets[index]
         # FIXME: This will likely require a wrapper
         return Task(
-            name, methods.concat, [expr._filtered_task(name, i) for i in bucket]
+            name, methods.concat, List(*(expr._filtered_task(name, i) for i in bucket))
         )
 
     @functools.cached_property

dask_expr/io/tests/test_distributed.py

@@ -63,4 +63,4 @@ def test_pickle_size(tmpdir, filesystem):
     df = read_parquet(tmpdir, filesystem=filesystem)
     from distributed.protocol import dumps
 
-    assert len(b"".join(dumps(df.optimize().dask))) <= 9000
+    assert len(b"".join(dumps(df.optimize().dask))) <= 9100