Improve Column Validation For Multiple Dataframes Pyspark-CodePudding

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