使用spark_apply计算经度/经度之间的距离

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import DoubleType
from math import sin, cos, sqrt, atan2, radians

spark = SparkSession.builder.appName("Distance Calculation").getOrCreate()

data = [(1, 40.7128, -74.0060), (2, 34.0522, -118.2437), (3, 51.5074, -0.1278)]
df = spark.createDataFrame(data, ["id", "latitude", "longitude"])

def calculate_distance(lat1, lon1, lat2, lon2):
    R = 6371.0  # 地球平均半径，单位为公里

    lat1_rad = radians(lat1)
    lon1_rad = radians(lon1)
    lat2_rad = radians(lat2)
    lon2_rad = radians(lon2)

    dlon = lon2_rad - lon1_rad
    dlat = lat2_rad - lat1_rad

    a = sin(dlat / 2)**2 + cos(lat1_rad) * cos(lat2_rad) * sin(dlon / 2)**2
    c = 2 * atan2(sqrt(a), sqrt(1 - a))

    distance = R * c
    return distance

distance_udf = udf(calculate_distance, DoubleType())

df = df.withColumn("distance", distance_udf(df.latitude, df.longitude, lit(0), lit(0)))