PySpark中实现出租车上下车点与纽约行政区多边形匹配的问题

解决方案 #

1. 预处理行政区数据并广播 #

先将行政区的几何对象转换为shapely.Polygon，并广播到所有Task，避免重复加载数据：

from shapely.geometry import Polygon, Point
from pyspark.sql import functions as F
from pyspark.sql.types import StringType

# 读取并清洗行政区数据
boroughs_df = spark.read.json("nyc-boroughs.geojson") \
    .drop("_corrupt_record") \
    .dropna()

# 转换为(行政区名称, Polygon对象)的列表
def parse_borough_geo(row):
    coords = row.geometry.coordinates
    # 处理MultiPolygon类型（纽约行政区可能是这种结构）
    if row.geometry.type == "MultiPolygon":
        # 取第一个子Polygon
        polygon = Polygon(coords[0][0])
    else:
        polygon = Polygon(coords[0])
    return (row.properties.borough, polygon)

boroughs_geo_list = boroughs_df.rdd.map(parse_borough_geo).collect()
# 广播数据
boroughs_broadcast = spark.sparkContext.broadcast(boroughs_geo_list)

2. 编写正确的行政区匹配UDF #

修正之前的UDF，直接使用广播的Polygon对象进行匹配：

def get_borough(lon, lat):
    # 处理空值
    if not lon or not lat:
        return None
    point = Point(lon, lat)
    # 遍历广播的行政区数据
    for borough_name, polygon in boroughs_broadcast.value:
        if point.within(polygon):
            return borough_name
    return None

# 注册UDF
get_borough_udf = F.udf(get_borough, StringType())

3. 给出租车行程数据添加行政区列 #

调用UDF生成上车和下车行政区列：

taxi_df = taxi_df \
    .withColumn("pickup_borough", get_borough_udf(F.col("pickup_longitude"), F.col("pickup_latitude"))) \
    .withColumn("dropoff_borough", get_borough_udf(F.col("dropoff_longitude"), F.col("dropoff_latitude")))

# 验证结果
taxi_df.select("pickup_longitude", "pickup_latitude", "pickup_borough", 
               "dropoff_longitude", "dropoff_latitude", "dropoff_borough").show(5)

关键说明 #

• 处理MultiPolygon：纽约部分行政区的GeoJSON是MultiPolygon类型，需要提取第一个子Polygon进行匹配
• 广播优化：广播行政区数据后，每个Task只需加载一次，大幅提升大数据集下的处理效率
• 空值防护：提前判断经纬度是否为空，避免创建Point对象时出错
• 简化转换流程：直接从原始GeoJSON坐标生成Polygon，避免中间WKT转换步骤，减少出错可能

内容的提问来源于stack exchange，提问作者John Stuart