pyspark：将多个数据框字段传递给 udf

Posted 2023-04-13

【中文标题】pyspark：将多个数据框字段传递给 udf

【英文标题】：pyspark: passing multiple dataframe fields to udf

【发布时间】：2017-12-15 02:08:59

【问题描述】：

我是 spark 和 python 的新手。任何帮助表示赞赏。

我有一个 UDF，并创建了一个带有美国 zipcd、纬度和经度的 spark 数据框

UDF：

import math
def distance(origin, destination):
lat1, lon1 = origin
lat2, lon2 = destination
radius = 6371 # km
dlat = math.radians(lat2-lat1)
dlon = math.radians(lon2-lon1)
a = math.sin(dlat/2) * math.sin(dlat/2) + math.cos(math.radians(lat1)) \
    * math.cos(math.radians(lat2)) * math.sin(dlon/2) * math.sin(dlon/2)
c = 2 * math.atan2(math.sqrt(a), math.sqrt(1-a))
d = radius * c
return d

示例 UDF 输出：

distance((101,121),(-121,-212)) 

15447.812243421227

数据框：

zip=spark.read.option("sep", ",").csv('wasb://[e-mail]/main/zip.txt')
zip1=zip.select(zip._c0,zip._c1.cast("Double"),zip._c2.cast("Double"))

zip1 数据示例：

zip1.first()        

行（_c0=u'00601', _c1=18.180555, _c2=-66.749961）

现在我试图将纬度和经度从 df zip1 传递到 udf 距离，但我收到了“需要浮点数”之类的错误。 我相信 udf 没有从 df 字段中获取数据，而是将 df 列作为常量值读取；因此我遇到了错误。

z=zip1.select(distance((zip1._c1,100.23),(zip1._c2,-99.21)))

回溯（最近一次通话最后一次）： 文件“”，第 1 行，在 文件“”，第 5 行，距离 TypeError：需要浮点数

请告诉我将 df 字段传递给 udf 的正确方法。

【参考方案1】：

我不太确定您的数据架构是什么。 但是下面的例子是使用udf 来得到你的例子答案的正确方法。

from pyspark.sql.functions import *
from pyspark.sql.types import *
import math

def distance(origin, destination):
    lat1, lon1 = origin
    lat2, lon2 = destination
    radius = 6371 # km
    dlat = math.radians(lat2-lat1)
    dlon = math.radians(lon2-lon1)
    a = math.sin(dlat/2) * math.sin(dlat/2) + math.cos(math.radians(lat1)) \
    * math.cos(math.radians(lat2)) * math.sin(dlon/2) * math.sin(dlon/2)
    c = 2 * math.atan2(math.sqrt(a), math.sqrt(1-a))
    d = radius * c
    return d

df = spark.createDataFrame([([101, 121], [-121, -212])], ["origin", "destination"])
filter_udf = udf(distance, DoubleType())
df.withColumn("distance", filter_udf(df.origin, df.destination))

+----------+------------+------------------+
|    origin| destination|          distance|
+----------+------------+------------------+
|[101, 121]|[-121, -212]|15447.812243421227|
+----------+------------+------------------+

【讨论】：

@vaira 你对上面的代码有什么问题吗？如果答案是你想要的，请选择它作为答案，谢谢。