[BUG][cudf-polars] Multi-rank `Union` (and `pl.concat`) interleaves child outputs across ranks, breaking row order

[madsbk]

Summary

When the streaming engine executes pl.concat([A, B]) (or any Union IR node) with N > 1 ranks, the resulting row order does not match Polars semantics. Polars guarantees that pl.concat([A, B]) yields all rows of A followed by all rows of B. Under multi-rank execution, the client instead observes rank-local concatenations interleaved at rank boundaries.

Reproducer

import polars as pl
from cudf_polars.experimental.rapidsmpf.frontend.ray import RayEngine
from cudf_polars.testing.asserts import assert_gpu_result_equal

with RayEngine(
    num_ranks=2,  # Work when `num_ranks=1`
    engine_options={"allow_gpu_sharing": True},
    executor_options={"max_rows_per_partition": 1_000},
) as streaming_engine:
    df = pl.LazyFrame({
        "x": range(30_000),
        "y": [1, 2, 3] * 10_000,
        "z": [1.0, 2.0, 3.0, 4.0, 5.0] * 6_000,
    })
    df2 = pl.concat([df, df])
    assert_gpu_result_equal(df2, engine=streaming_engine, check_row_order=True)

Currently xfailed in:
python/cudf_polars/tests/experimental/test_dataframescan.py::test_dataframescan_concat

Observed vs. expected order (2 ranks, children A and B)

• Expected (Polars CPU): A_rank0, A_rank1, B_rank0, B_rank1
• Actual (streaming engine): A_rank0, B_rank0, A_rank1, B_rank1

Each rank correctly processes child A before B, but the client-side concatenation across ranks does not enforce a barrier between children. As a result, chunks from child A on rank 1 may arrive after chunks from child B on rank 0.

Root cause (sketch)

cudf_polars/experimental/rapidsmpf/union.py::union_node iterates over chs_in in order and forwards child[0] then child[1] per rank, but emits chunks asynchronously. The downstream collector concatenates chunks in arrival order across ranks. There is no cross-rank synchronization between children, and chunks carry no child_id metadata to allow reordering on the client.

Possible fixes

1. Add a per-child cross-rank barrier in union_node: emit all chunks for child A, synchronize across ranks via the communicator, then proceed to child B. Simple, but serializes cross-child streaming.

2. Tag chunks with child_id metadata and have the client-side collector group and concatenate per child before combining children. Preserves streaming overlap, but requires extending metadata on TableChunk/messages.

Option (2) is preferred unless profiling shows the barrier cost is negligible.

Scope

Affects any Union/pl.concat query under the rapidsmpf streaming engine with N > 1 ranks. Not specific to DataFrameScan.

[wence-]

I think @rjzamora may have thoughts here since Union is used in lots of places in the rewrite phases of the IR to mean local (single-rank) concatenation rather than global concatenation.

The barrier idea doesn't work to fix anything because the semantics of reconstructing a full dataframe from a channel is to first concat the chunks in a channel on a rank, and then concatenate across ranks. Since union is a function

$$ U : \bigotimes_i C^{\text{in}}_i \to C^{\text{out}} $$

The consumer of a union just sees a single $C^\text{out}$. Inserting a barrier in between emitting each C^\text{in} doesn't change that.

Tagging the channel with metadata that a "concat-transpose" has to be performed to reconstruct the correct result sounds appealing, but I think it would pervade the evaluation of huge swathes of the query plan. The contract currently is that each rank can compute locally on its part and that is a contiguous slice of the dataframe. But concat breaks that because without communicating each rank has a block-cyclic slice of the dataframe.

So I suspect that we must either generalise the permissible data layouts and handle them in all our algorithms, or else suck up the transpose and data movement implied by concat.

[rjzamora]

I think @rjzamora may have thoughts here since Union is used in lots of places in the rewrite phases of the IR to mean local (single-rank) concatenation rather than global concatenation.

Ah, right - sorry. We have known about this concat/Union issue for a long time. However, it's not actually used in any re-write phases anymore. We use Union.do_evaluate within the _concat utility, but don't actually use an explicit Union in the physical plan. The only reason we didn't implement a proper union_actor yet is that the ordering doesn't matter for tpch (there is at least one explicit concat, but you sort at the end anyway).

At one point, I think we floated the idea that we should just warn the user that the Union won't preserve ordering and carry on. I suppose we can also just rebalance all the data. It's expensive, but not that hard to do (the ranks can allgather the expected chunk count for each element of the Union and send the chunks to the new rank).

[wence-]

At one point, I think we floated the idea that we should just warn the user that the Union won't preserve ordering and carry on.

That is OK if you actually don't care about the order. But what if you're doing pl.concat([A, B]).rolling(...). Now if A and B are ordered correctly for rolling, their concatenation will not be.

[wence-]

I note as well polars now has the concept of pl.union (unordered concatenation) and pl.concat (ordered concatenation). Those both get represented as the same Union node with different options (maintain_order or not), but we don't currently expose that flag.

[rjzamora]

That is OK if you actually don't care about the order.

Right - If concat doesn't preserve ordering, you would need to sort before the rolling operation. This example also illustrates that there are cases where the "meaning" of the ordering after a concat is significant, and others where it might not be. For example, if we just hash-partitioned one of the tables, then the concatenated order doesn't seem very meaningful to me. However, if the tables are sorted (like your example), then the order may be very meaningful.

I'm still wondering if there are clean ways to avoid unnecessary rebalancing. Perhaps rapidsai/rapidsmpf#853 will help us here?

[rjzamora]

I note as well polars now has the concept of pl.union (unordered concatenation) and pl.concat (ordered concatenation). Those both get represented as the same Union node with different options (maintain_order or not), but we don't currently expose that flag.

Oh! Very interesting. I didn't realize this.