pyspark udf 帶多個 arguments

Question

我正在使用 python function 來計算給定經度和緯度的兩點之間的距離。

def haversine(lon1, lat1, lon2, lat2):

    lon1, lat1, lon2, lat2 = map(np.radians, [lon1, lat1, lon2, lat2])

    newlon = lon2 - lon1
    newlat = lat2 - lat1

    haver_formula = np.sin(newlat/2.0)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(newlon/2.0)**2

    dist = 2 * np.arcsin(np.sqrt(haver_formula))
    miles = 3958 * dist 
    return miles

我的 dataframe 有 4 列 - lat、long、merch_lat、merch_long。

當我創建這樣的 UDF 時，它會引發錯誤。 我不知道我哪里錯了。

udf_haversine = udf(haversine, FloatType())
data = data.withColumn("distance", udf_haversine("long", "lat", "merch_long","merch_lat"))

錯誤是：

An error occurred while calling o1499.showString.
: org.apache.spark.SparkException: Job aborted due to stage failure:

如何創建一個需要多列並返回單個值的 udf？

Answer 1

您可能在numpy.dtype和序列化方面遇到問題。

由於miles的類型是numpy.float64嘗試返回float(miles) 。

有效的完整示例：

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.types import DoubleType
import numpy as np


def haversine(lon1, lat1, lon2, lat2):
    lon1, lat1, lon2, lat2 = map(np.radians, [lon1, lat1, lon2, lat2])
    newlon = lon2 - lon1
    newlat = lat2 - lat1
    haver_formula = (
        np.sin(newlat / 2.0) ** 2
        + np.cos(lat1) * np.cos(lat2) * np.sin(newlon / 2.0) ** 2
    )
    dist = 2 * np.arcsin(np.sqrt(haver_formula))
    miles = 3958 * dist
    return float(miles)


spark = SparkSession.builder.getOrCreate()
data = [
    {
        "long": 18.427238,
        "lat": 19.510083,
        "merch_long": 93.710735,
        "merch_lat": 52.182011,
    }
]
df = spark.createDataFrame(data)
udf_haversine = F.udf(haversine, DoubleType())
df = df.withColumn("distance", udf_haversine("long", "lat", "merch_long", "merch_lat"))