Is there another way to add a column to a groupby when using pandas?

Question

I defined a session as a set of songs that play that don’t have a break of at least 15 minutes among them. My goal is to find the average session lengths for each user.

So far, I have used python and pandas to group the below data by user id, and then sort each of those groups by the start time stamp.

Input data:

My code so far:

start_end_song.groupby('user_id').apply(lambda x: x.sort_values('start_timestamp'))

Output of above code:

Next I want to calculate the break between the end time stamp of the first song and the start of the next time stamp.

However, this doesn't work:

start_end_song.groupby('user_id')\
.apply(lambda x: x.sort_values('start_timestamp'))\
.apply(lambda x: x['break']= start_end_song['end_timestamp']- start_end_song['start_timestamp'].shift(-1))

SyntaxError: lambda cannot contain assignment

Is there another way to add a column to a groupby?

Answer 1

You can use pandas.DataFrame.shift and pandas.DataFrame.cumsum to get 'islands' songs:

>>> df = pd.DataFrame({'user_id': [1, 1, 1, 1, 2, 2, 2, 2], 'start_timestamp': [1, 3, 20, 26, 1, 5, 40, 42], 'end_timestamp': [2, 4, 25, 27, 2, 10, 41, 50]}, columns=['user_id', 'start_timestamp', 'end_timestamp'])
>>> df
   user_id  start_timestamp  end_timestamp
0        1                1              2
1        1                3              4
2        1               20             25
3        1               26             27
4        2                1              2
5        2                5             10
6        2               40             41
7        2               42             50

>>> df['session_break'] = (df['start_timestamp'] - df.groupby('user_id')['end_timestamp'].shift(1) >= 15).astype('int')
>>> df
   user_id  start_timestamp  end_timestamp  session_break
0        1                1              2              0
1        1                3              4              0
2        1               20             25              1
3        1               26             27              0
4        2                1              2              0
5        2                5             10              0
6        2               40             41              1
7        2               42             50              0
>>> df['session_label'] = df.groupby('user_id')['session_break'].cumsum()
>>> df
   user_id  start_timestamp  end_timestamp  session_break  session_label
0        1                1              2              0              0
1        1                3              4              0              0
2        1               20             25              1              1
3        1               26             27              0              1
4        2                1              2              0              0
5        2                5             10              0              0
6        2               40             41              1              1
7        2               42             50              0              1

update

To get average session duration you can do this:

>>> g = df.groupby(['user_id', 'session_label']).agg({'end_timestamp' : np.max, 'start_timestamp' : np.min})
>>> g
                       start_timestamp  end_timestamp
user_id session_label                                
1       0                            1              4
        1                           20             27
2       0                            1             10
        1                           40             50

>>> (g['end_timestamp'] - g['start_timestamp']).groupby(level=0).mean()
user_id
1    5.0
2    9.5