dask-polars: use splat everywhere. (#18243)

In the dask-graph, use the splat operator instead of list arguments everywhere. The splat operator is used in most places already, only a few functions are missing.
This makes it easier to transform the dask-graph after its creation, which will be useful for prototyping.

Authors:
  - Mads R. B. Kristensen (https://github.com/madsbk)

Approvers:
  - Lawrence Mitchell (https://github.com/wence-)

URL: #18243

Author: madsbk

python/cudf_polars/cudf_polars/experimental/base.py

@@ -9,7 +9,7 @@
 from cudf_polars.dsl.ir import Union
 
 if TYPE_CHECKING:
-    from collections.abc import Iterator, Sequence
+    from collections.abc import Iterator
 
     from cudf_polars.containers import DataFrame
     from cudf_polars.dsl.expr import NamedExpr
@@ -44,6 +44,6 @@ def get_key_name(node: Node) -> str:
     return f"{type(node).__name__.lower()}-{hash(node)}"
 
 
-def _concat(dfs: Sequence[DataFrame]) -> DataFrame:
+def _concat(*dfs: DataFrame) -> DataFrame:
     # Concatenate a sequence of DataFrames vertically
     return Union.do_evaluate(None, *dfs)

python/cudf_polars/cudf_polars/experimental/groupby.py

@@ -251,8 +251,8 @@ def _(
     return new_node, partition_info
 
 
-def _tree_node(do_evaluate, batch, *args):
-    return do_evaluate(*args, _concat(batch))
+def _tree_node(do_evaluate, nbatch, *args):
+    return do_evaluate(*args[nbatch:], _concat(*args[:nbatch]))
 
 
 @generate_ir_tasks.register(GroupBy)
@@ -289,7 +289,8 @@ def _(
             graph[(name, j, i)] = (
                 _tree_node,
                 ir.do_evaluate,
-                batch,
+                len(batch),
+                *batch,
                 *ir._non_child_args,
             )
             new_keys.append((name, j, i))
@@ -298,7 +299,8 @@ def _(
     graph[(name, 0)] = (
         _tree_node,
         ir.do_evaluate,
-        keys,
+        len(keys),
+        *keys,
         *ir._non_child_args,
     )
     return graph

python/cudf_polars/cudf_polars/experimental/join.py

@@ -309,6 +309,6 @@ def _(
             if len(_concat_list) == 1:
                 graph[(out_name, part_out)] = graph.pop(_concat_list[0])
             else:
-                graph[(out_name, part_out)] = (_concat, _concat_list)
+                graph[(out_name, part_out)] = (_concat, *_concat_list)
 
         return graph

python/cudf_polars/cudf_polars/experimental/parallel.py

@@ -150,7 +150,7 @@ def task_graph(
     key_name = get_key_name(ir)
     partition_count = partition_info[ir].count
     if partition_count > 1:
-        graph[key_name] = (_concat, list(partition_info[ir].keys(ir)))
+        graph[key_name] = (_concat, *partition_info[ir].keys(ir))
         return graph, key_name
     else:
         return graph, (key_name, 0)

python/cudf_polars/cudf_polars/experimental/shuffle.py

@@ -159,7 +159,7 @@ def _simple_shuffle_graph(
                 (split_name, part_in),
                 part_out,
             )
-        graph[(name_out, part_out)] = (_concat, _concat_list)
+        graph[(name_out, part_out)] = (_concat, *_concat_list)
     return graph