如何更快地迭代Pandas数据框？

Question

我有一个字符串列表，比如说：

fruit_list = ["apple", "banana", "coconut"]

我有一些Pandas数据框，例如：

import pandas as pd

data = [['Apple farm', 10], ['Banana field', 15], ['Coconut beach', 14], ['corn field', 10]]
df = pd.DataFrame(data, columns = ['fruit_source', 'value'])

我想要根据现有列 'fruit_source' 的文本搜索来填充一个新的列。我想要填充的是与数据框中特定列匹配的元素。一种编写方式如下：

df["fruit"] = NaN

for index, row in df.iterrows():
    for fruit in fruit_list:
        if fruit in row['fruit_source']:
            df.loc[index,'fruit'] = fruit
        else:
            df.loc[index,'fruit'] = "fruit not found"

在这里，DataFrame会增加一列来显示水果来源采集的水果种类。

然而，在扩展到较大的DataFrame时，这种迭代可能会对性能造成问题。原因是，随着引入更多行，由于同时迭代列表，迭代会变得非常慢。

是否存在更有效的方法可以实现?

- user7298979

Answer 1

你可以让Pandas这样做：

# Prime series with the "fruit not found" value
df['fruit'] = "fruit not found"
for fruit in fruit_list:
    # Generate boolean series of rows matching the fruit
    mask = df['fruit_source'].str.contains(fruit, case=False)
    # Replace those rows in-place with the name of the fruit
    df['fruit'].mask(mask, fruit, inplace=True)

print(df) 将会输出：

    fruit_source  value            fruit
0     Apple farm     10            apple
1   Banana field     15           banana
2  Coconut beach     14          coconut
3     corn field     10  fruit not found

- AKX

Answer 2

使用正则表达式模式的 str.extract 来避免循环:

import re

pattern = fr"({'|'.join(fruit_list)})"
df['fruit'] = df['fruit_source'].str.extract(pattern, flags=re.IGNORECASE) \
                                .fillna('fruit not found')

输出：

>>> df
    fruit_source  value            fruit
0     Apple farm     10            Apple
1   Banana field     15           Banana
2  Coconut beach     14          Coconut
3     corn field     10  fruit not found

>>> pattern
'(apple|banana|coconut)'

- Corralien

---

我比较了答案，看起来你的方法在处理大约450k条记录的数据框时最有效率。这似乎是最高效的方法。谢谢！ - user7298979