Having the following DF of user events:

    id                timestamp
0    1  2021-11-23 11:01:00.000
1    1  2021-11-23 11:02:00.000
2    1  2021-11-23 11:10:00.000
3    1  2021-11-23 11:11:00.000
4    1  2021-11-23 11:22:00.000
5    1  2021-11-23 11:40:00.000
6    1  2021-11-23 11:41:00.000
7    1  2021-11-23 11:42:00.000
8    1  2021-11-23 11:43:00.000
9    1  2021-11-23 11:44:00.000
10   2  2021-11-23 11:01:00.000
11   2  2021-11-23 11:02:00.000
12   2  2021-11-23 11:10:00.000
13   2  2021-11-23 11:11:00.000
14   2  2021-11-23 11:22:00.000
15   2  2021-11-23 11:40:00.000
16   2  2021-11-23 11:41:00.000
17   2  2021-11-23 11:42:00.000
18   2  2021-11-23 11:43:00.000
19   2  2021-11-23 11:44:00.000

I calculate the average session time per row as follows:

1. Each session is a sequence of events less than 5 minutes apart.
2. I calculate the number of seconds between the first event in the session and the current event.
3. I then calculate expanding().mean() for each user.

Here is my code:

def average_session_time(**kwargs):
    df = kwargs['df'].copy()
    df['timestamp'] = pd.to_datetime(df.timestamp)
    df['session_grp'] = df.groupby('id').apply(
        lambda x: (x.groupby([pd.Grouper(key="timestamp", freq='5min', origin='start')])).ngroup()).reset_index(
        drop=True).values.reshape(-1)

    # Getting relevant 5min groups
    ng = df.groupby(['id', 'session_grp'])
    df['fts'] = ng['timestamp'].transform('first')
    df['delta'] = df['timestamp'].sub(df['fts']).dt.total_seconds()

    return df.groupby('id')['delta'].expanding().mean().reset_index(drop=True)

And the output is:

0      0.000000
1     30.000000
2     20.000000
3     15.000000
4     12.000000
5     10.000000
6      8.571429
7     15.000000
8     26.666667
9     42.000000
10     0.000000
11    30.000000
12    20.000000
13    15.000000
14    12.000000
15    10.000000
16     8.571429
17    15.000000
18    26.666667
19    42.000000
Name: delta, dtype: float64

The code works fine, but when it's run on a large data set, the performance suffers and it takes a long time to calculate. I tried tweaking the code, but couldn't gain more performance. How can I write this function differently to improve performance?

Here is a Colab with the running code.