Pyspark dataframe row-wise null columns list

user_15900609 0
Viewed 365

I have a spark dataframe and I want to create a new column that contains the columns name having null in each row. For example

Original dataframe is:

|col_1|col_2|col_3|

+-----+-----+-----+

|62.45|null |62.49|

|56.45|null |null |

|null |null |null |

| 56.4|57.32|48.39|

+-----+-----+-----+------------------+

The final dataframe with the result column is:

|col_1|col_2|col_3|            Result|

+-----+-----+-----+------------------+

|62.45|null |62.49|             col_2|

|56.45|null |null |      col_2, col_3|

|null |null |null |col_1, col2, col_3|

| 56.4|57.32|48.39|                  |

+-----+-----+-----+------------------+

I know to get the number of null columns but looking for row-wise column names that can be different in each row. Any guidance will be appreciated.

pyspark python
2 Answers
  (df.withColumn('Result', array(*[when(col(c).isNull(), lit(c)) for c in df.columns]))#Lists Columns with null

.withColumn('Result', expr("filter(Result, x -> x is not null)"))#Exludes nulls from list
 
 .show())
user_89869750

The easiest way is collecting each null column with column function isNull then add them into an array with SQL function array

(df
    .withColumn('Result', F.array(
        F.when(F.col('col_1').isNull(), 'col_1'),
        F.when(F.col('col_2').isNull(), 'col_2'),
        F.when(F.col('col_3').isNull(), 'col_3'),
    ))
    .show()
)

# +-----+-----+-----+---------------------+
# |col_1|col_2|col_3|Result               |
# +-----+-----+-----+---------------------+
# |62.45|null |62.49|[null, col_2, null]  |
# |56.45|null |null |[null, col_2, col_3] |
# |null |null |null |[col_1, col_2, col_3]|
# |56.4 |57.32|48.39|[null, null, null]   |
# +-----+-----+-----+---------------------+

You're obviously want to get rid of the nulls so apply filter would do the job (Note: filter only available since 3.1.0)

(df
    .withColumn('Result', F.array(
        F.when(F.col('col_1').isNull(), 'col_1'),
        F.when(F.col('col_2').isNull(), 'col_2'),
        F.when(F.col('col_3').isNull(), 'col_3'),
    ))
    .withColumn('Result', F.filter(F.col('Result'), lambda c: c.isNotNull()))
 
    .show(10, False)
)

# +-----+-----+-----+---------------------+
# |col_1|col_2|col_3|Result               |
# +-----+-----+-----+---------------------+
# |62.45|null |62.49|[col_2]              |
# |56.45|null |null |[col_2, col_3]       |
# |null |null |null |[col_1, col_2, col_3]|
# |56.4 |57.32|48.39|[]                   |
# +-----+-----+-----+---------------------+

Finally, you also can improve code's flexibility by not hardcoding the columns but using a for loop instead, and still got the same result:

(df
    .withColumn('Result', F.array(
        *[F.when(F.col(c).isNull(), c) for c in df.columns]
    ))
    .withColumn('Result', F.filter(F.col('Result'), lambda c: c.isNotNull()))
 
    .show(10, False)
)
user_34415100
Related