Advertisements
I have the following 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),
],
],
}
)
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] │
└─────┴───────────────────────────────────────────────────────────────────┘
I want to know the total duration of each id, so I do:
df = df.with_columns(
pl.col("timestamp")
.apply(lambda x: [x[i + 1] - x[i] for i in range(len(x)) if i + 1 < len(x)])
.alias("duration")
)
Which gives me:
shape: (2, 2)
┌─────┬─────────────────────┐
│ idx ┆ duration │
│ --- ┆ --- │
│ i64 ┆ list[duration[μs]] │
╞═════╪═════════════════════╡
│ 259 ┆ [6m 7s, 38m 43s] │
│ 123 ┆ [16s, 9s, … 3m 48s] │
└─────┴─────────────────────┘
Now, in Pandas, I would have used total_seconds when calling apply
and sum
the list, like this:
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)
)
Which would give me the expected result:
print(df[["idx", "duration"]])
# Output
idx duration
0 259 2690
1 123 50271
What would be the equivalent, idiomatic way, to do this in Polars?
>Solution :
There is an arr.diff
method for list types, which then can be summed, and the total seconds can be calculated with dt.seconds
:
df.select(
"idx",
duration=pl.col("timestamp")
.arr.diff(null_behavior="drop")
.arr.sum()
.dt.seconds(),
)
┌─────┬──────────┐
│ idx ┆ duration │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪══════════╡
│ 259 ┆ 2690 │
│ 123 ┆ 50271 │
└─────┴──────────┘
An equivalent expression if you really don’t need the intermediate durations at all, that should perform better, would be subtracting the first element of the list from the last:
duration=(
pl.col("timestamp").arr.last() - pl.col("timestamp").arr.first()
).dt.seconds()