Как я могу удалить повторяющиеся строки из DataFrame Polars в соответствии с пользовательской функцией? ⇐ Python

[Anonymous]

У меня есть фрейм данных, в котором я пытаюсь удалить все дубликаты в определенном столбце, агрегируя неповторяющиеся значения.
Функция .unique позволяет мне выбрать только одно из {'first', 'last', 'any', 'none'. Однако я хочу применить функцию среднего ко всем числовым значениям и функцию режима ко всем категориальным значениям.
Я могу сделать это, используя group_by в интересующем меня столбце, как в примере ниже:

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

import polars as pl

df = pl.DataFrame(
{
"id": [0, 0, 0, 1, 1],
"color": ["red", "green", "green", "red", "red"],
"shape": ["square", "triangle", "square", "triangle", "square"],
"size": [2, 4, 6, 1, 3]
}
)

df_list = []
for gkey, group in df.group_by("id"):
g = group.select(pl.col("id"),
pl.exclude("id", "size").mode().first(),
pl.col("size").mean()
).unique()
df_list.append(g)

df_dedup = pl.concat(df_list)

Это дает мне ожидаемый результат:

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

> print(df_dedup)
shape: (2, 4)
┌─────┬───────┬──────────┬──────┐
│ id  ┆ color ┆ shape    ┆ size │
│ --- ┆ ---   ┆ ---      ┆ ---  │
│ i64 ┆ str   ┆ str      ┆ f64  │
╞═════╪═══════╪══════════╪══════╡
│ 1   ┆ red   ┆ triangle ┆ 2.0  │
│ 0   ┆ green ┆ square   ┆ 4.0  │
└─────┴───────┴──────────┴──────┘

Проблема в том, что эта реализация (что неудивительно) довольно медленная, поэтому мне было интересно, есть ли лучший способ сделать это или можно ли каким-то образом оптимизировать имеющийся у меня код.

Подробнее здесь: https://stackoverflow.com/questions/765 ... ustom-func