[Python] How can I perform a calculation on a rolling window over a partition in polars?

I have a Dataset containing GPS Coordinates of a few planes. I would like to calculate the bearing of each plane at every point in time.

The Dataset as among others these columns:

1. event_uid
2. plane_no
3. timestamp
4. gps_lat
5. gps_lon

I can calculate the bearing for a single plane_no by doing something as follows:

lf.sort("timestamp")
.with_row_index('new_index')
.with_columns(
pl.struct(pl.col('gps_lat'), pl.col('gps_lon'))
.rolling('new_index', period='2i')
.shift(-1)
.alias("latlon")
).with_columns(
pl.col("latlon")
.map_elements(lambda val: calculate_bearing(val[0]['gps_lat'], val[0]['gps_lon'], val[1]['gps_lat'], val[1]['gps_lon']), return_dtype=pl.Float64)
.alias("bearing")
).drop(["latlon", "new_index"])


In order to work on my entire dataset though I need to run this on every plane_no as a partition. I can't combine rolling with over in polars. Is there an alternative? Any possible improvement ideas or overall comments are also welcome.

Continue reading...