PySpark与GeoDataFrame实现点面匹配的高效方案咨询

一、给多边形加空间索引，缩小候选范围 #

先给GeoPandas的多边形GeoDataFrame创建R树空间索引，这能直接把“逐个判断所有多边形”改成“先筛选可能包含点的少量候选多边形”，计算量能砍一大半：

import geopandas as gpd
from shapely.geometry import Point

# 读取多边形数据
poly_gdf = gpd.read_feather("polygons.feather")
# 创建空间索引
spatial_index = poly_gdf.sindex

后续匹配时，先通过空间索引的intersection方法拿到候选多边形，再做精确的包含判断，比全量遍历快很多。

二、弃用Python UDF，用Spark原生空间操作 #

Python UDF的序列化/反序列化开销极大，直接用Spark原生的空间处理框架才是正道：

1. 先把GPS数据转成Spark的Point类型：

from pyspark.sql.functions import udf, col
from pyspark.sql.geometry import PointType, Point

def create_point(lon, lat):
    return Point(lon, lat)

point_udf = udf(create_point, PointType())
gps_df = gps_df.withColumn("point", point_udf(col("longitude"), col("latitude")))

2. 把多边形GeoDataFrame转成Spark DataFrame，用空间JOIN+分组取最大值：

# 多边形转Spark DataFrame
poly_spark_df = spark.createDataFrame(poly_gdf)
# 注册临时表方便写SQL
poly_spark_df.createOrReplaceTempView("polygons")
gps_df.createOrReplaceTempView("gps_points")

# 用ST_Contains做空间匹配，左连接后分组取最大MAX_SPD，没匹配到就用-1填充
result_df = spark.sql("""
    SELECT 
        g.*,
        COALESCE(MAX(p.MAX_SPD), -1) AS max_spd
    FROM gps_points g
    LEFT JOIN polygons p
    ON ST_Contains(p.geometry, g.point)
    GROUP BY g.id, g.longitude, g.latitude, g.point  -- 替换成GPS表的所有字段
""")

Spark的空间操作基于JVM，不用在Python和JVM之间来回折腾，数据量越大优势越明显。

三、给数据分区，对齐空间范围 #

不管是GPS还是多边形数据，按空间范围分区能避免全量笛卡尔积：

• 多边形数据：按经纬度网格分区，让每个分区只存特定区域的多边形
• GPS数据：用repartitionByRange按经度、纬度分区，和多边形分区对齐：

# 按经纬度分20个分区，数量可根据集群调整
gps_df = gps_df.repartitionByRange(20, col("longitude"), col("latitude"))

这样每个分区的GPS点只需要和对应分区的多边形匹配，减少跨区域计算。

四、预处理多边形，减少无效匹配 #

如果有大量重叠或包含关系的多边形，先做预处理：保留每个区域里MAX_SPD最大的多边形，直接砍掉后续不需要的匹配：

def keep_max_spd_polygons(gdf):
    # 按MAX_SPD降序排序，保留每个空间区域的最大MAX_SPD多边形
    # 这里用geometry去重是简化逻辑，实际可结合空间索引找重叠多边形后再筛选
    processed_gdf = gdf.sort_values("MAX_SPD", ascending=False).drop_duplicates(subset=["geometry"])
    return processed_gdf

poly_gdf = keep_max_spd_polygons(poly_gdf)

预处理后多边形数量减少，后续匹配的计算量自然下降。

五、一定要用UDF就用批量处理 #

如果实在绕不开Python UDF，别单条点处理，改成批量UDF，减少Python和JVM的交互次数：

from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import ArrayType, IntegerType
import pandas as pd

@pandas_udf(ArrayType(IntegerType()))
def batch_match(lons: pd.Series, lats: pd.Series) -> pd.Series:
    results = []
    for lon, lat in zip(lons, lats):
        point = Point(lon, lat)
        # 空间索引找候选
        possible_matches_idx = list(spatial_index.intersection(point.bounds))
        possible_matches = poly_gdf.iloc[possible_matches_idx]
        # 精确判断包含
        matches = possible_matches[possible_matches.contains(point)]
        results.append(matches["MAX_SPD"].max() if not matches.empty else -1)
    return pd.Series(results)

gps_df = gps_df.withColumn("max_spd", batch_match(col("longitude"), col("latitude")))

批量UDF的性能比单条UDF至少提升几倍，本质是减少了序列化的次数。

内容的提问来源于stack exchange，提问作者Gi Yeon Shin