[Python] Polars conditional merge of rows

I have some time data, that I want to merge if two intervals are close enough (and the type is the same), currently I'm doing sth. like this:

Note: time_start, time_end are datetimes, type a categorical string

buffer = dt.timedelta(minutes=15)
_temp_df = (
data.sort("time_start").select([
pl.col("*"),
((pl.col("time_start").shift(-1)) - pl.col("time_end")).fill_null(0).apply(lambda x: (dt.timedelta(seconds=x // 1000))).alias("offset"),
pl.col("type").shift(-1).alias("type_next")
]).select([
pl.col("time_start"), pl.col("time_end"), pl.col("type"),
((pl.col("offset") < buffer) & (pl.col("type") == pl.col("type_next"))).alias("can_merge")
])
).to_dicts()

_after_df = []
if len(_temp_df) != 0:
last_data = _temp_df[0]
for i in range(1, len(_temp_df)):
_dp = _temp_df
if last_data['can_merge']:
last_data['time_end'] = _dp['time_end']
last_data['can_merge'] = _dp['can_merge']
else:
_after_df.append(last_data)
last_data = _dp

_after_df.append(last_data)

data = pl.from_dicts(_after_df)


Whilst this works, it feels extremely off, that I have to convert polars -> dicts -> polars, is there a way to do this over the aggregations? The group_by_dynamic and rolling seem to only accepts times, but I want to merge consecutive rows by offset.

This is Polars: 0.13.59

Continue reading...