Fix order-preservation in cudf-polars groupby

[wence-]

Description

When we are requested to maintain order in groupby aggregations we must post-process the result by computing a permutation between the wanted order (of the input keys) and the order returned by the groupby aggregation. To do this, we can perform a join between the two unique key tables. Previously, we assumed that the gather map returned in this join for the left (wanted order) table was the identity. However, this is not guaranteed, in addition to computing the match between the wanted key order and the key order we have, we must also apply the permutation between the left gather map order and the identity.

• Closes [BUG] in polars-gpu, group_by(maintain_order=True) is not ordered #16893

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-]

I have pointed this at 24.10, but have marked as do not merge since I don't think this is a critical fix. If others agree, I can retarget to 24.12.

[brandon-b-miller]

LGTM

[bdice]

@wence- You said a similar fix is needed in cudf.pandas? This adds a fair number of steps to the computation. Should we do this at the C++ level?

I am happy to merge this for 24.10 and optimize in 24.12, or just target 24.12.

[wence-]

@wence- You said a similar fix is needed in cudf.pandas?

Yes: #16908

This adds a fair number of steps to the computation. Should we do this at the C++ level?

It's one additional sort_by_key (sorting an integer column by an integer key column, which should be a radix sort). So I don't think it's that many more steps.

Note this is only done if the user requests maintain_order=True in the groupby (which is not the default for polars).

I am happy to merge this for 24.10 and optimize in 24.12, or just target 24.12.

The optimisations that is plausibly worthwhile is that we know that both tables have distinct rows and so we could use the new distinct_hash_join code in libcudf.

As you note, since this kind of thing is a somewhat common pattern, it might be worth implementing a libcudf primitive that, given two tables, returns the gather map that is the permutation from one table to the other.

I'd need to look if the join ordering code has the same invariants as the groupby ordering code.

[vyasr]

After discussion, we're moving this to 24.12.

[wence-]

After discussion, we're moving this to 24.12.

Done.

[wence-]

/merge