如何使用PySpark从十亿级点数据中提取多边形内的点？

我来给你一步步拆解这个问题——10亿行的数据量确实得靠分布式计算才能高效搞定，PySpark正好是干这个的好手。下面是具体的实现步骤、优化技巧和避坑指南：

1. 先搞定依赖环境 #

首先确保你的Spark集群所有节点都安装了shapely库（用来处理空间几何判断），可以在每个节点执行pip install shapely，或者把依赖打包成虚拟环境提交Spark任务。

2. 加载并处理多边形数据 #

你的多边形CSV是2000个经纬度点，先把它转换成Shapely的Polygon对象，还要注意坐标顺序和闭合性：

from pyspark.sql import SparkSession
from shapely.geometry import Polygon, Point
from pyspark.sql.functions import udf, col
from pyspark.sql.types import BooleanType

# 初始化SparkSession，根据你的集群配置调整参数
spark = SparkSession.builder.appName("PointInPolygon").getOrCreate()

# 读取多边形CSV，假设表头是"latitude", "longitude"
polygon_df = spark.read.csv("path/to/polygon_points.csv", header=True, inferSchema=True)

# 提取经纬度坐标，注意Shapely要求坐标是「经度, 纬度」的顺序！
polygon_coords = polygon_df.select("longitude", "latitude").collect()
# 确保多边形是闭合的：如果首尾点不同，手动补上第一个点
if polygon_coords[0] != polygon_coords[-1]:
    polygon_coords.append(polygon_coords[0])

# 创建Polygon对象，然后广播到所有节点——避免每个Task重复加载，节省内存和网络
polygon = Polygon(polygon_coords)
broadcast_polygon = spark.sparkContext.broadcast(polygon)

3. 编写点-in-多边形判断的UDF #

用UDF封装Shapely的空间判断逻辑，核心是利用广播变量里的多边形：

def is_point_in_polygon(lat, lon):
    # 创建Shapely Point对象，顺序是「经度, 纬度」
    point = Point(lon, lat)
    # 从广播变量中获取多边形，判断点是否在内部
    return point.within(broadcast_polygon.value)

# 注册UDF，指定返回类型为布尔值
point_in_polygon_udf = udf(is_point_in_polygon, BooleanType())

4. 高效加载10亿行的点数据集 #

读取大文件时，分区数的设置直接影响计算效率，建议按「每个分区100-200万行」来调整（对应CSV大概100-200MB）：

# 读取点数据集，假设表头是"ID", "latitude", "longitude"
points_df = spark.read.csv(
    "path/to/large_points.csv",
    header=True,
    inferSchema=True,
    # 10亿行的话，设置5000左右的分区（根据你的集群核心数调整）
    numPartitions=5000
)

# 先做一步预过滤：只保留在多边形外接矩形内的点，能砍掉大量无效计算！
min_lon, min_lat, max_lon, max_lat = broadcast_polygon.value.bounds
filtered_points = points_df.filter(
    (col("longitude") >= min_lon) & 
    (col("longitude") <= max_lon) & 
    (col("latitude") >= min_lat) & 
    (col("latitude") <= max_lat)
)

5. 过滤结果并保存 #

最后应用UDF过滤，再把结果存成高效的格式：

# 过滤出多边形内的点
result_df = filtered_points.filter(
    point_in_polygon_udf(col("latitude"), col("longitude"))
)

# 优先用Parquet格式保存（比CSV快N倍，还省空间）
result_df.write.parquet("path/to/result_points.parquet", mode="overwrite")

# 如果一定要CSV，记得设置合并选项避免生成几千个小文件
# result_df.coalesce(10).write.csv("path/to/result_points.csv", header=True, mode="overwrite")

关键优化技巧 #

• 广播变量必用：多边形数据很小，广播到所有节点能避免每个Task重复加载，大幅降低内存占用和网络开销。
• 先过滤外接矩形：先排除掉明显不在多边形范围内的点，能减少后续UDF计算的量级，提升整体效率30%以上。
• 选对存储格式：读取和写入用Parquet/ORC比CSV高效太多，大数据场景下一定要用列式存储格式。
• 分区数要合理：分区太少会导致单个Task压力过大，分区太多会增加调度开销，建议按「每个分区100-200MB」来设置。

避坑指南 #

• 别搞反坐标顺序：Shapely的Point和Polygon都是「经度在前，纬度在后」，搞反的话判断结果完全错误！
• 确保多边形闭合：必须让首尾点相同，否则Shapely无法识别为闭合多边形，判断逻辑会失效。
• 数据类型要正确：读取时一定要让经纬度字段是float类型，用inferSchema=True或者手动指定Schema，避免字符串类型导致计算错误。

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