Сортировка выражения group_by при взятии первой строки с сохранением всех столбцов ⇐ Python

[Anonymous]

Учитывая следующий кадр данных, я хотел бы сгруппировать по «foo», отсортировать по «bar», а затем сохранить всю строку.

Код: Выделить всё

df = pl.DataFrame(
{
"foo": [1, 1, 1, 2, 2, 2, 3],
"bar": [5, 7, 6, 4, 2, 3, 1],
"baz": [1, 2, 3, 4, 5, 6, 7],
}
)

df_desired = pl.DataFrame({"foo": [1, 2, 3], "bar": [5, 2, 1], "baz": [1,5,7]})

Код: Выделить всё

>>> df_desired
shape: (3, 3)
┌─────┬─────┬─────┐
│ foo ┆ bar ┆ baz │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ 1   ┆ 5   ┆ 1   │
│ 2   ┆ 2   ┆ 5   │
│ 3   ┆ 1   ┆ 7   │
└─────┴─────┴─────┘

Я могу сделать это, отсортировав заранее, но это затратно по сравнению с сортировкой группы:

Код: Выделить всё

df_solution = df.sort("bar").group_by("foo", maintain_order=True).first().sort(by="foo")

assert df_desired.equals(df_solution)

Я могу сортировать по "foo" в агрегировании, как в этом ответе SO:

Код: Выделить всё

>>> df.group_by("foo").agg(pl.col("bar").sort().first()).sort(by="foo")
shape: (3, 2)
┌─────┬─────┐
│ foo ┆ bar │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 1   ┆ 5   │
│ 2   ┆ 2   │
│ 3   ┆ 1   │
└─────┴─────┘

но тогда я получаю только этот столбец. Как мне сохранить значение строки «baz»? Любые дополнительные записи в .agg([]) не зависят от нового pl.col("bar").sort().


Подробнее здесь: https://stackoverflow.com/questions/744 ... ll-columns