在 pandas DataFrame 中查找值的更快方法？

Question

我正在尝试将我的一些 R 脚本“翻译”成 Python，但我注意到，在 Python 中处理数据帧比在 R 中执行它要慢得多，例如根据某些条件提取细胞。

我做了一点调查，这是在 Python 中查找特定值需要多少时间：

import pandas as pd
from timeit import default_timer as timer

code = 145896

# real df is way bigger
df = pd.DataFrame(data={
    'code1': [145896, 800175, 633974, 774521, 416109],
    'code2': [100, 800, 600, 700, 400],
    'code3': [1, 8, 6, 7, 4]}
    )

start = timer()
for _ in range(100000):
    desired = df.loc[df['code1']==code, 'code2'][0]
print(timer() - start) # 19.866242500000226 (sec)

在 R 中：

code <- 145896

df <- data.frame("code1" = c(145896, 800175, 633974, 774521, 416109),
           "code2" = c(100, 800, 600, 700, 400),
           "code3" = c(1, 8, 6, 7, 4))

start <- Sys.time()
for (i in 1:100000) {
  desired <- df[df$code1 == code, "code2"]
}
print(Sys.time() - start) # Time difference of 1.140949 secs

我对 Python 比较陌生，我可能遗漏了一些东西。有什么方法可以加快这个过程吗？也许将这个脚本转移到 Python 的整个想法毫无意义？在其他操作中，Python 更快（即处理字符串），一旦需要处理数据帧，在两个或多个脚本之间跳转会非常不方便。请问这方面有什么帮助吗？

更新 真正的脚本块迭代初始数据帧的行（相当大，500-1500k 行）并创建一个新的行，其中包含来自原始列“code1”的值和对应于另一个数据帧的代码，以及许多其他值，是新创建的。我相信，我可以用图片澄清它：

稍后在脚本中，我还需要根据不同的条件在循环中搜索特定值。因此，搜索速度至关重要。

Answer 1

只需重用过滤器表达式，您就可以将其减少一半。

In [1]: import pandas as pd

In [2]: code = 145896
   ...: df = pd.DataFrame(data={
   ...:     'code1': [145896, 800175, 633974, 774521, 416109],
   ...:     'code2': [100, 800, 600, 700, 400],
   ...:     'code3': [1, 8, 6, 7, 4]
   ...: })

In [3]: %timeit df.loc[df['code1'] == code, 'code2'][0]
197 µs ± 5.14 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [4]: filter_expr = df['code1'] == code

In [5]: %timeit df.loc[filter_expr, 'code2'][0]
106 µs ± 3.3 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

按键列索引数据帧（假设查找频繁）应该是可行的方法，因为数据帧的索引是一个哈希表（有关详细信息，请参阅this answer 和these slides）。

In [6]: df_idx = df.set_index('code1')

In [7]: %timeit df_idx.loc[code]['code2']
72.7 µs ± 1.58 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

并且根据您拥有的其他用例，拥有真正的嵌入式（内存中）数据库、SQLite 或DuckDB 的数据（可以直接对 Pandas 数据运行查询，而无需导入或复制任何数据 em>)，也可能是一个解决方案。

Answer 2

R-lang 可能是围绕这种操作设计的。 Pandas 是 python 中的第三方库，因此它有额外的限制和开销需要解决。即，在每个中间步骤上，都会生成一个新的数据帧或序列。几乎每个等号或括号都是中间步骤。

如果你真的想从单个列中一次提取单个元素，你可以尝试设置索引：

df2 = df.set_index('code1')
def getel(df2,code):
    desired = None
    if code in df2.index: 
        desired = df2['code2'][code]
        if isinstance(desired, pd.Series):
            desired = desired.iloc[0]
    return code

如果值被复制，这将比原始值快三倍。如果该值是唯一的，则desired = df2['code2'][code] 不会生成新系列并且代码比原始代码快 17 倍。此外，请注意，在其他所有内容之前选择列往往会显着减少 pandas 所做的不必要工作。

如果您想对不同的值执行类似的操作，那么您可能应该查看groupby。或者至少一次过滤所有要处理的值：

codes = {145896,774521}
pad = df['code1'].apply(lambda x: x in codes) #this is good when there are many codes
#or
pad = df['code1'].isin(codes) #this is linear time in len(codes)
result = df[pad].apply(do_stuff, axis = 1)
#or df[pad].transform(do_stuff, axis = 1)

Answer 3

由于您希望从 DataFrame 中选择一个单个值，因此您可以采取一些措施来提高性能。

1. 使用.item() 代替[0]，这有一个小而不错的改进，尤其是对于较小的DataFrame。
2. 屏蔽整个DataFrame然后选择一个已知的系列是很浪费的。而是仅屏蔽系列并选择值。尽管您可能会认为“哦，这是链式的 - 被禁止的 ][”，但它只是链式 assignment 令人担忧，而不是链式选择。
3. 使用 numpy。由于索引和对齐，Pandas 有很多开销。但是您只想从一个矩形数据结构中选择一个值，所以下拉到numpy 会更快。

---

以下是选择数据的不同方法的时序图[下面每种方法都有自己的方法]。到目前为止，使用 numpy 是最快的，尤其是对于像您的示例中这样的较小 DataFrame。对于那些，它将比您选择数据的原始方式快 20 倍以上，看看您与 R 的初始比较应该会使它比在 R 中选择数据略快。随着 DataFrame 变大， numpy 解决方案的相对性能不太好，但它仍然是最快的方法（如图所示）。

import perfplot
import pandas as pd
import numpy as np

def DataFrame_Slice(df, code=0):
    return df.loc[df['code1'] == code, 'code2'][0]

def DataFrame_Slice_Item(df, code=0):
    return df.loc[df['code1'] == code, 'code2'].item()

def Series_Slice_Item(df, code=0):
    return df['code2'][df['code1'] == code].item()

def with_numpy(df, code=0):
    return df['code2'].to_numpy()[df['code1'].to_numpy() == code].item()


perfplot.show(
    setup=lambda N: pd.DataFrame({'code1': range(N),
                                  'code2': range(50, N+50),
                                  'code3': range(100, N+100)}),
    kernels=[
        lambda df: DataFrame_Slice(df),
        lambda df: DataFrame_Slice_Item(df),
        lambda df: Series_Slice_Item(df),
        lambda df: with_numpy(df)
    ],
    labels=['DataFrame_Slice', 'DataFrame_Slice_Item', 'Series_Slice_Item', 'with_numpy'],
    n_range=[2 ** k for k in range(1, 21)],
    equality_check=np.allclose,  
    relative_to=3,
    xlabel='len(df)'
)