Spark大数据集下OSMNX nearest_edges函数优化方案咨询

你的核心问题在于每个UDF调用都重复构建局部道路图，这带来了巨大的IO和计算开销，同时宽泛的异常处理导致空值原因无法定位。以下是针对性的优化方案：

一、彻底避免重复构建道路图 #

原代码里每个点都调用ox.graph_from_point下载/构建局部图，这是性能灾难。正确的做法是：

• 在Spark Driver端一次性加载整个目标区域的道路图（比如城市、省份范围），而不是每个点都拉取局部数据。
• 对加载后的图做投影处理，然后通过Spark的broadcast机制把图分发到所有工作节点，避免每个任务重复加载。

示例代码：

import osmnx as ox
from pyspark.sql import SparkSession

# Driver端预加载目标区域道路图
# 可用graph_from_bbox或graph_from_place指定范围
G = ox.graph_from_place("北京市", network_type='all', simplify=True, retain_all=True)
Gp = ox.project_graph(G)  # 提前完成投影，避免重复计算

# 广播投影后的图到所有节点
broadcast_Gp = spark.sparkContext.broadcast(Gp)

二、使用批量Vectorized UDF替代单条记录UDF #

OSMNX的nearest_edges本身支持批量坐标输入，配合Spark的Pandas UDF（Vectorized UDF），可以大幅减少函数调用开销，效率比普通UDF高一个数量级。

示例批量处理UDF：

from pyspark.sql.functions import pandas_udf
import pandas as pd
import numpy as np
from shapely.geometry import Point

@pandas_udf("double")
def batch_get_distance_to_road(lat_dd: pd.Series, long_dd: pd.Series) -> pd.Series:
    # 获取广播的投影图
    Gp = broadcast_Gp.value
    crs = Gp.graph['crs']
    
    # 批量投影坐标
    points = [Point(lon, lat) for lat, lon in zip(lat_dd, long_dd)]
    projected_points, _ = ox.projection.project_geometry(points, to_crs=crs)
    
    # 批量提取x、y坐标
    xs = [p.x for p in projected_points]
    ys = [p.y for p in projected_points]
    
    # 批量查询最近边及距离
    _, distances = ox.nearest_edges(Gp, xs, ys, return_dist=True)
    
    # 处理空值（比如点在图范围外的情况）
    distances = np.round(distances, 2)
    return pd.Series(distances).replace({np.nan: None})

三、优化空值问题排查与处理 #

• 提前过滤无效坐标：调用UDF前，先过滤掉经纬度不在目标区域范围内的点，避免无意义计算。
• 细化异常捕获：不要用宽泛的except Exception，可以捕获ValueError（坐标无效）、IndexError（无匹配边）等具体异常，针对性返回标记值，方便后续排查。
• 验证图范围：确保预加载的图覆盖所有数据点的范围，否则会出现大量空值。

四、额外性能优化建议 #

• 按区域分区数据集：如果数据覆盖多个大区域，可按经纬度范围分区，每个分区加载对应子图，进一步减少内存占用。
• 简化道路图：如果不需要所有类型道路，指定network_type='drive'或'walk'，缩小图的规模。
• 缓存投影后的图：Driver端加载图后缓存起来，避免重复加载。

内容的提问来源于stack exchange，提问作者William Fabyan