Vectorizing a variable length look-ahead loop in pandas - pandas

Question

This is a very simplified version of my data:
+----+---------+---------------------+
| | user_id | seconds_since_start |
+----+---------+---------------------+
| 0 | 1 | 10 |
| 1 | 1 | 12 |
| 2 | 1 | 15 |
| 3 | 1 | 52 |
| 4 | 1 | 60 |
| 5 | 1 | 67 |
| 6 | 1 | 120 |
| 7 | 2 | 55 |
| 8 | 2 | 62 |
| 9 | 2 | 105 |
| 10 | 3 | 200 |
| 11 | 3 | 206 |
+----+---------+---------------------+
And this is the data I would like to produce:
+----+---------+---------------------+-----------------+------------------+
| | user_id | seconds_since_start | session_ordinal | session_duration |
+----+---------+---------------------+-----------------+------------------+
| 0 | 1 | 10 | 1 | 5 |
| 1 | 1 | 12 | 1 | 5 |
| 2 | 1 | 15 | 1 | 5 |
| 3 | 1 | 52 | 2 | 15 |
| 4 | 1 | 60 | 2 | 15 |
| 5 | 1 | 67 | 2 | 15 |
| 6 | 1 | 120 | 3 | 0 |
| 7 | 2 | 55 | 1 | 7 |
| 8 | 2 | 62 | 1 | 7 |
| 9 | 2 | 105 | 2 | 0 |
| 10 | 3 | 200 | 1 | 6 |
| 11 | 3 | 206 | 1 | 6 |
+----+---------+---------------------+-----------------+------------------+
My notion of a session is a group of events from a single user which occur not more than 10 seconds apart, and a session's duration is defined as the difference between the first event in the session and the last event (in seconds).
I have written working Python that achieves what I want.
import pandas as pd
events_data = [[1, 10], [1, 12], [1, 15], [1, 52], [1, 60], [1, 67], [1, 120],
[2, 55], [2, 62], [2, 105],
[3, 200], [3, 206]]
events = pd.DataFrame(data=events_data, columns=['user_id', 'seconds_since_start'])
def record_session(index_range, ordinal, duration):
for i in index_range:
events.at[i, 'session_ordinal'] = ordinal
events.at[i, 'session_duration'] = duration
session_indexes = []
current_user = previous_time = session_start = -1
session_num = 0
for i, row in events.iterrows():
if row['user_id'] != current_user or (row['seconds_since_start'] - previous_time) > 10:
record_session(session_indexes, session_num, previous_time - session_start)
session_indexes = [i]
session_num += 1
session_start = row['seconds_since_start']
if row['user_id'] != current_user:
current_user = row['user_id']
session_num = 1
previous_time = row['seconds_since_start']
session_indexes.append(i)
record_session(session_indexes, session_num, previous_time - session_start)
My problem is the length of time this takes to run. As I said, this is a very simplified version of my data, my actual data has 70,000,000 rows. Is there a way to vectorize (and thus speed-up) algorithms like this that formulate additional columns based on variable length look-aheads?

Answer 1

You could try:
# Create a helper boolean Series
s = df.groupby('user_id')['seconds_since_start'].diff().gt(10)
df['session_ordinal'] = s.groupby(df['user_id']).cumsum().add(1).astype(int)
df['session_duration'] = (df.groupby(['user_id', 'session_ordinal'])['seconds_since_start']
.transform(lambda x: x.max() - x.min()))
[output]
user_id seconds_since_start session_ordinal session_duration
0 1 10 1 5
1 1 12 1 5
2 1 15 1 5
3 1 52 2 15
4 1 60 2 15
5 1 67 2 15
6 1 120 3 0
7 2 55 1 7
8 2 62 1 7
9 2 105 2 0
10 3 200 1 6
11 3 206 1 6

Answer 2

Chris A's answer here is great. It contains several techniques or calls I was unfamiliar with. This answer copies his and adds copious annotations.
We start by building a helper Boolean series. This series records which events start additional sessions for any user. This is OK as a Boolean series because in numeric contexts they behave like the integers 0 and 1 (quoting from here). Let's put the series together bit by bit.
starts_session = events.groupby('user_id')['seconds_since_start'].diff().gt(10)
First we group events by user_id (documentation) and then choose the 'seconds_since_start' column and call diff (documentation) on that. The result of events.groupby('user_id')['seconds_since_start'].diff()is
+----+----------------------+
| | seconds_since_start |
+----+----------------------+
| 0 | NaN |
| 1 | 2.0 |
| 2 | 3.0 |
| 3 | 37.0 |
| 4 | 8.0 |
| 5 | 7.0 |
| 6 | 53.0 |
| 7 | NaN |
| 8 | 7.0 |
| 9 | 43.0 |
| 10 | NaN |
| 11 | 6.0 |
+----+----------------------+
I can see that the start of each group is already picking up the correct NaN difference as there's no previous event from that user to give a delta from.
Then using the element-wise greater than gt(10) (documentation) we get
+----+----------------------+
| | seconds_since_start |
+----+----------------------+
| 0 | False |
| 1 | False |
| 2 | False |
| 3 | True |
| 4 | False |
| 5 | False |
| 6 | True |
| 7 | False |
| 8 | False |
| 9 | True |
| 10 | False |
| 11 | False |
+----+----------------------+
(N.B. The column heading is odd, but it is not used and so does not matter.)
events['session_ordinal'] = starts_session.groupby(events['user_id']).cumsum().add(1).astype(int)
We then re-group starts_session by the user_ids in events and then do the cumulative sum cumsum (documentation) over each group. The grouping does the work for us here ensuring that each user's events are restarted at zero. We need the session ordinal to start at 1 not zero so we simply add one add(1) (documentation) and we cast them to int as none of them are NaN astype(int) (documentation). This gives the derived session_ordinal column I wanted.
events['session_duration'] = events.groupby(['user_id', 'session_ordinal'])['seconds_since_start'].transform(lambda x: x.max() - x.min())
To derive each session's duration we first group events by both the user_id and the new session_ordinal, i.e. we group them into sessions. Using transform (documentation) we find the minimum and maximum value of seconds_since_start for each group (i.e. each session) and the difference between them is the session's duration. This pattern, applying transform to grouped data is used extensively in the split-apply-combine process.
Thanks Chris.