Databricks 火花 UDF 不適用於過濾的 dataframe

Question

我在使用 Pyspark 的 Databricks 中遇到了一個問題，如果我在這里遺漏了一些概念性的東西，我試圖理解為什么這個實現不起作用。 我想要做的是在 dataframe 的列上運行 UDF，但只能在非空值上運行。

如果我將 lstrip_udf 調用替換為“Val123”之類的固定值，那么它可以正常工作，但它不適用於 UDF。 如果我在 UDF 中實現 null 檢查，但實現方式略有不同，那么它也可以工作。 但即使有when和IsNotNull它仍然會引發以下錯誤。

有人可以解釋為什么或我在這里缺少什么來完成這項工作嗎？

代碼：

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)

錯誤：

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

謝謝

Answer 1

這可能是 Spark 中的一個錯誤，所以這里是對 UDF 的一個小修改，它解決了這個問題：

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

或者您可以使用 Spark SQL 來執行此操作，這比使用 UDF 更有效：

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