diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py index abb2e7082f01..b36b07342ced 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py +++ b/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, diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/union.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/union.py index b4cb6a922b93..2484620234df 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/union.py +++ b/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(), diff --git a/python/cudf_polars/cudf_polars/experimental/select.py b/python/cudf_polars/cudf_polars/experimental/select.py index 25d0189fdf6d..9ab30f9be136 100644 --- a/python/cudf_polars/cudf_polars/experimental/select.py +++ b/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( diff --git a/python/cudf_polars/tests/experimental/test_dataframescan.py b/python/cudf_polars/tests/experimental/test_dataframescan.py index 57684734fea9..dbf228488240 100644 --- a/python/cudf_polars/tests/experimental/test_dataframescan.py +++ b/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), diff --git a/python/cudf_polars/tests/test_groupby.py b/python/cudf_polars/tests/test_groupby.py index a14177b9f0cf..f14160a1043e 100644 --- a/python/cudf_polars/tests/test_groupby.py +++ b/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]})