[Python] Improving polars statement that adds a column applying a lambda function on each row

I am trying to add a column using map_rows in polars. The equivalent of pandas is as follows:

import pandas as pd
df = pd.DataFrame({"ref": [-1, 2, 8], "v1": [-1, 5, 0], "v2": [-1, 5, 8]})
df['count'] = df.apply(lambda r: len([i for i in r if i == r[0]]) - 1, axis=1)
df = df.drop('ref', axis=1)
df


v1 v2 count
0 -1 -1 2
1 5 5 0
2 0 8 1


The following is the sample code that I have with polars. Though it works as desired, it looks ugly and probably can be improved as well.

import polars as pl

df = pl.DataFrame({"ref": [-1, 2, 8], "v1": [-1, 5, 0], "v2": [-1, 5, 8]})

x = df.map_rows(lambda r: len([i for i in r if i == r[0]]) - 1).rename({'map': 'count'})
df = df.hstack([x.to_series()]).drop('ref')

df


shape: (3, 3)
┌─────┬─────┬───────┐
│ v1 ┆ v2 ┆ count │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═══════╡
│ -1 ┆ -1 ┆ 2 │
│ 5 ┆ 5 ┆ 0 │
│ 0 ┆ 8 ┆ 1 │
└─────┴─────┴───────┘


What bothers me is the rename part and hstack that I clobbered together to work. I would be grateful for any improvements in the above code.

TIA

Continue reading...