Keep rows based in more than 2 columns

Question

I have a database similar to the one I created below (with 982977 rows × 10 columns), and I wanted to keep the rows where the exams of the same patient (ID) that are different from "COVID" have been performed in a specific period based on the date of the "COVID" exam.

I created 2 columns, one with dates 7 days before and one with 30 days after the original exam date. Ex: If the patient had an iron exam between 7 days before and 30 days after the date of their COVID exam, then I would keep that patient, otherwise, I would remove.

I did a for loop, but since the database is big, it took almost 6h to complete and when it finished, I lost the connection to the server, and I couldn't continue to manipulate the data

Is there a simpler and/or faster way to do this?

ID = ['1','1','1','2','2']
Exam = ['COVID', 'Iron', 'Ferritin', 'COVID', 'Iron']
Date = [2021-02-22,2021-02-20,2021-06-22,2021-05-22,2021-05-29]
Date7 = [2021-02-15,2021-02-13,2021-06-15,2021-05-15,2021-05-22]
Date30 = [2021-03-24,2021-03-22,2021-07-24,2021-05-22,2021-06-29]

teste = list(zip(ID, Exam, Date, Date7, Date30))

teste2 = pd.DataFrame(teste, columns=['ID','Exam','Date', 'Date7', 'Date30'])

All the dates columns are in datetime already

pacients = []
for pacient in teste2.ID.unique():
    a = teste2[teste2.ID==pacient]
    b = a[a.Exam!="COVID"]
    c = a[a.Exam=="COVID"]
    for exam_covid in b.Data:
        for covid_7 in c.Data7:
            for covid_30 in c.Data30:
                if covid_7 < exam_covid < covid_30:
                    pacients.append(pacient)
pacients = set(pacientes)
pacients = list(pacientes)

user_14311263 · Answer

With the following sample dataframe named df

ID = ['1','1','1','2','2']
Exam = ['COVID', 'Iron', 'Ferritin', 'COVID', 'Iron']
Date = ['2021-02-22','2021-02-20','2021-06-22','2021-05-22','2021-06-29']
df = pd.DataFrame({'ID': ID, 'Exam': Exam, 'Date': pd.to_datetime(Date)})

you could try the following:

Step: Create a dataframe df_cov that covers all the time intervals around the Covid exams:

df_cov = df[['ID', 'Date']][df.Exam.eq('COVID')]
df_cov = df_cov.assign(
    Before=df_cov.Date - pd.Timedelta(days=7),
    After=df_cov.Date + pd.Timedelta(days=30)
).drop(columns='Date')

Step: merge the non-Covid-exams in df with df_cov on the column ID, then select the exams that are within the intervals (here with query), and then extract the remaining unique IDs:

patients = (
    df[df.Exam.ne('COVID')].merge(df_cov, on='ID', how='left')
    .query('(Before < Date) & (Date < After)')
    .ID.unique()
)

Result for the sample (I've changed the last exam date such that it won't fall into the required time interval):

array(['1'], dtype=object)

Keep rows based in more than 2 columns

1 Answers