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() 的情况下，对指定分组（分区）执行聚合或转换操作。
核心思路是：
- pl.col("Timestamp").diff(): 计算 Timestamp 列的差值。diff() 函数会返回当前行与上一行之间的差值。对于每个分组的第一行，diff() 的结果将是 null。
- .dt.total_seconds(): 将时间差（Duration 类型）转换为总秒数（整数或浮点数）。
- .fill_null(0): 将每个分组第一行产生的 null 值填充为 0，因为它们是该 ID 的第一个会话，没有“上一个会话”。
- .over("ID"): 这是关键步骤，它告诉 Polars 在计算 diff() 之前，先根据 ID 列进行分区。这意味着 diff() 操作将在每个 ID 组内独立进行，而不是在整个 DataFrame 上。
- .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 中进行复杂分组计算的能力。
