如何在Polars DataFrame中计算时长总和？

问题背景 #

现有如下Polars DataFrame：

import datetime
import polars as pl

df = pl.DataFrame(
    {
        "idx": [259, 123],
        "timestamp": [
            [
                datetime.datetime(2023, 4, 20, 1, 45),
                datetime.datetime(2023, 4, 20, 1, 51, 7),
                datetime.datetime(2023, 4, 20, 2, 29, 50),
            ],
            [
                datetime.datetime(2023, 4, 19, 6, 0, 1),
                datetime.datetime(2023, 4, 19, 6, 0, 17),
                datetime.datetime(2023, 4, 19, 6, 0, 26),
                datetime.datetime(2023, 4, 19, 19, 53, 29),
                datetime.datetime(2023, 4, 19, 19, 54, 4),
                datetime.datetime(2023, 4, 19, 19, 57, 52),
            ],
        ],
    }
)

目标是计算每个idx对应的时间戳列表中，相邻时间差的总秒数（转为整数）。用户最初尝试用map_elements生成差值列表，但需要进一步求和；对应的Pandas实现为：

df["duration"] = (
    df["timestamp"]
    .apply(
        lambda x: sum(
            [(x[i + 1] - x[i]).total_seconds() for i in range(len(x)) if i + 1 < len(x)]
        )
    )
    .astype(int)
)

期望结果：

idx  duration
0  259      2690
1  123     50271

Polars的惯用实现 #

方法一：利用首尾时间差（最优方案） #

相邻时间差的总和等价于列表最后一个时间戳减去第一个时间戳，直接通过Polars的列表内置函数计算，完全向量化，性能最优：

df = df.with_columns(
    (pl.col("timestamp").list.last() - pl.col("timestamp").list.first())
    .dt.total_seconds()
    .cast(pl.Int32)
    .alias("duration")
)

方法二：显式计算相邻差值再求和 #

如果需要显式遍历计算每个相邻时间差后求和，推荐用list.eval替代map_elements（避免Python lambda的性能损耗）：

df = df.with_columns(
    pl.col("timestamp")
    .list.eval(
        pl.element().diff().drop_nulls().dt.total_seconds().sum()
    )
    .list.first()
    .cast(pl.Int32)
    .alias("duration")
)

两种方法执行后都能得到期望结果：

shape: (2, 2)
┌─────┬──────────┐
│ idx ┆ duration │
│ --- ┆ ---      │
│ i64 ┆ i32      │
╞═════╪══════════╡
│ 259 ┆ 2690     │
│ 123 ┆ 50271    │
└─────┴──────────┘

注意：尽量避免使用map_elements，因为它会逐行调用Python函数，无法利用Polars的向量化优化，大数据量下性能差距明显。

内容的提问来源于stack exchange，提问作者Laurent