使用 Polars 高效计算 DataFrame 中按 ID 分组的时间间隔

在数据分析和处理中，我们经常会遇到需要计算时间序列数据中事件间时间间隔的场景。尤其是在处理用户行为日志、会话数据或传感器读数时，按特定实体（如用户ID、设备ID）分组并计算其内部连续事件的时间差，是常见的分析需求。Polars 作为一个高性能的 DataFrame 库，提供了强大的表达式系统和窗口函数功能，能够以极高的效率完成此类任务。

1. 问题背景与 Polars 优势

假设我们有一个包含 ID 和 Timestamp 列的 DataFrame，其中 Timestamp 表示某个会话或事件的结束时间。我们的目标是为每个唯一的 ID 计算其连续会话之间的时间间隔，并将结果存储在一个新列 time_between_sessions 中。

传统的 Pandas 操作可能会倾向于使用 groupby() 结合 apply() 或 transform()，但在大数据集上，这些方法可能效率不高。Polars 的设计理念是充分利用多核 CPU 和并行化，其表达式系统和惰性计算能力使其在处理大规模数据时表现出色。对于按组计算的需求，Polars 提供了更为优化的窗口函数（Window Functions）。

2. 核心概念：窗口函数 pl.Expr.over()

Polars 的 pl.Expr.over() 是实现分组计算的关键。它允许您在不显式执行 group_by() 操作的情况下，对数据进行分组并在每个组内应用表达式。这与 SQL 中的 OVER 子句概念类似，使得聚合、排名或差值计算等操作能够作用于数据的特定分区。

当与 diff() 函数结合使用时，over("ID") 会确保 diff() 操作仅在其所属的 ID 分组内进行，从而正确计算每个 ID 内部的连续时间差。

3. 解决方案详解

我们将通过一个最小可复现示例来演示如何使用 Polars 高效地计算每个 ID 的会话间隔。

3.1 准备示例数据

首先，我们创建一个包含 ID 和 Timestamp 的示例 DataFrame，并将其转换为 Polars DataFrame，确保 Timestamp 列是 Polars 的 Datetime 类型。

import polars as pl
import pandas as pd

# 创建一个示例 Pandas DataFrame
data = {
    'ID': ['A', 'A', 'A', 'B', 'B', 'B'],
    'Timestamp': ['2023-01-01 10:00:00', '2023-01-01 10:30:00' ,'2023-01-01 11:00:00', '2023-01-01 12:00:00', '2023-01-01 12:30:00', '2023-01-01 13:00:00']
}

df_pandas = pd.DataFrame(data)

# 转换为 Polars DataFrame 并确保 Timestamp 为 Datetime 类型
sessions_features = pl.from_pandas(df_pandas).with_columns(
   pl.col("Timestamp").str.to_datetime()
)

print("原始 Polars DataFrame:")
print(sessions_features)

输出示例:

美间AI

美间AI：让设计更简单

45

查看详情

原始 Polars DataFrame:
shape: (6, 2)
┌─────┬─────────────────────┐
│ ID  ┆ Timestamp           │
│ --- ┆ ---                 │
│ str ┆ datetime[μs]        │
╞═════╪═════════════════════╡
│ A   ┆ 2023-01-01 10:00:00 │
│ A   ┆ 2023-01-01 10:30:00 │
│ A   ┆ 2023-01-01 11:00:00 │
│ B   ┆ 2023-01-01 12:00:00 │
│ B   ┆ 2023-01-01 12:30:00 │
│ B   ┆ 2023-01-01 13:00:00 │
└─────┴─────────────────────┘

3.2 使用 pl.Expr.over() 计算时间间隔

现在，我们将使用 with_columns() 结合 pl.Expr.over() 来创建新的时间间隔列。

# 计算每个 ID 内部的会话时间间隔
result_df = sessions_features.with_columns(
  pl.col("Timestamp")
    .diff() # 计算当前行与上一行的时间差 (结果为 Duration 类型)
    .dt.total_seconds() # 将 Duration 转换为总秒数 (i64 类型)
    .fill_null(0) # 每个 ID 的第一个会话没有前一个会话，diff 结果为 null，此处填充为 0
    .over("ID") # 关键：此表达式在每个唯一的 "ID" 分组内执行
    .alias("time_between_sessions") # 为新列命名
)

print("\n计算时间间隔后的 DataFrame:")
print(result_df)

代码解析:

