如何使用withColumn将额外的参数传递给UDF
df = spark.createDataFrame([
["aaa","1"],
["bbb","2"],
["ccc","5"]
]).toDF("text","id")
def title(x,y):
if y:
x = x.title()
return x
title_udf = udf(lambda x: title(x,y), StringType())
spark.udf.register('title_udf', title_udf)
df = df.withColumn('text_title',title_udf('text',True)
当我尝试这个时，我得到一个错误： Invalid argument, not a string or column....
udf只能识别行元素。因此，要传递固定参数，您必须使用lit（）函数。另外，您的udf定义也必须更正。试试这个：
import pyspark.sql.functions as F
from pyspark.sql.types import *
df = spark.createDataFrame([
["aaa","1"],
["bbb","2"],
["ccc","5"]
]).toDF("text","id")
def title(x,y):
if y:
x = x.title()
return x
title_udf = F.udf(title, StringType())
df = df.withColumn('text_title',title_udf('text',F.lit(True)))
df.show()
+----+---+----------+
|text| id|text_title|
+----+---+----------+
| aaa| 1| Aaa|
| bbb| 2| Bbb|
| ccc| 5| Ccc|
+----+---+----------+
如注释中@powers所示，如果此输出是您的最终目的，则可以使用initcap（）函数在没有udf的情况下执行此操作
df = df.withColumn("text_title",F.when(F.lit(True),F.initcap(F.col('text'))).otherwise(F.col('text')))
您还可以使用其他列作为条件，例如“ id”列
df = df.withColumn("text_title",F.when(F.col('id')>2,F.initcap(F.col('text'))).otherwise(F.col('text')))
本文收集自互联网，转载请注明来源。
如有侵权，请联系 [email protected] 删除。
我来说两句