Varying quantile by group is broken

[orlp]

Consider the following example:

df = pl.DataFrame(
    {
        "value": [1, 2, 1, 2],
        "quantile": [0, 0, 1, 1],
    }
)
df.group_by(pl.col.quantile).agg(pl.col.value.quantile(pl.col.quantile.first()))

Output:

shape: (2, 2)
┌──────────┬───────┐
│ quantile ┆ value │
│ ---      ┆ ---   │
│ i64      ┆ f64   │
╞══════════╪═══════╡
│ 0        ┆ 1.0   │
│ 1        ┆ 1.0   │
└──────────┴───────┘

I would expect the value of the quantile = 1 group to be 2.0, not 1.0.

[npennequin]

I also ran into this issue. Was about to open an issue.

I found that using map_groups returns the correct results, but it's slower:

(
    df
    .group_by("quantile")
    .map_groups(lambda group_df: group_df.select(
        pl.col.quantile.first(),
        pl.col.value.quantile(pl.col.quantile.first())
    ))
)

Output:

shape: (2, 2)
┌──────────┬───────┐
│ quantile ┆ value │
│ ---      ┆ ---   │
│ i64      ┆ f64   │
╞══════════╪═══════╡
│ 0        ┆ 1.0   │
│ 1        ┆ 2.0   │
└──────────┴───────┘