[BUG] Explicit-comms shuffle does not obey "_partitions" column

[rjzamora]

While debugging a data-curation workflow, I discovered that the explicit-comms shuffle has a subtle bug in the logic used to assign data to the final partitions when "_partitions" is specified. For example:

from dask_cuda import LocalCUDACluster
from distributed import Client
import dask.dataframe as dd
from dask.datasets import timeseries

from dask_cuda.explicit_comms.dataframe.shuffle import (
    shuffle as explicit_comms_shuffle,
)

client = Client(LocalCUDACluster(n_workers=2))

ddf = timeseries().reset_index(drop=True).to_backend("cudf")
ddf["_partitions"] = 0

result = explicit_comms_shuffle(ddf, ["_partitions"])
result.partitions[0].compute()

Since all rows should be 0 in the "_partitions" columns, then all data should be moved to partition 0 after the shuffle. However, I get an empty DataFrame when I execute this:

Empty DataFrame
Columns: [name, id, x, y, _partitions]
Index: []

As far as I can tell, this problem is caused by the fact that shuffle_result[rank] is not in the same order as rank_to_out_part_ids[rank] in this client.submit loop (the order is reversed).

[madsbk]

Good catch @rjzamora !