add temporary serialization changes

python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py

@@ -124,6 +124,18 @@ def __init__(
         "linear": plc.types.Interpolation.LINEAR,
     }
 
+    def __reduce__(self):
+        """Pickle an Agg object."""
+        return (
+            type(self),
+            (
+                self.dtype,
+                self.name,
+                self.options,
+                *self.children,
+            ),
+        )
+
     def collect_agg(self, *, depth: int) -> AggInfo:
         """Collect information about aggregations in groupbys."""
         if depth >= 1:

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -840,6 +840,21 @@ def do_evaluate(
         return DataFrame(columns)
 
 
+class AggInfoWrapper:
+    """Wrap an GropuBy AggInfo sequence."""
+
+    agg_requests: Sequence[expr.NamedExpr]
+    agg_infos: Sequence[expr.AggInfo]
+
+    def __init__(self, agg_requests: Sequence[expr.NamedExpr]):
+        self.agg_requests = tuple(agg_requests)
+        self.agg_infos = [req.collect_agg(depth=0) for req in self.agg_requests]
+
+    def __reduce__(self):
+        """Pickle an AggInfoWrapper object."""
+        return (AggInfoWrapper, (self.agg_requests,))
+
+
 class GroupBy(IR):
     """Perform a groupby."""
 
@@ -883,13 +898,15 @@ def __init__(
             raise NotImplementedError("dynamic group by")
         if any(GroupBy.check_agg(a.value) > 1 for a in self.agg_requests):
             raise NotImplementedError("Nested aggregations in groupby")
-        self.agg_infos = [req.collect_agg(depth=0) for req in self.agg_requests]
+        # self.agg_infos = [req.collect_agg(depth=0) for req in self.agg_requests]
         self._non_child_args = (
             self.keys,
             self.agg_requests,
             maintain_order,
-            options,
-            self.agg_infos,
+            # options,
+            # self.agg_infos,
+            {"slice": options.slice},
+            AggInfoWrapper(self.agg_requests),
         )
 
     @staticmethod
@@ -927,7 +944,8 @@ def do_evaluate(
         agg_requests: Sequence[expr.NamedExpr],
         maintain_order: bool,  # noqa: FBT001
         options: Any,
-        agg_infos: Sequence[expr.AggInfo],
+        # agg_infos: Sequence[expr.AggInfo],
+        agg_info_wrapper: AggInfoWrapper,
         df: DataFrame,
     ):
         """Evaluate and return a dataframe."""
@@ -947,7 +965,7 @@ def do_evaluate(
         # TODO: uniquify
         requests = []
         replacements: list[expr.Expr] = []
-        for info in agg_infos:
+        for info in agg_info_wrapper.agg_infos:
             for pre_eval, req, rep in info.requests:
                 if pre_eval is None:
                     # A count aggregation, doesn't touch the column,
@@ -1012,7 +1030,7 @@ def do_evaluate(
                     ordered_table.columns(), broadcasted, strict=True
                 )
             ]
-        return DataFrame(broadcasted).slice(options.slice)
+        return DataFrame(broadcasted).slice(options.get("slice"))
 
 
 class ConditionalJoin(IR):

python/cudf_polars/cudf_polars/experimental/benchmarks/tpch.py

@@ -41,7 +41,7 @@
     "--executor",
     default="dask",
     type=str,
-    choices=["dask", "dask-experimental", "pylibcudf", "polars"],
+    choices=["dask", "dask-cuda", "dask-experimental", "pylibcudf", "polars"],
     help="Executor.",
 )
 parser.add_argument(
@@ -216,6 +216,23 @@ def q18(args):
 
 def run(args):
     """Run the benchmark once."""
+    executor = args.executor
+
+    if executor == "dask-cuda" and not args.debug:
+        from dask_cuda import LocalCUDACluster
+        from distributed import Client
+
+        client = Client(
+            LocalCUDACluster(
+                n_workers=1,
+                dashboard_address=":8585",
+                rmm_pool_size=0.9,
+                jit_unspill=True,
+            )
+        )
+    else:
+        client = None
+
     t0 = time.time()
 
     q_id = args.query
@@ -230,7 +247,6 @@ def run(args):
     else:
         raise NotImplementedError(f"Query {q_id} not implemented.")
 
-    executor = args.executor
     if executor == "polars":
         result = q.collect()
     else:
@@ -256,6 +272,9 @@ def run(args):
     print(result)
     print(f"time is {t1-t0}")
 
+    if client is not None:
+        client.close()
+
 
 if __name__ == "__main__":
     run(args)

python/cudf_polars/cudf_polars/experimental/parallel.py

@@ -128,9 +128,28 @@ def task_graph(
         return graph, (key_name, 0)
 
 
+def _register_serialize():
+    """Register Dask/cudf-polars serializers in calling process."""
+    from cudf_polars.experimental.dask_serialize import register
+
+    register()
+
+
 def evaluate_dask(ir: IR) -> DataFrame:
     """Evaluate an IR graph with Dask."""
-    from dask import get
+    from dask import get as _get
+    from distributed import get_client
+
+    _register_serialize()
+
+    try:
+        client = get_client()
+        client.run(_register_serialize)
+        client.run_on_scheduler(_register_serialize)
+
+        get = client.get
+    except ValueError:
+        get = _get
 
     ir, partition_info = lower_ir_graph(ir)