Databricks 火花 UDF 不适用于过滤的 dataframe

Question

I came across an issue in Databricks with Pyspark and I'm trying to understand why this implementation is not working, if I'm missing something conceptual here. What I'm trying to do is to run a UDF on a column in a dataframe but only on non-null values.

If I replace the lstrip_udf call with a fix value like "Val123" then it works fine, however it doesn't work for UDF. If I implement a null check inside the UDF with a bit different implementation then it works too. But even though there is the when and the IsNotNull it still throws the below error.

Can someone explain why or what am I missing here to make this work?

Code:

from pyspark.sql.types import StructType, StructField, IntegerType, StringType, BooleanType, TimestampType
inputschema = StructType([StructField("testcol", StringType(), True),
                          StructField("testcol2", StringType(), True)
                         ]
                        )
inputfile = spark.createDataFrame([("012121212","Ref #1"),
                                   ("0034343434","Ref #2"),
                                   ("0034343434","Ref #3"),
                                   (None,"Ref #4"),
                                   (None,"Ref #5"),
                                   ("00998877","Ref #6")
                                  ],
                                  schema = inputschema
                                 )
#display(inputfile)

from pyspark.sql.functions import col, when, lit
column_name = "testcol"
lstrip_udf = udf(lambda s: s.lstrip().lstrip("0"), StringType())
outputfile = (inputfile.withColumn(column_name,
                                  when(col(column_name).isNotNull(),
                                       lstrip_udf(col(column_name)) #replace this line with "Val123" and it works
                                      )
                                 ))
display(outputfile)

Error:

File "<command-3701821159856508>", line 18, in <lambda>
AttributeError: 'NoneType' object has no attribute 'lstrip'

Thanks

Answer 1

It's probably a bug in Spark, so here is a minor modification to the UDF which solves the problem:

lstrip_udf = udf(lambda s: s.lstrip().lstrip("0") if s is not None else None, StringType())

Or you can use Spark SQL to do this, which is more efficient than using UDF:

outputfile = (
    inputfile.withColumn(column_name,
        F.when(col(column_name).isNotNull(),
            F.expr("ltrim('0', ltrim('',testcol))")
        )
    )
)