How can I drop duplicate rows from a Polars DataFrame according to a custom function?

I have a dataframe where I am trying to remove all duplicates in a specific column while aggregating the non-duplicated values.

The .unique function lets me choose only one of {‘first’, ‘last’, ‘any’, ‘none’}. However what I want is to apply the mean function to all numeric values, and the mode function to all categorical values.

I can do this by using groupby on the column I'm interested in, like in the example below:

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.groupby("id"):
    g = group.select(pl.col("id"),
       pl.all().exclude("id", "size").mode().first(),
       pl.col("size").mean()
    ).unique()
    df_list.append(g)

df_dedup = pl.concat(df_list)

Which gives me the output I'm expecting:

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

The problem is that this implementation is (unsurprisingly) quite slow, so I was wondering if there was a better way of doing this, or if it's possible to optimize the code I have in some way.

How about

In [22]: df.groupby("id").agg(
    ...:     pl.col(["color", "shape"]).mode().sort(descending=True).first(),
    ...:     pl.col("size").mean(),
    ...: )
    ...:
Out[22]:
shape: (2, 4)
┌─────┬───────┬──────────┬──────┐
│ id  ┆ color ┆ shape    ┆ size │
│ --- ┆ ---   ┆ ---      ┆ ---  │
│ i64 ┆ str   ┆ str      ┆ f64  │
╞═════╪═══════╪══════════╪══════╡
│ 1   ┆ red   ┆ triangle ┆ 2.0  │
│ 0   ┆ green ┆ square   ┆ 4.0  │
└─────┴───────┴──────────┴──────┘

Or, using the column selectors:

In [30]: import polars.selectors as cs

In [31]: df.groupby("id").agg(
    ...:     cs.string().mode().sort(descending=True).first(),
    ...:     cs.numeric().mean(),
    ...: )
Out[31]:
shape: (2, 4)
┌─────┬───────┬──────────┬──────┐
│ id  ┆ color ┆ shape    ┆ size │
│ --- ┆ ---   ┆ ---      ┆ ---  │
│ i64 ┆ str   ┆ str      ┆ f64  │
╞═════╪═══════╪══════════╪══════╡
│ 0   ┆ green ┆ square   ┆ 4.0  │
│ 1   ┆ red   ┆ triangle ┆ 2.0  │
└─────┴───────┴──────────┴──────┘

You can select columns by type, e.g. pl.col(pl.Utf8) for all str columns.

There's also the new polars.selectors helper module.

import polars.selectors as cs

df.groupby('id').agg(
   cs.string().mode().first(),
   cs.numeric().mean()
)

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