python col函数_Pyspark中的UDF和python函数

listA = [(1,'AAA','USA'),(2,'XXX','CHN'),(3,'KKK','USA'),(4,'PPP','USA'),(5,'EEE','USA'),(5,'HHH','THA')]

df = spark.createDataFrame(listA, ['id', 'name','country'])

我创造了一本字典:

thedict={"USA":"WASHINGTON","CHN":"BEIJING","DEFAULT":"KEY NOT FOUND"}

然后我创建了一个UDF,从dictionary中获取匹配的键值。

def my_func(letter):

if(thedict.get(letter) !=None):

return thedict.get(letter)

else:

return thedict.get("DEFAULT")

尝试调用函数时出现以下错误:

df.withColumn('CAPITAL',my_func(df.country))

File "", line 1, in

File "/usr/hdp/current/spark2-client/python/pyspark/sql/dataframe.py", line 1848, in withColumn

assert isinstance(col, Column), "col should be