Add columns to df1 that are in df2 but not in df1 and vice versa in pyspark

Viewed 61

I have 2 spark dfs that are of different column size and partly different columns.

sdf1 = ['colA','colB','colC','colD','coldE','colF','colG','colH']
sdf2 = ['colA','colB','colI','colJ']

I would like to add columns from sdf1 that are missing in sdf2, and vice versa what I expect to get after that are:

sdf1 = ['colA','colB','colC','colD','coldE','colF','colG','colH','colI','colJ']
sdf2 = ['colA','colB','colI','colC','colD','coldE','colF','colG','colH','colJ']

when I followed the suggested approached from https://www.geeksforgeeks.org/merge-two-dataframes-with-different-amounts-of-columns-in-pyspark/ or from https://sparkbyexamples.com/spark/spark-merge-two-dataframes-with-different-columns/ ,

#Add missing columns to df1
from pyspark.sql.functions import lit
for column in [column for column in sdf2.columns if column not in sdf1.columns]:
    df1 = sdf1.withColumn(column, lit(None))

#Add missing columns to df2
for column in [column for column in sdf1.columns if column not in sdf2.columns]:
    df2 = sdf2.withColumn(column, lit(None))

sdf1 resulted to what I expect but not with sdf2, where I get:

df2 = ['colA','colB','colI','colJ','colH']

only the last column name was copied. Am I missing something that explains this? My goal is to merge both and I think I cannot do it without the same columns in pyspark

2 Answers

I'm using following function for such purpose (source):

def add_missing_columns(df1: DataFrame, df2: DataFrame) -> DataFrame:
    """
    Adds to first dataframe columns from the second dataframe that don't exist in first one.
    Columns get null values casted to respective data types
    :param df1: first dataframe
    :param df2: second dataframe
    :return: new dataframe with added columns
    """
    additional_cols = [F.lit(None).cast(field.dataType).alias(field.name)
                       for field in df2.schema.fields if field.name not in df1.columns]
    return df1.select("*", *additional_cols)

So if you want to add columns to both of them, then you just do two calls to that function, reverting the order of arguments:

df1_new = add_missing_columns(df1, df2)
df2_new = add_missing_columns(df2, df1)

Try the approach presented in this link: pyspark adding columns to dataframe that are already not present from a list where you list the column names and save it as a list, then use it as your reference list for the columns of your sdf2.

all_cols = ['colA','colB','colC','colD','coldE','colF','colG','colH','colI','colJ']

for cols in all_cols:
    if cols not in sdf2.columns:
         df = sdf2.withColumns(cols,lit(None))
Related