Как суммировать длительность в кадре данных Polars? ⇐ Python

[Anonymous]

У меня есть следующий фрейм данных:

Код: Выделить всё

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),
],
],
}
)

Код: Выделить всё

print(df)
# Output
shape: (2, 2)
┌─────┬───────────────────────────────────────────────────────────────────┐
│ idx ┆ timestamp                                                         │
│ --- ┆ ---                                                               │
│ i64 ┆ list[datetime[μs]]                                                │
╞═════╪═══════════════════════════════════════════════════════════════════╡
│ 259 ┆ [2023-04-20 01:45:00, 2023-04-20 01:51:07, 2023-04-20 02:29:50]   │
│ 123 ┆ [2023-04-19 06:00:01, 2023-04-19 06:00:17, … 2023-04-19 19:57:52] │
└─────┴───────────────────────────────────────────────────────────────────┘

Мне нужно знать общую продолжительность каждого идентификатора, поэтому я это и делаю:

Код: Выделить всё

df = df.with_columns(
pl.col("timestamp")
.map_elements(lambda x: [x[i + 1] - x[i] for i in range(len(x)) if i + 1 < len(x)])
.alias("duration")
)

Что дает мне:

Код: Выделить всё

shape: (2, 2)
┌─────┬─────────────────────┐
│ idx ┆ duration            │
│ --- ┆ ---                 │
│ i64 ┆ list[duration[μs]]  │
╞═════╪═════════════════════╡
│ 259 ┆ [6m 7s, 38m 43s]    │
│ 123 ┆ [16s, 9s, … 3m 48s] │
└─────┴─────────────────────┘

Теперь в Pandas я бы использовал total_секунды при вызове Apply и суммирования списка, например:

Код: Выделить всё

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)
)

Это даст мне ожидаемый результат:

Код: Выделить всё

print(df[["idx", "duration"]])
# Output

idx  duration
0  259      2690
1  123     50271

Каким идиоматическим способом можно было бы сделать это в Polars?


Подробнее здесь: https://stackoverflow.com/questions/760 ... -dataframe