Python Pandas: Trying to speed-up a per row per date in date_range operation

Viewed 216

I have a dataframe of the following form where each row corresponds to a job run on a machine:

import pandas as pd
df = pd.DataFrame({
    'MachineID': [4, 3, 2, 2, 1, 1, 5, 3],
    'JobStartDate': ['2020-01-01', '2020-01-01', '2020-01-01', '2020-01-01', '2020-01-02', '2020-01-03', '2020-01-01', '2020-01-03'], 
    'JobEndDate': ['2020-01-03', '2020-01-03', '2020-01-04', '2020-01-02', '2020-01-04', '2020-01-05', '2020-01-02', '2020-01-04'], 
    'IsTypeAJob': [1, 1, 0, 1, 0, 0, 1, 1]
    })
    
df
>>>    MachineID JobStartDate JobEndDate  IsTypeAJob
0          4      2020-01-01  2020-01-03      1
1          3      2020-01-01  2020-01-03      1
2          2      2020-01-01  2020-01-04      0
3          2      2020-01-01  2020-01-02      1
4          1      2020-01-02  2020-01-04      0
5          1      2020-01-03  2020-01-05      0
6          5      2020-01-01  2020-01-02      1
7          3      2020-01-03  2020-01-04      1

In my data there are two types of jobs that can be run on a machine, either type A or type B. My goal is to count the number of type A and type B jobs per machine per day. Thus the desired result would look something like

        MachineID    Date         TypeAJobs  TypeBJobs
0           1     2020-01-02          0          1
1           1     2020-01-03          0          2
2           1     2020-01-04          0          2
3           1     2020-01-05          0          1
4           2     2020-01-01          1          1
5           2     2020-01-02          1          1
6           2     2020-01-03          0          1
7           2     2020-01-04          0          1
8           3     2020-01-01          1          0
9           3     2020-01-02          1          0
10          3     2020-01-03          2          0
11          3     2020-01-04          1          0
12          4     2020-01-01          1          0
13          4     2020-01-02          1          0
14          4     2020-01-03          1          0
15          5     2020-01-01          1          0
16          5     2020-01-02          1          0

I have tried approaches found here and here with a resample() and apply() method, but the computing time is too slow. This has to do with the fact that some date ranges span multiple years in my set, meaning one row can blow up into 2000+ new rows during resampling (my data contains around a million rows to begin with). Thus something like creating a new machine/date row for each date in the range of a certain job is too slow (with the goal of doing a group_by(['MachineID', 'Date']).sum() at the end).

I am currently thinking about a new approach where I begin by grouping by MachineID then finding the earliest job start date and latest job end date for that machine. Then I could create a date range of days between these two dates (incrementing by day) which I would use to index a new per machine data frame. Then for each job for that MachineID I could potentially sum over a range of dates, ie in pseudocode:

df['TypeAJobs'][row['JobStartDate']:row['JobEndDate']] += 1 if it is a type A job or

df['TypeBJobs'][row['JobStartDate']:row['JobEndDate']] += 1 otherwise.

This seems like it would avoid creating a bunch of extra rows for each job as now we are creating extra rows for each machine. Furthermore, the addition operations seem like they would be fast since we are adding to an entire slice of a series at once. However, I don't know if something like this (indexing by date) is possible in Pandas. Maybe there is some conversion that can be done first? After doing the above, ideally I would have a number of data frames similar to the desired result but only with one MachineID, then I would concatenate these data frames to get the result.

I would love to hear any suggestions about the feasibility/effectiveness of this approach or another potential algorithm. Thanks so much for reading!

2 Answers

IIUC, try using pd.date_range and explode to create 'daily' rows, then groupby dates and IsTypeAJob and rename columns:

df_out = df.assign(JobDates=df.apply(lambda x: pd.date_range(x['JobStartDate'], 
                                                             x['JobEndDate'], freq='D'), 
                                     axis=1))\
           .explode('JobDates')

df_out = df_out.groupby([df_out['MachineID'], 
                         df_out['JobDates'].dt.floor('D'), 
                         'IsTypeAJob'])['MachineID'].count()\
               .unstack()\
               .rename(columns={0:'TypeBJobs', 1:'TypeAJobs'})\
               .fillna(0).reset_index()

df_out

Output:

IsTypeAJob  MachineID   JobDates  TypeBJobs  TypeAJobs
0                   1 2020-01-02        1.0        0.0
1                   1 2020-01-03        2.0        0.0
2                   1 2020-01-04        2.0        0.0
3                   1 2020-01-05        1.0        0.0
4                   2 2020-01-01        1.0        1.0
5                   2 2020-01-02        1.0        1.0
6                   2 2020-01-03        1.0        0.0
7                   2 2020-01-04        1.0        0.0
8                   3 2020-01-01        0.0        1.0
9                   3 2020-01-02        0.0        1.0
10                  3 2020-01-03        0.0        2.0
11                  3 2020-01-04        0.0        1.0
12                  4 2020-01-01        0.0        1.0
13                  4 2020-01-02        0.0        1.0
14                  4 2020-01-03        0.0        1.0
15                  5 2020-01-01        0.0        1.0
16                  5 2020-01-02        0.0        1.0

pd.concat([pd.DataFrame({'JobDates':pd.date_range(r.JobStartDate, r.JobEndDate, freq='D'),
                        'MachineID':r.MachineID, 
                        'IsTypeAJob':r.IsTypeAJob}) for i, r in df.iterrows()])

Here is another way to do the job, the idea is similar to use str.get_dummies on both columns start and end, but done with array broadcasting. Use cumsum do get one between start and end and 0 otherwise. Create a dataframe with the columns as dates and the index as both Machine and Type. Then do similar operation than the answer from @Scott Boston to get the expected output shape.

#get all possible dates
dr = pd.date_range(df['JobStartDate'].min(), 
                   df['JobEndDate'].max()).strftime("%Y-%m-%d").to_numpy()

df_ = (pd.DataFrame(
            np.cumsum((df['JobStartDate'].to_numpy()[:, None] == dr).astype(int) 
                       - np.pad(df['JobEndDate'].to_numpy()[:, None]==dr,((0,0),(1,False)), 
                                mode='constant')[:, :-1],  # pad is equivalent to shift along columns
                      axis=1), 
            index=pd.MultiIndex.from_frame(df[['MachineID', 'IsTypeAJob']]), 
            columns=dr,)
         .sum(level=['MachineID', 'IsTypeAJob']) #equivalent to groupby(['MachineID', 'IsTypeAJob']).sum()
         .replace(0, np.nan) #to remove extra dates per original row during the stack
         .stack()
         .unstack(level='IsTypeAJob', fill_value=0)
         .astype(int)
         .reset_index()
         .rename_axis(columns=None)
         .rename(columns={'level_1':'Date', 0:'TypeBJobs', 1:'TypeAJobs'})
      )

and you get

    MachineID        Date  TypeBJobs  TypeAJobs
0           1  2020-01-02          1          0
1           1  2020-01-03          2          0
2           1  2020-01-04          2          0
3           1  2020-01-05          1          0
4           2  2020-01-01          1          1
5           2  2020-01-02          1          1
6           2  2020-01-03          1          0
7           2  2020-01-04          1          0
8           3  2020-01-01          0          1
9           3  2020-01-02          0          1
10          3  2020-01-03          0          2
11          3  2020-01-04          0          1
12          4  2020-01-01          0          1
13          4  2020-01-02          0          1
14          4  2020-01-03          0          1
15          5  2020-01-01          0          1
16          5  2020-01-02          0          1
Related