如何在 pyarrow 中使用自定义函数过滤行

Question

我有一个镶木地板数据集，其中包含作为单独列的纬度和经度值。我想过滤多边形内的那些行，我可以在 pandas 数据框中执行此操作，但无法在 pyarrow 表中执行此操作。

我正在使用 pyarrow 读取镶木地板文件，因为它非常快。

以下是我在熊猫中的做法：

import pyarrow as pa
from shapely.geometry import shape, Point

def point_in_polygon(df, polygon): 
    return df.apply(lambda x: shape(polygon).intersects(Point(x.lon, x.lat)), axis=1)

res: pa.Table = ParquetDataset(....)
res.to_pandas().loc[lambda df: point_in_polygon(df, polygon)]

但是上述方法的问题是它很慢。我知道 pyarrow 和 pyarrow.compute 中的过滤器，但无法弄清楚如何实现这一点。

如果需要更多信息，请告诉我:)
谢谢:)

Answer 1

我不确定

pyarrow

的计算是否有apply方法。 Arrow 项目最初是作为内存分析而非计算的规范开始的，并以不同的语言实现。后来添加了一些对计算的支持。

Arrow 实现主要由其他库使用，而不是最终用户使用。如果您有兴趣，这里有一篇关于 Arrow 的精彩文章：https://www.datawill.io/posts/apache-arrow-2022-reflection/

为此，我建议使用优化的 DataFrame 库，例如 Polars (https://www.pola.rs/)。 Polars 是一个面向最终用户的数据框库，使用 Rust 编写并基于 Apache Arrow 规范。

现在关于你的问题：

使用 apply 调用外部库通常比使用内置的 Polars 函数慢。 Polars 有很多功能，你可以在这里找到：https://pola-rs.github.io/polars/py-polars/html/reference/（我不熟悉地理数据分析，所以我无法推荐要使用的特定功能）

如果您需要使用 apply 和 shapely 库，Polars 允许您在多线程上运行该函数。这可以大大加快计算速度。

资源：

• 将 Pyarrow 转换为 Polars：https://pola-rs.github.io/polars/py-polars/html/reference/api/polars.from_arrow.html#polars-from-arrow

• 在 Polars 数据帧上运行应用程序（您可以查看策略：线程）：https://pola-rs.github.io/polars/py-polars/html/reference/expressions/api/polars.Expr。 apply.html#polars.Expr.apply

Answer 2

它很慢，因为 Python 函数调用在标量点上循环。 Shapely 支持 numpy 数组上的 universal functions。关键是获得一个内嵌循环的点数组。假设您有

lons

和

lats

.

的数组（numpy 或 pyarrow）

points = shapely.from_ragged_array(shapely.GeometryType.POINT, np.array([lons, lats]).T)
shape(polygon).intersects(points)