Multi-rank fixes for cudf-polars streaming

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

@@ -205,9 +205,19 @@ async def _collect_small_side_for_broadcast(
             for s_id in range(len(chunks)):
                 inserter.insert(s_id, chunks.pop(0))
         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
+        # empty table for the small side so downstream joins still match the
+        # expected schema.
+        table = (
+            empty_table_chunk(ir, context, stream).table_view()
+            if gathered.num_columns() == 0 and len(ir.schema) > 0
+            else gathered
+        )
         dfs = [
             DataFrame.from_table(
-                await allgather.extract_concatenated(stream),
+                table,
                 list(ir.schema.keys()),
                 list(ir.schema.values()),
                 stream,

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

@@ -24,6 +24,7 @@
 )
 
 if TYPE_CHECKING:
+    from rapidsmpf.communicator.communicator import Communicator
     from rapidsmpf.streaming.core.channel import Channel
     from rapidsmpf.streaming.core.context import Context
 
@@ -34,6 +35,7 @@
 @define_actor()
 async def union_node(
     context: Context,
+    comm: Communicator,
     ir: Union,
     ir_context: IRExecutionContext,
     ch_out: Channel[TableChunk],
@@ -46,6 +48,9 @@ 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.
     ir
         The Union IR node.
     ir_context
@@ -61,14 +66,19 @@ async def union_node(
         # Merge and forward metadata.
         # Union loses partitioning/ordering info since sources may differ.
         # TODO: Warn users that Union does NOT preserve order?
-        total_local_count = 0
-        duplicated = True
         metadata = await gather_in_task_group(
             *(recv_metadata(ch, context) for ch in chs_in)
         )
-        for meta in metadata:
-            total_local_count += meta.local_count
-            duplicated = duplicated and meta.duplicated
+        # 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)
+        )
+        duplicated = all(meta.duplicated for meta in metadata)
         await send_metadata(
             ch_out,
             context,
@@ -79,21 +89,22 @@ async def union_node(
         )
 
         seq_num_offset = 0
-        for ch_in in chs_in:
+        for ch_in, drop in zip(chs_in, skip, strict=True):
             num_ch_chunks = 0
             while (msg := await ch_in.recv(context)) is not None:
-                num_ch_chunks += 1
-                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
+                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
             seq_num_offset += num_ch_chunks
 
         await ch_out.drain(context)
@@ -116,6 +127,7 @@ def _(
     nodes[ir] = [
         union_node(
             rec.state["context"],
+            rec.state["comm"],
             ir,
             rec.state["ir_context"],
             channels[ir].reserve_input_slot(),

python/cudf_polars/cudf_polars/experimental/select.py

@@ -431,13 +431,17 @@ 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.
+        input_ir: IR = Empty({})
         new_node = Select(
             {named_expr.name: named_expr.value.dtype},
             [named_expr],
             should_broadcast=True,
-            df=child,
+            df=input_ir,
         )
-        partition_info[new_node] = PartitionInfo(count=1)
+        partition_info[input_ir] = partition_info[new_node] = PartitionInfo(count=1)
         return new_node, partition_info
 
     if not any(

python/cudf_polars/tests/experimental/test_dataframescan.py

@@ -60,6 +60,15 @@ def test_parallel_dataframescan(df, streaming_engine_factory, max_rows_per_parti
         assert count == 1
 
 
+@pytest.mark.xfail(
+    reason=(
+        "Multi-rank Union interleaves child outputs across ranks: client "
+        "receives [rank0_A, rank0_B, rank1_A, rank1_B] instead of the "
+        "polars-CPU [A, B]. Tracked in "
+        "https://github.com/rapidsai/cudf/issues/22376."
+    ),
+    strict=False,
+)
 def test_dataframescan_concat(df, streaming_engine_factory):
     streaming_engine = streaming_engine_factory(
         StreamingOptions(max_rows_per_partition=1_000),

python/cudf_polars/tests/test_groupby.py

@@ -501,8 +501,13 @@ def test_groupby_sum_decimal_null_group(engine: pl.GPUEngine) -> None:
 
 
 @pytest.mark.xfail(
-    raises=AssertionError,
-    reason="https://github.com/rapidsai/cudf/issues/19610",
+    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)."
+    ),
 )
 def test_groupby_literal_agg(engine: pl.GPUEngine):
     df = pl.LazyFrame({"c0": [True, False]})