Add support for maintain_order param in joins

[Matt711]

Description

Closes #17696

Checklist

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

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[wence-]

Thanks @Matt711, I think we are doing too much work in the "none" case.

[Matt711]

For the reviewer: This PR needs more work, but I'm opening it up for review so I can get some help handling a special case: full joins (where we maintain the order of the right table). Specifically, the case where the test fails is when there are unmatched keys in the left dataframe. Any advice on how to handle this?

Example:

left = pl.LazyFrame(
    {
        "a": [1, 2, 3, 1, None],
        "b": [1, 2, 3, 4, 5],
        "c": [2, 3, 4, 5, 6],
    }
)
right = pl.LazyFrame(
    {
        "a": [1, 4, 3, 7, None, None, 1],
        "c": [2, 3, 4, 5, 6, 7, 8],
        "d": [6, None, 7, 8, -1, 2, 4],
    }
)
q = left.join(right, on=pl.col("a"), how="full", maintain_order="right")
q.collect(engine="gpu")

The dataframe differ at column "a"

AssertionError: DataFrames are different (value mismatch for column 'a')
[left]:  [1, 1, None, 3, None, None, None, 1, 1, **None, 2]**
[right]: [1, 1, None, 3, None, None, None, 1, 1, **2, None]**

The a=2 entry is unmatched in the right dataframe, so it should be appended to the end of the result, not included with the other matches.

[wence-]

I think this is expected. Because those last two rows have the same sort key in the right table column, there's no disambiguator to decide which order the left column result comes in.

e.g.

(No GPU engine involved):

In [18]: q = left.join(right, on=pl.col("a"), how="full", maintain_order="right")

In [19]: q.collect(engine="cpu")
Out[19]: 
shape: (11, 2)
┌──────┬─────────┐
│ a    ┆ a_right │
│ ---  ┆ ---     │
│ i64  ┆ i64     │
╞══════╪═════════╡
│ 1    ┆ 1       │
│ 1    ┆ 1       │
│ null ┆ 4       │
│ 3    ┆ 3       │
│ null ┆ 7       │
│ …    ┆ …       │
│ null ┆ null    │
│ 1    ┆ 1       │
│ 1    ┆ 1       │
│ null ┆ null    │
│ 2    ┆ null    │
└──────┴─────────┘

In [20]: q.collect(engine="cpu")
Out[20]: 
shape: (11, 2)
┌──────┬─────────┐
│ a    ┆ a_right │
│ ---  ┆ ---     │
│ i64  ┆ i64     │
╞══════╪═════════╡
│ 1    ┆ 1       │
│ 1    ┆ 1       │
│ null ┆ 4       │
│ 3    ┆ 3       │
│ null ┆ 7       │
│ …    ┆ …       │
│ null ┆ null    │
│ 1    ┆ 1       │
│ 1    ┆ 1       │
│ 2    ┆ null    │
│ null ┆ null    │
└──────┴─────────┘

Notice how the last two rows flip around between the two runs.

[Matt711]

Should I open an issue in polars to add maintain_order="left_right" or some other value than none? That way we wouldn't have to xfail some of these tests. WDYT?
cc. @wence-

[wence-]

Oh maybe, yes

[Matt711]

I tested this PR with Polars 1.19 and these tests now pass

[Matt711]

I think the problem is related to left joins preserving the order of the left table by default. There are two options I see

1. If I'm right about left joins, we can try handling that case in this PR (ie. how="left" and maintain_order=None => preserve the left table order)
2. Keep the PR as is and do a polars version bump
   Thoughts?

[wence-]

Let's bump the version I think. I would like to be aligned with the most recent polars release before we go into burndown anyway

[vyasr]

Is this PR gated on #17771 going in? If there are Polars version-specific behaviors that we care about?