Pandas，快速从字典列提取信息到新列

Question

我有一个很大的数据表，里面有四列。其中一列是一些单词，另一列是这些单词在字典里的键。我需要再添加一列，从字典中提取出我感兴趣的单词对应的值。举个例子：

ID        ID2     words      dict1
x12_12    12984   apple      {'apple': 5, 'pear': 10}
x12_12    12984   pear       {'apple': 5, 'pear': 10}
x12_12    20934   orange     {'orange': 5, 'pear': NaN}
x12_12    20934   pear       {'orange': 5, 'pear': NaN}

我需要创建一个叫做“value”的新列，从dict1中提取信息。

ID        ID2     words      dict1                         value
x12_12    12984   apple      {'apple': 5, 'pear': 10}      5
x12_12    12984   pear       {'apple': 5, 'pear': 10}      10
x12_12    20934   orange     {'orange': 20, 'pear': NaN}   20
x12_12    20934   pear       {'orange': 20, 'pear': NaN}   NaN

我有一段代码可以得到我想要的结果，但运行起来很慢，因为我的数据集非常大。我知道在处理大数据时，使用'apply'方法效率不高。

df['value'] = df.apply(lambda row: row['dict1'][row['words']], axis=1)

有没有更快的方法呢？我试过用np.vectorize，但它对nan值有问题，导致我一直出错。

Answer 1

最简单的方法是使用列表推导式和 zip：

df['value'] = [d.get(w) for w,d in zip(df['words'], df['dict1'])]

另外，可以结合 json_normalize 和 索引查找：

idx, cols = pd.factorize(df['words'])

df['value'] = (pd.json_normalize(df['dict1'])
                 .reindex(cols, axis=1).to_numpy()
               [np.arange(len(df)), idx]
              )

或者使用 groupby.transform：

df['value'] = (df.groupby('words', as_index=False)['dict1']
                 .apply(lambda g: g.str[g.name]).droplevel(0)
              )

输出结果：

       ID    ID2   words                       dict1  value
0  x12_12  12984   apple    {'apple': 5, 'pear': 10}    5.0
1  x12_12  12984    pear    {'apple': 5, 'pear': 10}   10.0
2  x12_12  20934  orange  {'orange': 5, 'pear': nan}    5.0
3  x12_12  20934    pear  {'orange': 5, 'pear': nan}    NaN

时间比较

在处理40万行数据时，最后发现列表推导式是最快的，这很合理，因为限制步骤本身就是基于循环的。apply也是一个循环，但由于每行都要创建一个Series，所以开销比较大。尽量在可以的情况下，用列表推导式替代 apply 和 axis=1。

# apply
2.58 s ± 33.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# list comprehension
127 ms ± 7.65 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# json_normalize + indexing lookup
811 ms ± 46 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# groupby.transform
237 ms ± 3.22 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)