PySpark: Create new column based on null values in other columns

Viewed 38

I am working on a PySpark transformation to create a new column based on null values in another columns. Below is the sample input dataframe:

Input DataFrame

This is the expected output dataframe:

Output Dataframe

1 Answers

Hi welcome to stack overflow. It is probably a good idea to read the question best practices. To build a column like that, the easiest way will probably be to build a column with your desired text corrosponding to nulls in each original column. for example

    cols = ["A", "B", "C", "D"]
    new_cols = ["A_nulls", "B_nulls", "C_nulls", "D_nulls", "new_col"]

    df = source_df.withColumn("new_col", F.lit("Null Columns:"))
    df = df.withColumn("A_nulls", F.when(
        F.col("A").isNotNull(), F.lit("A,"))
        .otherwise("")
    )
    df = df.withColumn("B_nulls", F.when(
        F.col("B").isNotNull(), F.lit("B,"))
        .otherwise(""))

    df = df.withColumn("C_nulls", F.when(
        F.col("C").isNotNull(), F.lit("C,"))
        .otherwise(""))

    df = df.withColumn("D_nulls", F.when(
        F.col("D").isNotNull(), F.lit("D,"))
        .otherwise(""))

    df = df.select(*cols, F.concat(*new_cols).alias("NewCol"))

If you want to remove the trailing , you can use F.regexp_replace("new_col", ",$", "") which should trim those.

Related