Move replicated-output dedup to the Dask and Ray frontends

python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py

@@ -386,9 +386,8 @@ def execute_ir_on_rank(
     stats: StatsCollector,
     collective_id_map: dict[IR, list[int]],
     *,
-    collect_metadata: bool = False,
     query_id: uuid.UUID,
-) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]:
+) -> tuple[pl.DataFrame, list[ChannelMetadata]]:
     """
     Execute a Polars IR query on a single rank's GPU.
 
@@ -414,8 +413,6 @@ def execute_ir_on_rank(
         Statistics collector.
     collective_id_map
         Mapping from IR nodes to their pre-allocated collective operation IDs.
-    collect_metadata
-        Whether to collect channel metadata during execution.
     query_id
         Unique identifier for the query, propagated into actor traces.
 
@@ -424,13 +421,12 @@ def execute_ir_on_rank(
     result
         This rank's output fragment as a Polars DataFrame.
     metadata
-        Collected channel metadata if ``collect_metadata`` is ``True``,
-        otherwise ``None``.
+        Collected channel metadata.
     """
     ir_context = IRExecutionContext(
         get_cuda_stream=ctx.get_stream_from_pool, query_id=query_id
     )
-    metadata_collector: list[ChannelMetadata] | None = [] if collect_metadata else None
+    metadata_collector: list[ChannelMetadata] = []
 
     nodes, output = generate_network(
         ctx,
@@ -606,9 +602,8 @@ def evaluate_on_rank(
     ir: IR,
     config_options: ConfigOptions[StreamingExecutor],
     *,
-    collect_metadata: bool = False,
     query_id: uuid.UUID,
-) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]:
+) -> tuple[pl.DataFrame, list[ChannelMetadata]]:
     """
     Evaluate a polars IR plan on a single rank.
 
@@ -632,8 +627,6 @@ def evaluate_on_rank(
         Root of the **pre-lowered** IR graph.
     config_options
         Executor configuration forwarded from the client.
-    collect_metadata
-        Whether to collect channel metadata during execution.
     query_id
         Unique identifier for the query, propagated into actor traces.
 
@@ -642,8 +635,7 @@ def evaluate_on_rank(
     result
         This rank's output fragment as a Polars DataFrame.
     metadata
-        Collected channel metadata if *collect_metadata* is ``True``,
-        otherwise ``None``.
+        Collected channel metadata.
     """
     stats = allgather_stats(comm, ctx.br(), ir, config_options)
     ir, partition_info = lower_ir_graph(ir, config_options, stats)
@@ -663,6 +655,5 @@ def evaluate_on_rank(
             config_options,
             stats,
             collective_id_map,
-            collect_metadata=collect_metadata,
             query_id=query_id,
         )

python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/dask.py

