Nested withColumn and otherwise Pyspark

Question

I have this code:

dfSpark = dfSpark\
    .withColumn('colA', when(col('colA') == True, 'S').otherwise('N')) \
    .withColumn('colB', when(col('colB') == True, 'S').otherwise('N')) \
    .withColumn('colC', when(col('colC') == True, 'S').otherwise('N')) \
    .withColumn('colD', when(col('colD') == True, 'S').otherwise('N')) \
    .withColumn('colE', when(col('colE') == True, 'S').otherwise('N')) \
    .withColumn('colF', when(col('colF') == True, 'S').otherwise('N')) \
    .withColumn('colG', when(col('colG') == True, 'S').otherwise('N')) \
    .withColumn('colH', when(col('colH') == True, 'S').otherwise('N')) \
    .withColumn('colI', when(col('colI') == True, 'S').otherwise('N')) \
    .withColumn('colJ', when(col('colJ') == True, 'S').otherwise('N'))

Is there any other way that could be more efficient and not redundant? Something like a lambda in python?

Thanks!

score 2 · Accepted Answer · answered Aug 03 '22 at 12:50

Use select with list comprehension:

from pyspark.sql import functions as F

update_cols = ["colA", "colB", "colC", "colD", "colE", "colF", "colG", "colH", "colI", "colJ"]

dfSpark = dfSpark.select(
    *[F.col(c) for c in df.columns if c not in update_cols],
    *[F.when(F.col(c), "S").otherwise("N") for c in update_cols]
)

Azhar Khan · Answer 2 · 2022-08-04T03:32:17.470

You can use UDF (user defined functions):

from pyspark.sql.functions import udf, col
from pyspark.sql.types import StringType

@udf(returnType=StringType())
def my_udf(col):
  if col == True:
    return "S"
  else:
    return "N"

and call it as:

dfSpark = dfSpark \
    .withColumn('colA', my_udf(col('colA')) \
    .withColumn('colB', my_udf(col('colB'))

But remember, UDFs involve repeated serialization, deserialization and data movement between Python interpreter and JVM, and may be very slow for huge datasets. The vectorised form (as you've coded) is more efficient.

There are vectorised UDFs, as well. Please check https://databricks.com/blog/2017/10/30/introducing-vectorized-udfs-for-pyspark.html

Performance: Spark functions vs UDF performance?

Nested withColumn and otherwise Pyspark

2 Answers2