ensure drop_duplicates with keep is stable

[fjetter]

Closes #10708

cc @rjzamora

[phofl]

This isn't stable unfortunately, we can't rely on the index of our original df being sorted:

if __name__ == "__main__":
    client = Client()
    pdf = pd.DataFrame(
        {"x": [1, 2, 3, 4, 5, 6] * 10, "y": np.random.randint(1, 1_000_000, (60, ))},
        index=pd.Series(list(range(0, 30)) * 2),
    )
    df = dd.from_pandas(pdf, npartitions=2, sort=False)
    result_pd = pdf.drop_duplicates(subset=["x"], keep="first")
    assert_eq(df.drop_duplicates(
            subset=["x"],
            keep="first",
            split_out=df.npartitions,
            shuffle="p2p",
        ).compute(), result_pd)

[rjzamora]

This isn't stable unfortunately, we can't rely on the index of our original df being sorted

Agree, we should be able to do this if we have sorted divisions. Otherwise, we would need to add a temporary "__order" column.

[fjetter]

Yes, you are right, of course.

Otherwise, we would need to add a temporary "__order" column.

I thought about this for a moment but couldn't come up with something that would actually give me this. We'd need internal ordering (can be done with a reset_index) and a global ordering (partition index). Is there an easy way to assign a column with the partition index?

[rjzamora]

Is there an easy way to assign a column with the partition index?

Easy-ish:

from dask.blockwise import BlockIndex

def _add_partid(x, val):
    return x.assign(__partid=val[0])

df2 = df.map_partitions(_add_partid, BlockIndex((df.npartitions,)))
df2.head()

   x  y  __partid
0  1  a         0
0  1  a         0
1  2  b         0
1  2  b         0
2  3  d         0

[fjetter]

BlockIndex

interesting, thank you.