基于PySpark实现BallTree完成两个DataFrame的最近位置匹配问题

• UDF运行在Executor进程中，禁止在UDF内部调用Spark DataFrame的collect()、filter()等算子，Executor侧没有可用的SparkSession上下文，会直接触发运行错误
• 将Spark DataFrame作为默认参数传入UDF会触发不可预期的序列化问题，无法在Executor侧正常解析
• 单条处理的普通UDF性能极低，无法支撑600万行规模的数据处理

如果你的门店数据量小于10万条（绝大多数线下门店场景都符合这个量级），优先使用以下方案：提前把门店数据拉到Driver端转成本地Python对象，通过闭包传入批量处理的Pandas UDF，全程不需要用到Spark广播能力，性能足够处理600万行数据。

实现代码 #

步骤1：预处理门店数据转成本地对象 #

from pyspark.sql import functions as f
from sklearn.neighbors import BallTree
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import IntegerType
import numpy as np
import pandas as pd

# 先把门店数据全部拉到Driver端转成本地数组，只要门店量<10万完全不占内存
store_arr = df_store.select('lat_store_rad', 'lon_store_rad', 'Store_ID').collect()
# 拆分坐标数组和ID列表
store_coords_rad = np.array([(row['lat_store_rad'], row['lon_store_rad']) for row in store_arr])
store_ids = np.array([row['Store_ID'] for row in store_arr])
EARTH_RADIUS = 6371
MAX_DISTANCE = 3 # 匹配阈值单位：公里

步骤2：定义批量处理Pandas UDF #

@pandas_udf(IntegerType())
def get_nearest_store(lat_rad: pd.Series, lon_rad: pd.Series) -> pd.Series:
    # 把输入的两个series转成(n,2)的坐标数组
    src_coords = np.column_stack([lat_rad.values, lon_rad.values])
    # 构建BallTree批量查询
    tree = BallTree(store_coords_rad, leaf_size=15, metric='haversine')
    distances, indices = tree.query(src_coords, k=1)
    # 转换为公里距离
    distances_km = distances.flatten() * EARTH_RADIUS
    # 匹配ID，超过3公里的替换为默认值
    result_ids = store_ids[indices.flatten()]
    result_ids[distances_km > MAX_DISTANCE] = 999999999
    return pd.Series(result_ids)

步骤3：调用UDF得到结果 #

# 预处理个人位置的弧度坐标
individual_positions = individual_positions.withColumn('lat_rad', f.radians(f.col('lat')))
individual_positions = individual_positions.withColumn('lon_rad', f.radians(f.col('lon')))
# 应用UDF得到匹配结果
individual_positions = individual_positions.withColumn('Store_ID', get_nearest_store('lat_rad', 'lon_rad'))
# 可选：删除中间计算的弧度字段
individual_positions = individual_positions.drop('lat_rad', 'lon_rad')

如果你的门店数据超过10万条，Driver端无法全量存储，可以用网格分箱方案：

• 对所有个人位置和门店位置计算Geohash前缀，选择精度对应3公里左右的前缀长度（一般是Geohash长度5，对应边长约5公里）
• 关联时同时匹配当前网格和周围8个相邻网格的门店，避免边界位置漏匹配
• 对同一个个人位置关联到的所有门店，计算真实哈弗辛距离，过滤出3公里内最近的门店即可

内容的提问来源于stack exchange，提问作者Alex Germain