[Python] How to filter sequential event data according to whether record is followed by...

I have some data with a timestamp column t, an event category column cat, and a user_id column. cat can take n values, including value A.

I want to select records which are followed (not necessarily immediately) by an A event within 5 minutes (for the same user).

In SQL, I would do the following:

select
*
from self
qualify count_if(category = 'A') over (partition by user_id order by t
range between current row and interval '5 minutes' following) > 0


How can you do this in Polars? (other than switching to .sql method).

More generally, does anyone have good pointers on the appropriate mental model from translating from a <function> over (partition by <X> order by <Y> rows / range between <A> and <B>) SQL syntax to Polars?

I tried the following:

df.filter(
pl.col('cat').any().over('user_id').rolling(
index_column='t',
period = '5m',
offset = '0m',
closed = 'both')
)


and variations thereof, but got an error 'window expression not allowed in aggregation'.

I also tried using .rolling directly on the DataFrame, but this results in an aggregation which I don't want (I want to preserve the grain of the table).

Continue reading...