如何在PyArrow中使用自定义函数筛选多边形内的行数据

核心思路 #

先通过PyArrow的数值筛选快速缩小数据范围（利用多边形的外接矩形），再用矢量化空间判断处理剩余数据，避免全量转Pandas和逐行apply的低效操作。

具体实现步骤 #

1. 导入依赖库

import pyarrow as pa
import pyarrow.parquet as pq
from shapely.geometry import shape
from shapely.vectorized import contains

2. 读取Parquet并做外接矩形预筛选

先筛选出经纬度落在多边形外接矩形内的行，大幅减少后续空间计算的数据量：

# 读取Parquet数据集
dataset = pq.ParquetDataset("你的Parquet文件/路径")

# 定义目标多边形（假设输入是GeoJSON格式的多边形）
polygon = shape(你的多边形GeoJSON)
min_lon, min_lat, max_lon, max_lat = polygon.bounds

# 用PyArrow Compute做预筛选
filtered_table = dataset.read().filter(
    pa.compute.and_(
        pa.compute.greater_equal(pa.field("lon"), min_lon),
        pa.compute.less_equal(pa.field("lon"), max_lon),
        pa.compute.greater_equal(pa.field("lat"), min_lat),
        pa.compute.less_equal(pa.field("lat"), max_lat)
    )
)

3. 矢量化空间判断筛选最终数据

用Shapely的矢量化contains函数批量判断点是否在多边形内，再转回PyArrow布尔数组完成筛选：

# 提取经纬度为Numpy数组（Shapely矢量化函数支持Numpy输入）
lon_array = filtered_table["lon"].to_numpy()
lat_array = filtered_table["lat"].to_numpy()

# 批量判断点是否在多边形内
in_polygon_mask = contains(polygon, lon_array, lat_array)

# 将掩码转为PyArrow布尔数组，筛选最终结果
final_table = filtered_table.filter(pa.array(in_polygon_mask, type=pa.bool_()))

4. 超大数据集优化：分批处理

如果数据集过大无法全量加载，用PyArrow Scanner分批处理，内存更友好：

# 创建Scanner，只扫描外接矩形内的数据
scanner = dataset.scanner(
    filter=pa.compute.and_(
        pa.compute.greater_equal(pa.field("lon"), min_lon),
        pa.compute.less_equal(pa.field("lon"), max_lon),
        pa.compute.greater_equal(pa.field("lat"), min_lat),
        pa.compute.less_equal(pa.field("lat"), max_lat)
    )
)

# 分批处理每个数据块
filtered_batches = []
for batch in scanner.to_batches():
    lon_np = batch["lon"].to_numpy()
    lat_np = batch["lat"].to_numpy()
    mask = contains(polygon, lon_np, lat_np)
    filtered_batches.append(batch.filter(pa.array(mask, type=pa.bool_())))

# 合并所有批次为最终Table
final_table = pa.Table.from_batches(filtered_batches)

为什么原方法慢？ #

1. 全量将PyArrow Table转成Pandas DataFrame会带来额外的内存和转换开销；
2. apply逐行处理是循环操作，没有利用矢量化计算的优势，数据量越大效率越低。

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