如何用Polars为Pandas DataFrame中每个ID计算行时间差

直接用Polars的**窗口函数（over）**就能实现按ID分组计算会话间隔，这是Polars里最高效的方式，比map_groups性能好很多：

完整代码示例 #

import polars as pl
import pandas as pd

# 创建样本数据
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 = pd.DataFrame(data)
sessions_features = pl.from_pandas(df)

# 转换Timestamp为datetime类型
sessions_features = sessions_features.with_columns(
   pl.col("Timestamp").str.to_datetime()
)

# 新增会话间隔列
sessions_features = sessions_features.with_columns(
    pl.col("Timestamp")
    .diff()  # 计算当前行与上一行的时间差
    .dt.total_seconds()  # 转换为秒数
    .over("ID")  # 按ID分组计算
    .fill_null(0)  # 每个ID的第一行填充为0
    .cast(pl.Int64)  # 转为整数类型，匹配预期结果的i64
    .alias("time_between_sessions")
)

print(sessions_features)

关键逻辑说明 #

• over("ID")：指定按ID分组执行后续的diff计算，相当于给每个ID单独计算时间差
• diff()：计算当前行Timestamp与同组上一行的差值，每个ID的第一行会得到null
• fill_null(0)：把每个ID的第一行间隔值填充为0，符合预期结果
• 这种方式不需要拆分数据或用map_groups，Polars会自动优化执行效率，处理大数据量也很顺畅

输出结果 #

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

内容的提问来源于stack exchange，提问作者Camilo Piñón