How to speed up rolling diff in Pandas when applied to segments of DataFrame

Question

I have the following code

from random import randrange, randint
from datetime import timedelta, datetime

def random_date(start, end):
    delta = end - start
    int_delta = (delta.days * 24 * 60 * 60) + delta.seconds
    random_second = randrange(int_delta)
    return start + timedelta(seconds=random_second)

from datetime import datetime
d1 = datetime.strptime('1/1/2008 1:30 PM', '%m/%d/%Y %I:%M %p')
d2 = datetime.strptime('1/1/2009 4:50 AM', '%m/%d/%Y %I:%M %p')

num_rows = 40000
num_users = 10000
events = ['page_view', 'session_start']

random_timestamps = [random_date(d1, d2).timestamp() for i in range(num_rows)]
random_users = [randint(0, num_users) for i in range(num_rows)]
random_events = [events[randint(0, 1)] for i in range(num_rows)]
df = pd.DataFrame({'event_timestamp': random_timestamps,
                   'user_pseudo_id': random_users,
                   'event_name': random_events
                   })


user_ids = df.user_pseudo_id.unique()
df.sort_values(['event_timestamp', 'event_name'], ascending=[True, False], inplace=True)

for user_id in user_ids:
  df.loc[df.user_pseudo_id == user_id, 'event_timestamp_diff'] = df[df.user_pseudo_id == user_id]['event_timestamp'].rolling(window=2).apply(np.diff)

df.event_timestamp_diff.fillna(0, inplace=True)

The df is events (new session, pageview, etc) from Google Analytics 4 tied to specific users pseudo_user_id . What I want to accomplish is to calculate timestamp diffs from prior events only for events tied to a specific user. Essentially, how long after the prior event did this event occur, for this user.

I have used rolling in very limited ways previously and was hoping there was either another option (eg shift ) or grouping logic that would help speed this up for instances where there are a large number of users.

Answer 1

for user_id in user_ids:
  df.loc[df.user_pseudo_id == user_id, 'event_timestamp_diff'] = df[df.user_pseudo_id == user_id]['event_timestamp'].rolling(window=2).apply(np.diff)

can be replaced with

df['event_timestamp_diff'] = df.groupby('user_pseudo_id')['event_timestamp'].rolling(window=2).apply(np.diff).reset_index(0,drop=True)