polars datetime 5 minutes floor

Question

我有 polars dataframe，其时间戳 folumn 类型为 datetime[ns]，其值为

2023-03-08 11:13:07.831

我想使用极地效率将时间戳四舍五入到 5 分钟底。

现在我做的：

import arrow

def timestamp_5minutes_floor(ts: int) -> int:
    return int(arrow.get(ts).timestamp() // 300000 * 300000)

df.with_columns([
    pl.col("timestamp").apply(lambda x: timestamp_5minutes_floor(x)).alias("ts_floor")
    ])

很慢。如何改进？

Answer 1

您可以尝试使用

.dt.truncate

：使用示例数据框

df = pl.DataFrame({
    "ts": ["2023-03-08 11:01:07.831", "2023-03-08 18:09:01.007"]
}).select(pl.col("ts").str.strptime(pl.Datetime, "%Y-%m-%d %H:%M:%S%.3f"))

┌─────────────────────────┐
│ ts                      │
│ ---                     │
│ datetime[ms]            │
╞═════════════════════════╡
│ 2023-03-08 11:01:07.831 │
│ 2023-03-08 18:09:01.007 │
└─────────────────────────┘

这个

df = df.select(pl.col("ts").dt.truncate("5m"))

结果

┌─────────────────────┐
│ ts                  │
│ ---                 │
│ datetime[ms]        │
╞═════════════════════╡
│ 2023-03-08 11:00:00 │
│ 2023-03-08 18:05:00 │
└─────────────────────┘