Pandas `isin` 函数的更快替代方案

Question

我有一个非常大的数据框 df，看起来像：

ID       Value1    Value2
1345      3.2      332
1355      2.2      32
2346      1.0      11
3456      8.9      322

我有一个包含 ID 子集ID_list 的列表。对于ID_list 中包含的ID，我需要df 的子集。

目前，我正在使用df_sub=df[df.ID.isin(ID_list)] 来执行此操作。但这需要很多时间。 ID_list 中包含的IDs 没有任何模式，因此不在一定范围内。 （而且我需要对许多相似的数据帧应用相同的操作。我想知道是否有更快的方法来做到这一点。如果将ID 作为索引会有很大帮助吗？

谢谢！

Answer 1

是的，isin 很慢。

相反，将ID 设为索引然后使用loc 会更快，例如：

df.set_index('ID', inplace=True)
df.loc[list_of_indices]

实际上将我带到此页面的原因是我需要根据另一个 df 中的索引在我的df 中创建一个标签：“如果 df_1 的索引与 df_2 的索引匹配，则将其标记为 1，否则为 NaN”，我完成了像这样：

df_2['label'] = 1  # Create a label column
df_1.join(df_2['label'])

这也非常快。

Answer 2

编辑 2：这是对各种 pandas 操作性能的最新研究的链接，尽管它似乎不包括迄今为止的合并和连接。

https://github.com/mm-mansour/Fast-Pandas

编辑 1：这些基准测试适用于相当旧的 pandas 版本，可能不再适用。请参阅下面 Mike 在merge 上的评论。

这取决于您的数据大小，但对于大型数据集DataFrame.join 似乎是要走的路。这要求您的 DataFrame 索引是您的“ID”，并且您要加入的 Series 或 DataFrame 的索引是您的“ID_list”。系列还必须有一个name 与join 一起使用，它作为一个名为name 的新字段被拉入。您还需要指定一个内连接来获得类似isin 的内容，因为join 默认为左连接。对于大型数据集，查询 in 语法似乎与 isin 具有相同的速度特征。

如果您使用的是小型数据集，您会得到不同的行为，实际上使用列表推导或应用于字典比使用 isin 更快。

否则，您可以尝试使用Cython 获得更快的速度。

# I'm ignoring that the index is defaulting to a sequential number. You
# would need to explicitly assign your IDs to the index here, e.g.:
# >>> l_series.index = ID_list
mil = range(1000000)
l = mil
l_series = pd.Series(l)

df = pd.DataFrame(l_series, columns=['ID'])


In [247]: %timeit df[df.index.isin(l)]
1 loops, best of 3: 1.12 s per loop

In [248]: %timeit df[df.index.isin(l_series)]
1 loops, best of 3: 549 ms per loop

# index vs column doesn't make a difference here
In [304]: %timeit df[df.ID.isin(l_series)]
1 loops, best of 3: 541 ms per loop

In [305]: %timeit df[df.index.isin(l_series)]
1 loops, best of 3: 529 ms per loop

# query 'in' syntax has the same performance as 'isin'
In [249]: %timeit df.query('index in @l')
1 loops, best of 3: 1.14 s per loop

In [250]: %timeit df.query('index in @l_series')
1 loops, best of 3: 564 ms per loop

# ID must be the index for DataFrame.join and l_series must have a name.
# join defaults to a left join so we need to specify inner for existence.
In [251]: %timeit df.join(l_series, how='inner')
10 loops, best of 3: 93.3 ms per loop

# Smaller datasets.
df = pd.DataFrame([1,2,3,4], columns=['ID'])
l = range(10000)
l_dict = dict(zip(l, l))
l_series = pd.Series(l)
l_series.name = 'ID_list'


In [363]: %timeit df.join(l_series, how='inner')
1000 loops, best of 3: 733 µs per loop

In [291]: %timeit df[df.ID.isin(l_dict)]
1000 loops, best of 3: 742 µs per loop

In [292]: %timeit df[df.ID.isin(l)]
1000 loops, best of 3: 771 µs per loop

In [294]: %timeit df[df.ID.isin(l_series)]
100 loops, best of 3: 2 ms per loop

# It's actually faster to use apply or a list comprehension for these small cases.
In [296]: %timeit df[[x in l_dict for x in df.ID]]
1000 loops, best of 3: 203 µs per loop

In [299]: %timeit df[df.ID.apply(lambda x: x in l_dict)]
1000 loops, best of 3: 297 µs per loop