Polars DataFrame 中如何在窗口内添加行号

本文介绍了如何在 Polars DataFrame 的每个窗口（分组）内添加行号。通过结合 int_range() 函数和窗口函数，可以方便地为每个分组生成连续的行号，从而实现更灵活的数据分析和处理。

在 Polars 中，虽然 with_row_numbers() 方法可以为整个 DataFrame 添加行号，但有时我们需要在特定的窗口（分组）内添加行号，以便进行更细粒度的数据分析。本文将介绍如何利用 Polars 的窗口函数和 int_range() 函数来实现这一目标。

核心思路：

核心思路是使用 pl.int_range(pl.len()) 生成一个从 0 开始，长度等于组大小的整数序列，然后将其加 1，使其成为从 1 开始的行号。再通过 over("groupings") 将此操作应用于每个分组窗口内。

具体实现：

以下是一个示例，演示了如何在 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)

代码解释：

1. pl.int_range(pl.len()): pl.len() 返回当前分组的大小。pl.int_range() 则生成一个从 0 开始，长度等于分组大小的整数序列。例如，如果一个分组有 3 行，则 pl.int_range(pl.len()) 将生成 [0, 1, 2]。

2. 1 + pl.int_range(pl.len()): 将生成的整数序列加 1，使其从 1 开始，作为行号。 在上述例子中，就变成了 [1, 2, 3]。

   

   Trae国内版

   国内首款AI原生IDE，专为中国开发者打造

   815

   查看详情

3. .over("groupings"): 这是一个窗口函数。它指定了将上述操作应用于哪个分组。在这个例子中，我们按照 "groupings" 列进行分组，并在每个分组内应用上述操作。

4. df.with_columns(count = ...): 使用 with_columns() 方法将新生成的 "count" 列添加到 DataFrame 中。

运行结果：

运行上述代码将得到以下 DataFrame：

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" 列包含了每个分组内的行号。

注意事项：

• int_range() 函数默认从 0 开始生成整数序列。为了生成从 1 开始的行号，需要在结果上加 1。
• over() 函数的参数指定了分组的列名。可以根据实际情况修改此参数。
• 确保数据已按照分组列排序，以保证行号的连续性。如果数据未排序，可以使用 sort() 方法进行排序。

总结：

通过结合 int_range() 函数和窗口函数，可以方便地在 Polars DataFrame 的每个分组内添加行号。这种方法灵活且高效，可以满足各种数据分析需求。 掌握这种技巧，能够更方便地进行分组统计、排名等操作。

以上就是Polars DataFrame 中如何在窗口内添加行号的详细内容，更多请关注php中文网其它相关文章！