Add sort_actor to cudf-polars + rapidsmpf

[rjzamora]

Description

Closes #20486
Depends on rapidsai/rapidsmpf#891 (or similar)

Adds sort_actor for the "rapidsmpf" runtime.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[wence-]

This is incorrect in a multi-rank setting. I think what you're trying to do is give each row in the (concatenated) stream of local table chunks a unique id.

seq_num * (1 << 32) is a good offset for this because by construction no sequence numbers from chunks on the same rank can now overlap.

But, it seems unnecessary (you have the number of rows when adding the column, so just keep track of that locally).

However, what this doesn't do is ensure that two different ranks have different sequence numbers. For that, you need to incorporate the rank in the high bits of the id, such that when you sort globally, rows from rank-1 (say) come after rows from rank-0 if their keys are equal.

[rjzamora]

Okay, thanks for explaining - I updated the logic here (hopefully in line with your suggestion).

[wence-]

This sort is wrong, because it sorts by all the columns in the table, whereas you only want to sort by the key columns (and the disambiguating seq_id column).

You want to be using sort_by_key

[rjzamora]

Makes sense - I decided to construct a Sort so I can use do_evaluate, but I can move to an explicit sort_by_key if you'd prefer.

[wence-]

These arguments are only used to construct the schema of the "empty" table for the case that the local candidates are empty, which is then used to provide the schema for chunk_to_frame.

Two things:

1. we only need the schema for all of those operations, so we should redo the utilities to take the schema.
2. Let's construct the empty table outside if local_candidates is empty.

So this function becomes:

async def _compute_sort_boundaries(
    context: Context,
    comm: Communicator,
    ir_context: IRExecutionContext,
    local_candidates: list[TableChunk],
    schema: Schema,
    num_partitions: int,
    column_order: list[plc.types.Order],
    null_order: list[plc.types.NullOrder],
    allgather_id: int
) -> plc.Table:
    stream = ...
    boundaries = _get_final_sort_boundaries(
        chunk_to_frame(
            await concat_batch(local_candidates, context, schema, ir_context)
        )
    )
    if comm.nranks > 1:
         ...
    boundaries = _get_final_sort_boundaries(
         ...
    )
    return boundaries.table

[rjzamora]

Okay, yeah - I revised this a bit.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[rjzamora]

Fiddle sticks. I accidentally merged in main - I guess this is slipping to 26.06 then :)