Polars Window 函数：为分组数据添加组内行号

本文介绍了如何使用 Polars 的 Window 函数为 DataFrame 中按分组划分的数据添加组内行号。通过 int_range() 函数和 over() 方法的结合，可以轻松实现对每个分组数据的行号计数，从而进行更细粒度的数据分析和处理。

在 Polars 中，为 DataFrame 添加行号通常使用 with_row_numbers() 方法，但该方法会为整个 DataFrame 添加行号，而不是针对特定分组。如果需要为每个分组单独添加行号（即组内行号），则需要借助 Window 函数来实现。

使用 int_range() 和 over() 函数实现组内行号

以下示例展示了如何使用 int_range() 函数和 over() 方法为 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}
])

df = df.with_columns(count = 1 + pl.int_range(pl.len()).over("groupings"))

print(df)

代码解释:

RecoveryFox AI

AI驱动的数据恢复、文件恢复工具

下载

1. pl.int_range(pl.len()): pl.len() 获取每个分组的长度，pl.int_range() 根据这个长度生成一个从 0 开始的整数序列。例如，如果一个分组的长度为 3，则会生成序列 [0, 1, 2]。
2. .over("groupings"): over("groupings") 指定了 Window 函数的作用范围，即按照 "groupings" 列进行分组。对于每个分组，int_range() 函数都会生成一个独立的整数序列。
3. 1 + ...: 由于 int_range() 生成的序列从 0 开始，因此需要加 1 才能得到从 1 开始的行号。
4. df.with_columns(count = ...): with_columns() 方法用于向 DataFrame 添加新列，这里添加了一个名为 "count" 的新列，其值为每个分组的行号。

输出结果:

运行上述代码后，DataFrame 将会增加一个名为 "count" 的列，其中包含每个分组的行号：

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     │
└───────────┴───────────────────────────┴───────┘

总结

通过结合 int_range() 和 over() 函数，可以灵活地为 Polars DataFrame 中的分组数据添加组内行号。这种方法在需要对分组数据进行排序、排名或进行其他基于行号的操作时非常有用。理解和掌握 Window 函数的使用，能够更有效地利用 Polars 进行数据分析和处理。