Applying UDF only on rows where value is not null or not an empty string not working as expected

user_439621 0
Viewed 258

What's the best (fastest) way to apply an UDF only when a value is not null or not an empty string.

I've added a simple example.

df = spark.createDataFrame(
    [["John Jones"], ["Tracey Smith"], [None], ["Amy Sanders"], [""]]
).toDF("Name")


def upperCase(str):
    return str.upper()


upperCaseUDF = udf(lambda z: upperCase(z), StringType())

df.withColumn(
    "Cureated Name",
    F.when(
        ((F.col("Name").isNotNull()) | (F.trim(F.col("name")) != "")),
        upperCaseUDF(F.col("Name")),
    ),
)

AttributeError: 'NoneType' object has no attribute 'upper'.

I don't think the when clause works properly (or at least not as I would expect).
I get an error for the Null value.

I expect the UDF not to be executed on a Null value.
It's not about solving the Null value, but why the when clause doesn't work as I would expect !

pyspark user-defined-functions python-3.x
1 Answers

I would advice you to consider that your UDF should apply to the whole dataframe and adapt the code in consequence:

@F.udf
def upperCase(in_string):
    return in_string.upper() if in_string else in_string


df.withColumn(
    "Created_Name",
    upperCase(F.col("Name")),
).show()

+------------+------------+
|        Name|Created_Name|
+------------+------------+
|  John Jones|  JOHN JONES|
|Tracey Smith|TRACEY SMITH|
|        null|        null|
| Amy Sanders| AMY SANDERS|
|            |            |
+------------+------------+

NB: Your UDF works if you filter out the bad lines:

df.where(F.col("Name").isNotNull()).select(upperCaseUDF(F.col("Name"))).show()
+--------------+                                                                
|<lambda>(Name)|
+--------------+
|    JOHN JONES|
|  TRACEY SMITH|
|   AMY SANDERS|
|              |
+--------------+
user_50137520
Related