Pandas .loc 需要很长时间

Question

我有一个 10 GB 的 csv 文件，其中包含 170,000,000 行和 23 列，我将其读入数据框，如下所示：

import pandas as pd 

d = pd.read_csv(f, dtype = {'tax_id': str})

我还有一个包含近 20,000 个独特元素的字符串列表：

h = ['1123787', '3345634442', '2342345234', .... ]

我想在数据框d 中创建一个名为class 的新列。只要d['tax_id'] 具有在字符串列表h 中找到的值，我就想分配d['class'] = 'A'。否则，我想要d['class'] = 'B'。

以下代码在我的数据框d 的 1% 样本上运行得非常快：

d['class'] = 'B'
d.loc[d['tax_num'].isin(h), 'class'] = 'A' 

但是，在完整的数据帧 d 上，此代码需要 48 小时（并且还在继续）才能在 32 核服务器上以批处理模式运行。我怀疑使用loc 进行索引会减慢代码速度，但我不确定它到底是什么。

总之：有没有更有效的方法来创建class 列？

Answer 1

如果您的税号是唯一的，我建议将tax_num 设置为索引，然后对其进行索引。就目前而言，您调用isin，这是一个线性运算。无论您的机器有多快，它都无法在合理的时间内对 1.7 亿条记录进行线性搜索。

df.set_index('tax_num', inplace=True) # df = df.set_index('tax_num')
df['class'] = 'B'
df.loc[h,  'class'] = 'A'

如果您仍然遇到性能问题，我建议您使用 dask 切换到分布式处理。

Answer 2

“我还有一个包含近 20,000 个独特元素的字符串列表”

好吧，首先，如果您打算将该列表用于成员资格测试，您应该将该列表设为set。 list 对象具有线性时间成员资格测试，set 对象具有非常优化的成员资格测试的恒定时间性能。这是这里最容易实现的目标。所以用

h = set(h) # convert list to set
d['class'] = 'B'
d.loc[d['tax_num'].isin(h), 'class'] = 'A'