python - 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))

在熊猫中，我可以使用转换生成这样的列“codeindex”。我怎么能在 py-polars 中做到这一点？

多谢！

Answer 1

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

Windows 表达式对您使用 分区的组进行操作.over()。它由表达式部分 likecol("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         │
└──────┴──────┴───────┴───────────┘