PyPolars：加速应用功能查找共同元素

Question

我试图在引用单元格的列表列中查找公共元素。我可以用一个小数据集来完成它，但我面临两个问题。即使对于 25 行样本数据（每个循环 20.7 s ± 52 ms），速度也非常慢，并且我无法通过

map_batches

找到更快的实现，它可以使用并行化，这与在单线程上工作的

map_elements

不同。

我现在正在使用的版本如下：

import polars as pl
import numpy as np

df = pl.DataFrame({'animal': ['goat','tiger','goat','tiger','lion','goat','tiger','lion'], 'food': ['grass','rabbit','carrots','deer','zebra','water','water','water']})
dl = df.group_by('animal', maintain_order=True).all()
dl

shape: (3, 2)
┌────────┬───────────────────────────────┐
│ animal ┆ food                          │
│ ---    ┆ ---                           │
│ str    ┆ list[str]                     │
╞════════╪═══════════════════════════════╡
│ goat   ┆ ["grass", "carrots", "water"] │
│ tiger  ┆ ["rabbit", "deer", "water"]   │
│ lion   ┆ ["zebra", "water"]            │
└────────┴───────────────────────────────┘

refn = dl['food'][1].to_numpy()

dl = dl.with_columns(
    pl.col('food').map_elements(lambda x: np.intersect1d(refn,x.to_numpy()), return_dtype=pl.List(pl.String))
)

dl

shape: (3, 2)
┌────────┬─────────────────────────────┐
│ animal ┆ food                        │
│ ---    ┆ ---                         │
│ str    ┆ list[str]                   │
╞════════╪═════════════════════════════╡
│ goat   ┆ ["water"]                   │
│ tiger  ┆ ["deer", "rabbit", "water"] │
│ lion   ┆ ["water"]                   │
└────────┴─────────────────────────────┘

任何帮助将不胜感激。 短暂性脑缺血发作。

Answer 1

如果您使用 python/numpy/sets 作为极坐标中的成员资格，那么您大多数时候都走在错误的道路上。这实际上对于一般使用

apply

很重要。

让我们从您离开的地方开始。您获取了

1

列的第

"food"

索引，其 dtype 为

List

。这样我们就将

Series

分配给

refn

。

dl = df.groupby('animal', maintain_order=True).agg_list()
refn = dl['food'][1]
refn

shape: (3,)
Series: 'food' [str]
[
    "rabbit"
    "deer"
    "water"
]

选项 1（较慢）

此解决方案在每个元素中进行

O(n)

搜索，因此这可能不是最快的。

我们简单地将原始

filter

与

df

的成员资格进行

refn

并再次聚合结果。

(df.filter(pl.col("food").is_in(refn))
   .groupby("animal").agg_list()
)

选项 2（最快）

更好的选择是使用

semi

连接。

semi

连接通过另一个

DataFrame

中的成员身份过滤

DataFrame

，听起来像我们想要的东西！

(df.join(refn.to_frame(), on="food", how="semi")
   .groupby("animal").agg_list()
)