I have a function that validates if the dataframe passed has a few columns and if it does not it creates them and fills the values with 0.0
.
This takes a bit of time to run and has several if statements. Is there any way this function can be improved? In truth, I run this for multiple dataframes but at the moment I need to run this function for each individually, is there a way to run for all of them at once?
This is the function I have:
def validate_columns(df):
if 'A' not in df.columns:
df = df.withColumn('A', lit(0.0))
if 'B' not in df.columns:
df = df.withColumn('B', lit(0.0))
if 'C' not in df.columns:
df = df.withColumn('C', lit(0.0))
if 'D' not in df.columns:
df = df.withColumn('D', lit(0.0))
df_to_return = df.select('A', 'B', 'C', 'D')
return df_to_return
CodePudding user response:
For a single dataframe, you can use a for-loop just to improve code understandability. You need to pass a list of columns to the function.
def validate_columns(df, cols_of_interest):
for c in cols_of_interest:
if c not in df.columns:
df = df.withColumn(c, lit(0.0))
result = df.select(*cols_of_interest)
return result