CODEWITHSUNDEEP
pythonpandasperformancedataframepandas-groupby
Ataxias
Ataxias

Reputation: 1193

Efficiently applying custom functions to groups in Pandas

This question is about efficiently applying a custom function on logical groups of rows in a Pandas dataframe, which share a value in some column.

Consider the following example of a dataframe:

sID = [1,1,1,2,4,4,5,5,5]
data = np.random.randn(len(sID))
dates = pd.date_range(start='1/1/2018', periods=len(sID))

mydf = pd.DataFrame({"subject_id":sID, "data":data, "date":dates})
mydf['date'][5] += pd.Timedelta('2 days')

which looks like:

       data       date  subject_id
0  0.168150 2018-01-01           1
1 -0.484301 2018-01-02           1
2 -0.522980 2018-01-03           1
3 -0.724524 2018-01-04           2
4  0.563453 2018-01-05           4
5  0.439059 2018-01-08           4
6 -1.902182 2018-01-07           5
7 -1.433561 2018-01-08           5
8  0.586191 2018-01-09           5

Imagine that for each subject_id, we want to subtract from each date the first date encountered for this subject_id. Storing the result in a new column "days_elapsed", the result will look like this:

       data       date  subject_id  days_elapsed
0  0.168150 2018-01-01           1             0
1 -0.484301 2018-01-02           1             1
2 -0.522980 2018-01-03           1             2
3 -0.724524 2018-01-04           2             0
4  0.563453 2018-01-05           4             0
5  0.439059 2018-01-08           4             3
6 -1.902182 2018-01-07           5             0
7 -1.433561 2018-01-08           5             1
8  0.586191 2018-01-09           5             2

One natural way of doing this is by using groupby and apply:

g_df = mydf.groupby('subject_id')
mydf.loc[:, "days_elapsed"] = g_df["date"].apply(lambda x: x - x.iloc[0]).astype('timedelta64[D]').astype(int)

However, if the number of groups (subject IDs) is big (e.g. 10^4), and let's say only 10 times smaller than the length of the dataframe, this very simple operation is really slow.

Is there any faster method?

PS: I have also tried setting the index to subject_id and then using the following list comprehension:

def get_first(series, ind):
    "Return the first row in a group within a series which (group) potentially can span multiple rows and corresponds to a given index"
    group = series.loc[ind]

    if hasattr(group, 'iloc'):
        return group.iloc[0]
    else: # this is for indices with a single element
        return group

hind_df = mydf.set_index('subject_id')
A = pd.concat([hind_df["date"].loc[ind] - get_first(hind_df["date"], ind) for ind in np.unique(hind_df.index)])

However, it's even slower.

Upvotes: 1

Views: 107

Answers (2)

PMende
PMende

Reputation: 5460

mydf['days_elapsed'] = (mydf['date'] - mydf.groupby(['subject_id'])['date'].transform('min')).dt.days

Upvotes: 1

jpp
jpp

Reputation: 164663

You can use GroupBy + transform with first. This should be more efficient as it avoids expensive lambda function calls.

You may also see a performance improvement by working with the NumPy array via pd.Series.values:

first = df.groupby('subject_id')['date'].transform('first').values

df['days_elapsed'] = (df['date'].values - first).astype('timedelta64[D]').astype(int)

print(df)

   subject_id      data       date  days_elapsed
0           1  1.079472 2018-01-01             0
1           1 -0.197255 2018-01-02             1
2           1 -0.687764 2018-01-03             2
3           2  0.023771 2018-01-04             0
4           4 -0.538191 2018-01-05             0
5           4  1.479294 2018-01-08             3
6           5 -1.993196 2018-01-07             0
7           5 -2.111831 2018-01-08             1
8           5 -0.934775 2018-01-09             2

Upvotes: 1

Related Questions