How do I vectorize a Pandas function that counts rows belonging to a group and falling between two dates?

user_3877029 0
Viewed 72

I have the following table and I want to count the number of active jobs, per client, on each day in 2020. A job is active if the date falls on or between its start_date and end_date.

job client start_date end_date
AA001 ALPHA 2020/12/19 2020/12/28
AA002 ALPHA 2020/04/03 2020/10/10
AA003 BRAVO 2020/10/11 2020/10/11
AA004 CHARLIE 2020/04/06 2020/11/15
AA005 ALPHA 2020/04/01 2020/04/30
AA006 CHARLIE 2020/05/01 2020/06/03
AA007 BRAVO 2020/06/04 2020/06/17
AA008 BRAVO 2020/06/18 2020/07/01
AA009 CHARLIE 2020/07/02 2020/08/04
AA010 ALPHA 2020/05/05 2020/08/06
AA011 BRAVO 2020/10/12 2020/11/04

For instance, here is how many jobs were active for client ALPHA at the beginning of April:

Date Client Active jobs
ALPHA 2020-04-01 1
ALPHA 2020-04-02 1
ALPHA 2020-04-03 2
ALPHA 2020-04-04 2
ALPHA 2020-04-05 2
ALPHA 2020-04-06 2
ALPHA 2020-04-07 2
ALPHA 2020-04-08 2
ALPHA 2020-04-09 2
ALPHA 2020-04-10 2

I can solve this problem using nested loops, e.g.

groups = df.groupby(["client"])   
dates = pd.date_range('2020-01-01','2020-12-01', freq='D')   
 
for client, jobs in groups:  
    for date in dates:  
        active_jobs = jobs.loc[(jobs.start_date <= date) & (jobs.end_date >= date)]  
        print(date,client,len(active_jobs))

(Explanation: group rows by client, construct a list of dates, then for each date for each client, find/count the rows where start_date <= date and end_date >= date.)

Of course my real data is much larger than this and looping is very inefficient. How do I rewrite my query to take advantage of vectorization?

pandas vectorization python
1 Answers

Approach with broadcasting

Check for the inclusion of dates between the start_date and end_date columns, this will create a boolean mask, now we create a new dataframe from this mask and assign the column names to the corresponding dates, then group this dataframe by client and aggregate using sum to count the number of active jobs for each client on each day

start, end = df[['start_date', 'end_date']].to_numpy().T
dates = pd.date_range('2020-01-01','2020-12-01', freq='D').to_numpy()

m = (start[:, None] <= dates) & (end[:, None] >= dates)
s = pd.DataFrame(m, columns=dates).groupby(df['client']).sum().stack()

After stacking the resulting series containing the counts of active_jobs will look like

>>> s

client             
ALPHA    2020-01-01    0
         2020-01-02    0
         2020-01-03    0
         2020-01-04    0
         2020-01-05    0
                      ..
CHARLIE  2020-11-27    0
         2020-11-28    0
         2020-11-29    0
         2020-11-30    0
         2020-12-01    0
Length: 1008, dtype: int64

Examining the active jobs for client ALPHA for the month of APRIL

>>> s.loc[pd.IndexSlice['ALPHA', '2020-04-01':]]

client            
ALPHA   2020-04-01    1
        2020-04-02    1
        2020-04-03    2
        2020-04-04    2
        2020-04-05    2
        2020-04-06    2
        2020-04-07    2
        2020-04-08    2
        2020-04-09    2
        2020-04-10    2
        2020-04-11    2
        2020-04-12    2
        2020-04-13    2
        2020-04-14    2
        2020-04-15    2
        2020-04-16    2
        2020-04-17    2
        2020-04-18    2
        2020-04-19    2
        2020-04-20    2
        2020-04-21    2
        2020-04-22    2
        2020-04-23    2
        2020-04-24    2
        2020-04-25    2
        2020-04-26    2
        2020-04-27    2
        2020-04-28    2
        2020-04-29    2
        2020-04-30    2
dtype: int64

PS: Although using broadcasting is faster but it will require sufficient amount of memory to hold the intermediate boolean mask in memory. One more thing you also have to convert the start_date and end_date column to pandas datetime format prior to using this approach

user_128331660
Related