1. pl.col("Timestamp"): 选择 Timestamp 列。
2. .diff(): 计算当前行 Timestamp 与其前一行 Timestamp 的差值。对于每个 ID 的第一条记录，由于没有前一行，diff() 的结果将是 null。
3. .dt.total_seconds(): diff() 的结果是一个 Duration（持续时间）类型。.dt.total_seconds() 将这个持续时间转换为总秒数，返回一个整数类型（i64）。
4. .fill_null(0): 由于每个 ID 分组的第一条记录的 diff() 结果是 null，我们使用 fill_null(0) 将这些 null 值填充为 0，符合预期结果。
5. .over("ID"): 这是核心步骤。它告诉 Polars，前面的 diff().dt.total_seconds().fill_null(0) 表达式应该在每个唯一的 ID 值所定义的“窗口”内独立执行。这意味着对于 ID='A' 的所有行，diff() 只会考虑 A 组内部的顺序；对于 ID='B' 的行，也只考虑 B 组内部的顺序。
6. .alias("time_between_sessions"): 将新生成的列命名为 time_between_sessions。

预期输出:

计算时间间隔后的 DataFrame:
shape: (6, 3)
┌─────┬─────────────────────┬───────────────────────┐
│ ID  ┆ Timestamp           ┆ time_between_sessions │
│ --- ┆ ---                 ┆ ---                   │
│ str ┆ datetime[μs]        ┆ i64                   │
╞═════╪═════════════════════╪═══════════════════════╡
│ A   ┆ 2023-01-01 10:00:00 ┆ 0                     │
│ A   ┆ 2023-01-01 10:30:00 ┆ 1800                  │
│ A   ┆ 2023-01-01 11:00:00 ┆ 1800                  │
│ B   ┆ 2023-01-01 12:00:00 ┆ 0                     │
│ B   ┆ 2023-01-01 12:30:00 ┆ 1800                  │
│ B   ┆ 2023-01-01 13:00:00 ┆ 1800                  │
└─────┴─────────────────────┴───────────────────────┘

4. 避免 map 或 apply 的原因

在 Polars 中，强烈建议避免使用 map_groups()、apply() 或其他 Python 原生循环的函数，原因如下：

• 性能瓶颈: map 或 apply 函数通常会在 Python 解释器层面操作，无法充分利用 Polars 底层的 Rust 优化和并行化能力。对于大型数据集，这会导致显著的性能下降。
• 惰性求值: Polars 的许多操作都支持惰性求值，这意味着它会构建一个执行计划，只有在需要结果时才真正执行计算。而 map 或 apply 会立即触发计算，打破了惰性求值的优势。
• 内存效率: Polars 的表达式 API 经过优化，通常能以更低的内存消耗完成任务。map 或 apply 可能会在 Python 对象之间进行不必要的转换，增加内存开销。
• 可读性与维护性: 虽然 map 看起来灵活，但当有原生表达式可以实现相同功能时，使用原生表达式通常代码更简洁、意图更明确，也更容易被 Polars 优化器理解和加速。

5. 注意事项与最佳实践

• 数据排序: diff() 函数计算的是当前行与其“前一行”的差值。因此，为了确保 time_between_sessions 的计算逻辑正确，数据在每个 ID 分组内部必须按照 Timestamp 列进行升序排序。如果您的原始数据未排序，请务必在执行 diff() 之前添加 .sort("Timestamp") 操作，例如：

  result_df = sessions_features.sort("ID", "Timestamp").with_columns(
    pl.col("Timestamp").diff().dt.total_seconds().fill_null(0).over("ID").alias("time_between_sessions")
  )

  登录后复制

  在我们的示例中，数据已经按 ID 和 Timestamp 排序，所以可以省略 sort。

• 数据类型: 确保时间戳列是 Polars 的 Datetime 类型。如果它是字符串，需要使用 str.to_datetime() 进行转换。
• 结果类型: dt.total_seconds() 返回的是 i64（64位整数），这对于大多数时间间隔计算是足够的。如果需要浮点数精度（例如毫秒），可以使用 dt.total_milliseconds() 或 dt.total_microseconds()。

总结

通过本教程，我们学习了如何利用 Polars 的 pl.Expr.over() 窗口函数，结合 diff() 和 dt.total_seconds() 等表达式，高效且优雅地计算 DataFrame 中按组划分的时间间隔。这种方法充分利用了 Polars 的高性能特性，避免了传统 map 或 apply 可能带来的性能瓶颈，是处理大规模时间序列分组数据的推荐实践。掌握 over() 函数的使用，将极大地提升您在 Polars 中处理复杂数据转换任务的能力。

以上就是使用 Polars 高效计算 DataFrame 中按 ID 分组的时间间隔的详细内容，更多请关注php中文网其它相关文章！