What is the most efficient way to check several conditions in columns in a pandas dataframe?

user_9202025 0
Viewed 129

I am working through a pandas dataframe with three relevant columns and 2.7 million rows. The structure is: 

key VisitLink   dx_filter   time
0   1   ddcde14 1   100
1   2   abcde11 1   140
2   3   absdf12 1   50
3   4   ddcde14 0   125
4   5   ddcde14 1   140

data = [[1,'ddcde14',1,100],[2,'abcde11',1,140],[3,'absdf12',1,50],[4,'ddcde14',0,125],[5,'ddcde14',1,140]]
df_example = pd.DataFrame(data,columns = ['key','VisitLink','dx_filter','time'])

I need 3 things to be true:
- VisitLink: matches between the two rows
- dx_filter: is 1 for the first event
- Time: the second event happens within 30 days of the first event Example: Key 1 will generate Key 4 as a matching record, as it meets all qualifications, but Key 4 will not generate Key 5 because its dx_filter = 0.

I ran a trial where I predicted my method would take 120+ hours to complete and am wondering if there is a way to shorten this to <10 hours or if that is not possible.

def add_readmit_id(df):
    df['readmit_id'] = np.nan
    def set_id(row):
        if row['dx_filter'] ==0:
            return np.nan
        else:
            relevant_df = df.loc[df['VisitLink']==row['VisitLink']]
            timeframe_df = relevant_df.loc[(relevant_df['time']>row['time'])&(relevant_df['time']<=row['time']+30)]
            next_timeframe = timeframe_df['time'].min()
            id_row = timeframe_df.loc[timeframe_df['time']==next_timeframe]
            if not id_row.empty:
                return id_row.iloc[0]['key']
            else:
                return np.nan
    df['readmit_id'] = df.apply(set_id,axis=1)
    return df

df_example = add_readmit_id(df_example)

See above for the code I used to run it #minimum reproducible.

pandas python
1 Answers

Here's my approach with groupby:

groups = df.groupby('VisitLink')
s = groups['time'].diff(-1).le(30) & df['dx_filter']
df['shifted'] = np.where(s, groups['key'].shift(-1), np.nan)

Output:

   key VisitLink  dx_filter  time  shifted
0    1   ddcde14          1   100      4.0
1    2   abcde11          1   140      NaN
2    3   absdf12          1    50      NaN
3    4   ddcde14          0   125      NaN
4    5   ddcde14          1   140      NaN
user_42384080
Related