[Python] Long time for Python Polars Dataframe operations

I have a reasonably large dataset ~500,000 rows. This contains information on patient readmissions and the procedures/diagnoses present for that admission identified through unique patient identifiers. There is also information on days between admission events.

This is a sample dataset (my actual dataset has 40 diagnosis columns and 30 procedure columns):

df = pd.DataFrame({
'Patient_identifier': ['A12', 'A12', 'A12', 'A14', 'A14', 'B12', 'B15', 'B15', 'B15','B15'],
'Day': [3, 7, 12, 5, 8, 21, 2, 4, 7, 121]
'Diagnosis': [code1, code2, code1, code3, code1, code7, code9, code3, code7, code7]
'Procedure': [procedure1, procedure1, procedure4, procedure2, procedure5, procedure2,
procedure5, procedure7, procedure7, procedure4]})


My goal is to calculate multiple complication rates for each patient readmission within a number of days which are specified as a parameter to this function by looking through each readmission. My current function is below, however, it takes a long period of time ~2-3 minutes for me to calculate each complication for this large of a dataset. Are there ways for me to speed the implementation of this function?

My function:

def complication_col_readmission(df: pl.LazyFrame,
icd_10_codes: list[str],
readmission_period: float,
column_name: str) -> pl.LazyFrame:

columns = ['Diagnosis', 'Procedure']


# Sort the dataframe by NRD_VISITLINK and NRD_DAYSTOEVENT
df = df.sort(['Patient_identifier', 'Day'], descending=[False, True])

# Identify the first visit per Patient_identifier
df_first_visits = df.group_by('Patient_identifier').agg(pl.col('Days').min().alias('min_DAYSTOEVENT'))
# Join the first visit information back to the main dataframe
df = df.join(df_first_visits, on='Patient_identifier', how='left')
conditions = []
max_shifts_df = df.group_by('Patient_identifier').agg(pl.count('Patient_identifier').alias('max_shifts'))

# Get the maximum number of shifts
max_shifts_value = int(max_shifts_df.select(pl.col('max_shifts').max()).collect().item())

for shift in range(1, max_shifts_value):
for code in icd_10_codes:
condition = (
(pl.col(columns).shift(-shift).str.starts_with(code)) &
(pl.col('Patient_identifier').shift(-shift) == pl.col('Patient_identifier')) &
((pl.col('Days').shift(-shift) - pl.col('min_DAYSTOEVENT')) < readmission_period)
)
conditions.append(condition)


condition_shifted = pl.any_horizontal(conditions)

# Calculate the readmission complication column
df = df.with_columns(
pl.when(
(
condition_shifted
)
)
.then(1.0)
.otherwise(0.0)
.alias('sub_complication')
)
complication_occured = df.group_by('Patient_identifier').agg(pl.col('sub_complication').max().alias(column_name))
df = df.join(complication_occured, on = 'Patient_identifier')
df = df.drop('min_DAYSTOEVENT', 'sub_complication')
print(df.collect())
return df

Continue reading...