Polars分组时间序列差值计算技巧

1. 背景与挑战

在数据分析任务中，我们经常需要处理时间序列数据，例如计算用户会话之间的时间间隔、订单处理时长等。当数据按某个标识符（如用户 ID）分组时，挑战在于如何高效地计算每个组内连续记录之间的时间差。传统的 Pandas groupby().apply() 方法在处理大数据集时可能效率低下。对于 Polars 而言，虽然 map_groups() 提供了一定的灵活性，但其性能通常不如原生表达式。

我们的目标是，给定一个包含 ID 和 Timestamp 列的 DataFrame，为每个 ID 计算一个新列 time_between_sessions，表示当前会话与上一个会话之间的时间间隔。

2. 准备数据

首先，我们创建一个示例 Pandas DataFrame，并将其转换为 Polars DataFrame。确保 Timestamp 列被正确解析为日期时间类型，这是进行时间计算的基础。

import polars as pl
import pandas as pd

# 创建一个示例 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)

# 将 Pandas DataFrame 转换为 Polars DataFrame
sessions_features = pl.from_pandas(df_pandas)

# 确保 Timestamp 列是日期时间类型
sessions_features = sessions_features.with_columns(
   pl.col("Timestamp").str.to_datetime()
)

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

输出：

原始 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. 使用 pl.Expr.over() 高效计算时间差

Polars 提供了强大的表达式 API 和窗口函数，其中 pl.Expr.over() 是解决此类分组计算问题的理想选择。它允许我们在不显式使用 group_by() 的情况下，对指定分组（分区）执行聚合或转换操作。

核心思路是：

1. pl.col("Timestamp").diff(): 计算 Timestamp 列的差值。diff() 函数会返回当前行与上一行之间的差值。对于每个分组的第一行，diff() 的结果将是 null。
2. .dt.total_seconds(): 将时间差（Duration 类型）转换为总秒数（整数或浮点数）。
3. .fill_null(0): 将每个分组第一行产生的 null 值填充为 0，因为它们是该 ID 的第一个会话，没有“上一个会话”。
4. .over("ID"): 这是关键步骤，它告诉 Polars 在计算 diff() 之前，先根据 ID 列进行分区。这意味着 diff() 操作将在每个 ID 组内独立进行，而不是在整个 DataFrame 上。
5. .alias("time_between_sessions"): 为新生成的列指定一个有意义的名称。

以下是实现代码：

# 计算每个 ID 的会话之间的时间差
sessions_features_with_time_diff = sessions_features.with_columns(
  pl.col("Timestamp")
    .diff()
    .dt.total_seconds()
    .fill_null(0)
    .over("ID")
    .alias("time_between_sessions")
)

print("\n计算时间差后的 Polars DataFrame:")
print(sessions_features_with_time_diff)

预期输出：

计算时间差后的 Polars 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                  │
└─────┴─────────────────────┴───────────────────────┘

从输出可以看出，对于 ID 'A' 和 'B'，它们各自的第一个会话的 time_between_sessions 值为 0，而后续会话的时间间隔则被正确计算为 1800 秒（30 分钟）。

4. 注意事项与最佳实践

• 避免 map 或 apply 函数: 尽可能使用 Polars 的原生表达式 API，如 over()、group_by() 配合表达式，而不是 map_groups() 或 apply()。原生表达式经过高度优化，能够利用 Polars 的并行计算能力，从而在处理大数据集时提供卓越的性能。
• 数据类型: 确保时间戳列为 Polars 的 Datetime 类型。如果数据源是字符串，需要使用 str.to_datetime() 进行转换。
• 排序: diff() 操作依赖于行的顺序。在进行时间差计算之前，如果数据未按 ID 和 Timestamp 排序，应先使用 sort(["ID", "Timestamp"]) 进行排序，以确保计算的正确性。本示例中，输入数据已是按 ID 和时间戳排序的。
• 空值处理: diff() 必然会在每个分组的第一行产生 null 值。根据业务需求，可以使用 fill_null() 将其填充为 0 或其他有意义的值。

5. 总结

通过利用 Polars 的 pl.Expr.over() 窗口函数，我们可以优雅且高效地在每个分组内执行时间序列的差值计算。这种方法避免了传统循环或低效的 apply 操作，充分发挥了 Polars 的列式和并行处理优势，是处理大规模时间序列数据时的首选方案。掌握 over() 的使用，将极大提升你在 Polars 中进行复杂分组计算的能力。