python poloars 中是否有与 pandas 中的 'transform' 类似的功能？

Question

在 pandas 中，我可以运行这样的代码：

day = ['day1','day2','day3','day4','day1','day2','day3','day1','day2']

code = ["a","a","a","a","b","b","b","c","c"]

price = [1,2,3,4,5,6,7,8,9]

df = pd.DataFrame({"date":day,"code":code,"price":price})


df['codeindex'] = df.groupby('code')['date'].transform(lambda x: range(0, len(x), 1))

在 pandas 我可以使用转换生成这样的列 'codeindex'。我如何在 py-polars 中做到这一点？

非常感谢！

Answer 1

您可以使用window expressions来处理只需要在组内应用的表达式。

windows 表达式对您使用 .over() 划分的组进行运算。它由 col("date").cumcount() 等表达式部分和 .over("code").

定义的分区部分定义

如果您使用 aggregation 结果将被广播以匹配组的大小。

代码如下所示：

day = ['day1','day2','day3','day4','day1','day2','day3','day1','day2']

code = ["a","a","a","a","b","b","b","c","c"]

price = [1,2,3,4,5,6,7,8,9]

df = pl.DataFrame({"date":day,"code":code,"price":price})

(df.select([
    pl.all(),
    pl.col("date").cumcount().over("code").alias("codeindex"),
]))

产出

shape: (9, 4)
┌──────┬──────┬───────┬───────────┐
│ date ┆ code ┆ price ┆ codeindex │
│ ---  ┆ ---  ┆ ---   ┆ ---       │
│ str  ┆ str  ┆ i64   ┆ u32       │
╞══════╪══════╪═══════╪═══════════╡
│ day1 ┆ a    ┆ 1     ┆ 0         │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ day2 ┆ a    ┆ 2     ┆ 1         │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ day3 ┆ a    ┆ 3     ┆ 2         │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ day4 ┆ a    ┆ 4     ┆ 3         │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ ...  ┆ ...  ┆ ...   ┆ ...       │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ day2 ┆ b    ┆ 6     ┆ 1         │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ day3 ┆ b    ┆ 7     ┆ 2         │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ day1 ┆ c    ┆ 8     ┆ 0         │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ day2 ┆ c    ┆ 9     ┆ 1         │
└──────┴──────┴───────┴───────────┘

python poloars 中是否有与 pandas 中的 'transform' 类似的功能？

is there any similar function in python poloars like 'transform' in pandas?

python

dataframe

python-polars

产出