本文介绍了如何使用Polars的窗口函数为DataFrame中的分组数据添加组内行号。通过结合`int_range()`和`over()`函数，可以轻松地为每个分组生成从1开始的序列，从而实现组内行号的添加。
在数据分析中，有时需要在DataFrame的分组数据中添加组内行号，以便进行后续的分析和处理。Polars提供了一种简洁高效的方法来实现这一需求，即结合使用窗口函数和int_range()函数。
实现原理
核心思路是利用Polars的窗口函数，对每个分组应用int_range()函数生成一个序列，然后将该序列作为新的列添加到DataFrame中。
- int_range(length): 生成一个从0开始，长度为length的整数序列。
- .over("groupings"): 指定窗口函数的作用范围，这里是按照"groupings"列进行分组。
具体步骤与代码示例
假设我们有一个DataFrame，其中包含"groupings"列和"target_count_over_windows"列，我们希望为每个"groupings"分组添加一个从1开始的行号列，命名为"count"。
首先，创建示例DataFrame：
import polars as pl
df = pl.DataFrame([
{'groupings': 'a', 'target_count_over_windows': 1},
{'groupings': 'a', 'target_count_over_windows': 2},
{'groupings': 'a', 'target_count_over_windows': 3},
{'groupings': 'b', 'target_count_over_windows': 1},
{'groupings': 'c', 'target_count_over_windows': 1},
{'groupings': 'c', 'target_count_over_windows': 2},
{'groupings': 'd', 'target_count_over_windows': 1},
{'groupings': 'd', 'target_count_over_windows': 2},
{'groupings': 'd', 'target_count_over_windows': 3}
])
print(df)输出结果：
shape: (9, 2) ┌───────────┬───────────────────────────┐ │ groupings ┆ target_count_over_windows │ │ --- ┆ --- │ │ str ┆ i64 │ ╞═══════════╪═══════════════════════════╡ │ a ┆ 1 │ │ a ┆ 2 │ │ a ┆ 3 │ │ b ┆ 1 │ │ c ┆ 1 │ │ c ┆ 2 │ │ d ┆ 1 │ │ d ┆ 2 │ │ d ┆ 3 │ └───────────┴───────────────────────────┘
接下来，使用窗口函数添加组内行号：
df = df.with_columns(count = 1 + pl.int_range(pl.len()).over("groupings"))
print(df)输出结果：
shape: (9, 3) ┌───────────┬───────────────────────────┬───────┐ │ groupings ┆ target_count_over_windows ┆ count │ │ --- ┆ --- ┆ --- │ │ str ┆ i64 ┆ i64 │ ╞═══════════╪═══════════════════════════╪═══════╡ │ a ┆ 1 ┆ 1 │ │ a ┆ 2 ┆ 2 │ │ a ┆ 3 ┆ 3 │ │ b ┆ 1 ┆ 1 │ │ c ┆ 1 ┆ 1 │ │ c ┆ 2 ┆ 2 │ │ d ┆ 1 ┆ 1 │ │ d ┆ 2 ┆ 2 │ │ d ┆ 3 ┆ 3 │ └───────────┴───────────────────────────┴───────┘
可以看到，新的"count"列包含了每个分组内的行号，从1开始递增。
注意事项
- pl.len() 返回的是每个分组的长度，int_range 会生成一个从0开始的序列，因此需要加1才能得到从1开始的行号。
- 窗口函数非常灵活，可以结合其他表达式和函数来实现更复杂的需求。
总结
使用Polars的窗口函数和int_range()函数可以方便快捷地为DataFrame中的分组数据添加组内行号。这种方法简洁高效，适用于各种数据分析场景。掌握这一技巧可以有效提升数据处理效率。
