本文介绍了如何在Polars DataFrame中使用窗口函数为每个分组添加行号。通过结合`int_range()`函数和`over()`方法，可以轻松地在每个窗口内生成递增的序列，从而实现分组行号的功能。本文提供详细的代码示例和解释，帮助读者理解并掌握该技巧。
在数据分析中，经常需要在每个分组内添加行号，以便进行后续的计算或分析。Polars 提供了强大的窗口函数功能，可以方便地实现这一需求。本文将详细介绍如何使用 Polars 的窗口函数和 int_range() 函数，为 DataFrame 中的每个分组添加行号。
核心思路
核心思路是利用 pl.int_range(pl.len()) 在每个分组内生成一个从 0 开始的整数序列，然后通过窗口函数 over("groupings") 将其应用到每个分组。最后，加 1 即可得到从 1 开始的行号。
具体实现
下面是一个具体的代码示例：
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}
])
df = df.with_columns(count = 1 + pl.int_range(pl.len()).over("groupings"))
print(df)代码解释
- 导入 Polars 库: import polars as pl 导入 Polars 库，并将其别名为 pl。
- 创建 DataFrame: 创建一个示例 DataFrame，包含 groupings 和 target_count_over_windows 两列。
-
使用 with_columns 添加新列: df.with_columns(count = 1 + pl.int_range(pl.len()).over("groupings")) 使用 with_columns 方法添加一个名为 count 的新列。
- pl.int_range(pl.len()): pl.int_range(pl.len()) 生成一个整数序列，序列的长度等于每个分组的长度。 pl.len() 函数用于获取每个分组的长度。
- .over("groupings"): .over("groupings") 指定窗口函数的作用范围，这里表示按照 groupings 列进行分组。
- 1 + ...: 将生成的整数序列加 1，使得行号从 1 开始。
运行结果
运行上述代码，将得到以下结果：
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 │ └───────────┴───────────────────────────┴───────┘
可以看到，DataFrame 中成功添加了 count 列，其中包含了每个分组内的行号。
注意事项
- pl.int_range() 函数生成的序列从 0 开始，因此需要加 1 才能得到从 1 开始的行号。
- over() 方法用于指定窗口函数的作用范围，必须指定分组的列名。
- 该方法适用于任何类型的分组列，例如字符串、整数等。
总结
本文介绍了如何使用 Polars 的窗口函数和 int_range() 函数，为 DataFrame 中的每个分组添加行号。该方法简单易懂，效率高，适用于各种数据分析场景。通过掌握该技巧，可以更加方便地进行分组计算和分析。