@@ -411,15 +411,24 @@ def _worker_evaluate(
     mp_ctx: _WorkerContext = getattr(dask_worker, f"_cudf_polars_mp_context_{uid}")
     if mp_ctx.ctx is None or mp_ctx.comm is None or mp_ctx.py_executor is None:
         raise RuntimeError("_setup_worker must be called before _worker_evaluate")
-    return evaluate_on_rank(
+    # evaluate_on_rank always collects metadata internally so we can read
+    # metadata[-1].duplicated to decide whether to suppress this rank's output.
+    # The client concatenates each rank's result, so without this dedup an
+    # output marked duplicated=True would appear N times. The external
+    # collect_metadata parameter still controls whether the collected list is
+    # returned to the client (see the return statement), which is the cost we
+    # care about saving when the caller doesn't need the metadata.
+    df, metadata = evaluate_on_rank(
         mp_ctx.ctx,
         mp_ctx.comm,
         mp_ctx.py_executor,
         ir,
         config_options,
-        collect_metadata=collect_metadata,
         query_id=query_id,
     )
+    if mp_ctx.comm.rank != 0 and metadata and metadata[-1].duplicated:
+        df = df.clear()
+    return df, metadata if collect_metadata else None
 
 
 def evaluate_pipeline_dask_mode(

python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py

@@ -368,21 +368,32 @@ def evaluate_polars_ir(
         RuntimeError
             If :meth:`setup_worker` has not been called first.
         """
-        if self._ctx is None:
+        if self._ctx is None or self._comm is None:
             raise RuntimeError("setup_worker must be called before evaluate_polars_ir")
         # Ray transfers the returned Polars DataFrame back to the client via the
         # object store (pickle / Arrow IPC). The DataFrame is already on CPU at
         # this point (to_polars() copies the result off-GPU), so no GPU memory
         # crosses process boundaries.
-        return evaluate_on_rank(
+        #
+        # evaluate_on_rank always collects metadata internally so we can read
+        # metadata[-1].duplicated to decide whether to suppress this rank's
+        # output. The client concatenates each rank's result, so without this
+        # dedup an output marked duplicated=True would appear N times. The
+        # external collect_metadata parameter still controls whether the
+        # collected list is returned to the client (see the return statement),
+        # which is the cost we care about saving when the caller doesn't need
+        # the metadata.
+        df, metadata = evaluate_on_rank(
             self._ctx,
             self._comm,
             self._py_executor,
             ir,
             config_options,
-            collect_metadata=collect_metadata,
             query_id=query_id,
         )
+        if self._comm.rank != 0 and metadata and metadata[-1].duplicated:
+            df = df.clear()
+        return df, metadata if collect_metadata else None
 
     def _run(self, func: Callable[..., T], *args: Any, **kwargs: Any) -> T:
         return func(*args, **kwargs)

python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py

@@ -108,15 +108,15 @@ def evaluate_pipeline_spmd_mode(
     context = config_options.executor.spmd_context.context
     py_executor = config_options.executor.spmd_context.py_executor
 
-    return evaluate_on_rank(
+    df, metadata = evaluate_on_rank(
         context,
         comm,
         py_executor,
         ir,
         config_options,
-        collect_metadata=collect_metadata,
         query_id=query_id,
     )
+    return df, metadata if collect_metadata else None
 
 
 def allgather_polars_dataframe(

python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py

@@ -207,7 +207,7 @@ async def _collect_small_side_for_broadcast(
         stream = ir_context.get_cuda_stream()
         gathered = await allgather.extract_concatenated(stream)
         # When every rank inserted zero chunks, the AllGather has no schema
-        # to infer and returns a 0-column table. Substitute a properly typed
+        # to infer and returns a 0 column table. Substitute a properly typed
         # empty table for the small side so downstream joins still match the
         # expected schema.
         table = (

python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py

@@ -732,20 +732,10 @@ async def metadata_drain_node(
     ):
         # Drain metadata channel (we don't need it after this point)
         metadata = await recv_metadata(ch_in, context)
-        send_empty = metadata.duplicated and comm.rank != 0
         if metadata_collector is not None:
             metadata_collector.append(metadata)
 
-        # Forward non-duplicated data messages
         while (msg := await ch_in.recv(context)) is not None:
-            if not send_empty:
-                await ch_out.send(context, msg)
-
-        # Send empty data if needed
-        if send_empty:
-            stream = ir_context.get_cuda_stream()
-            await ch_out.send(
-                context, Message(0, empty_table_chunk(ir, context, stream))
-            )
+            await ch_out.send(context, msg)
 
         await ch_out.drain(context)

python/cudf_polars/cudf_polars/experimental/rapidsmpf/union.py

@@ -17,6 +17,7 @@
 from cudf_polars.experimental.rapidsmpf.nodes import define_actor, shutdown_on_error
 from cudf_polars.experimental.rapidsmpf.utils import (
     ChannelManager,
+    empty_table_chunk,
     gather_in_task_group,
     process_children,
     recv_metadata,
@@ -49,8 +50,7 @@ async def union_node(
     context
         The rapidsmpf context.
     comm
-        The communicator. Used to suppress duplicated children's chunks on
-        non-root ranks so they aren't emitted twice cluster-wide.
+        The communicator.
     ir
         The Union IR node.
     ir_context
@@ -69,15 +69,9 @@ async def union_node(
         metadata = await gather_in_task_group(
             *(recv_metadata(ch, context) for ch in chs_in)
         )
-        # When a child has duplicated=True, every rank has produced the same
-        # data and only rank 0 should forward it -- otherwise the downstream
-        # client-side concat would over-count by `nranks - 1` for each
-        # duplicated chunk.
-        skip = tuple(meta.duplicated and comm.rank != 0 for meta in metadata)
-        total_local_count = sum(
-            0 if drop else meta.local_count
-            for meta, drop in zip(metadata, skip, strict=True)
-        )
+        # Chunk counts on the wire are uniform across ranks, so report the
+        # full sum.
+        total_local_count = sum(meta.local_count for meta in metadata)
         duplicated = all(meta.duplicated for meta in metadata)
         await send_metadata(
             ch_out,
@@ -88,23 +82,29 @@ async def union_node(
             ),
         )
 
+        # When a child has duplicated=True, every rank has produced the same
+        # rows, so we drop them everywhere except rank 0 to avoid N counting.
+        suppress = tuple(meta.duplicated and comm.rank != 0 for meta in metadata)
+
         seq_num_offset = 0
-        for ch_in, drop in zip(chs_in, skip, strict=True):
+        for ch_in, drop in zip(chs_in, suppress, strict=True):
             num_ch_chunks = 0
             while (msg := await ch_in.recv(context)) is not None:
-                if not drop:
-                    await ch_out.send(
-                        context,
-                        Message(
-                            msg.sequence_number + seq_num_offset,
-                            TableChunk.from_message(
-                                msg, br=context.br()
-                            ).make_available_and_spill(
-                                context.br(), allow_overbooking=True
-                            ),
-                        ),
-                    )
-                    num_ch_chunks += 1
+                if drop:
+                    stream = ir_context.get_cuda_stream()
+                    out_chunk = empty_table_chunk(ir, context, stream)
+                else:
+                    out_chunk = TableChunk.from_message(
+                        msg, br=context.br()
+                    ).make_available_and_spill(context.br(), allow_overbooking=True)
+                await ch_out.send(
+                    context,
+                    Message(
+                        msg.sequence_number + seq_num_offset,
+                        out_chunk,
+                    ),
+                )
+                num_ch_chunks += 1
             seq_num_offset += num_ch_chunks
 
         await ch_out.drain(context)

python/cudf_polars/cudf_polars/experimental/select.py

@@ -432,8 +432,9 @@ def _(
         named_expr = expr.NamedExpr(ir.exprs[0].name or "len", lit_expr)
 
         # Use Empty as the input so the streaming network's metadata flows
-        # `duplicated=True` end-to-end. Without that, every rank emits the
-        # literal once and the client concatenates N copies.
+        # duplicated=True end to end. Without that, the literal expression
+        # would be evaluated chunkwise over child and emit one row per input
+        # chunk instead of a single row globally.
         input_ir: IR = Empty({})
         new_node = Select(
             {named_expr.name: named_expr.value.dtype},

python/cudf_polars/tests/experimental/test_spmd.py

@@ -22,6 +22,7 @@
     SPMDEngine,
     allgather_polars_dataframe,
 )
+from cudf_polars.testing.asserts import assert_gpu_result_equal
 from cudf_polars.utils.config import MemoryResourceConfig
 
 if TYPE_CHECKING:
@@ -294,6 +295,24 @@ def test_run(spmd_engine: SPMDEngine) -> None:
     assert result == [os.getpid()]
 
 
+def test_sort_slice_over_union_of_duplicated_streams(
+    spmd_engine: SPMDEngine,
+) -> None:
+    """Sort+head over a concat of two group-by branches returns the global result on every rank."""
+    lf1 = (
+        pl.LazyFrame({"name": ["alice"], "score": [1.0]})
+        .group_by("name")
+        .agg(pl.col("score").sum())
+    )
+    lf2 = (
+        pl.LazyFrame({"name": ["bob"], "score": [2.0]})
+        .group_by("name")
+        .agg(pl.col("score").sum())
+    )
+    lf = pl.concat([lf1, lf2]).sort("score").head(10)
+    assert_gpu_result_equal(lf, engine=spmd_engine, check_row_order=False)
+
+
 def test_reset_keeps_comm_alive(comm: Communicator) -> None:
     """``_reset`` must not rebuild the communicator."""
     with SPMDEngine(

python/cudf_polars/tests/experimental/test_union.py

@@ -17,3 +17,19 @@ def test_union_shared_fanout_no_deadlock(streaming_engine):
     project = df.select("key", "val")
     q = pl.concat([gb, project])
     assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False)
+
+
+def test_sort_slice_over_union_of_duplicated_streams(streaming_engine):
+    # Sort+head over a concat of two group-by branches.
+    lf1 = (
+        pl.LazyFrame({"name": ["alice"], "score": [1.0]})
+        .group_by("name")
+        .agg(pl.col("score").sum())
+    )
+    lf2 = (
+        pl.LazyFrame({"name": ["bob"], "score": [2.0]})
+        .group_by("name")
+        .agg(pl.col("score").sum())
+    )
+    q = pl.concat([lf1, lf2]).sort("score").head(10)
+    assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False)

python/cudf_polars/tests/test_groupby.py

@@ -484,13 +484,8 @@ def test_groupby_sum_decimal_null_group(engine: pl.GPUEngine) -> None:
 
 
 @pytest.mark.xfail(
-    raises=(AssertionError, pl.exceptions.SchemaError),
-    reason=(
-        "https://github.com/rapidsai/cudf/issues/19610 — in-memory engine "
-        "fails with AssertionError (wrong values); multi-rank streaming "
-        "fails earlier with SchemaError (literal agg yields a divergent "
-        "schema after cross-rank concat)."
-    ),
+    raises=AssertionError,
+    reason="https://github.com/rapidsai/cudf/issues/19610",
 )
 def test_groupby_literal_agg(engine: pl.GPUEngine):
     df = pl.LazyFrame({"c0": [True, False]})