使用NumPy数据类型的Python字典查找速度问题

Question

背景

我有许多数值消息代码存储在NumPy数组中，需要快速将它们转换为字符串。我在性能方面遇到了一些问题，希望了解为什么会出现这种情况以及如何快速处理。

一些基准测试

I - 简单的方法

import numpy as np

# dictionary to use as the lookup dictionary
lookupdict = {
     1: "val1",
     2: "val2",
    27: "val3",
    35: "val4",
    59: "val5" }

# some test data
arr = np.random.choice(lookupdict.keys(), 1000000)

# create a list of words looked up
res = [ lookupdict[k] for k in arr ]

字典查询需要我咖啡休息时间中的大部分时间，758毫秒。（我也尝试过 res = map(lookupdict.get, arr) 但那更糟糕。）

II - 没有 NumPy

import random

# dictionary to use as the lookup dictionary
lookupdict = {
     1: "val1",
     2: "val2",
    27: "val3",
    35: "val4",
    59: "val5" }

# some test data
arr = [ random.choice(lookupdict.keys()) for _ in range(1000000) ]

# create a list of words looked up
res = [ lookupdict[k] for k in arr ]

时间结果相比之前变化了很多，现在为76毫秒！

需要注意的是我只对查找时间感兴趣。随机生成只是为了创建一些测试数据。它花费了多少时间并不重要。这里给出的所有基准结果仅针对100万次查找。

III - 将NumPy数组转换为列表

我最初的猜想是这与列表与数组的问题有关。但是，通过修改使用列表的NumPy版本：

res = [ lookupdict[k] for k in list(arr) ]

我得到了778毫秒的时间，其中大约110毫秒用于转换列表，570毫秒用于查找。因此，查找速度略快，但总时间相同。

IV - 从np.int32转换为int类型

由于唯一的其他区别似乎是数据类型(np.int32 vs. int)，因此我尝试了即时转换类型。这有点愚蠢，因为可能dict也会这样做：

res = [ lookupdict[int(k)] for k in arr ]

然而，这似乎有些有趣，因为时间降至266毫秒。看起来，几乎但不完全相同的数据类型会对字典查找产生恶劣影响，而字典代码在转换方面效率不高。

V - 将字典键转换为np.int32

为了测试这一点，我修改了NumPy版本，使用与字典键和查找完全相同的数据类型：

import numpy as np

# dictionary to use as the lookup dictionary
lookupdict = {
     np.int32(1): "val1",
     np.int32(2): "val2",
    np.int32(27): "val3",
    np.int32(35): "val4",
    np.int32(59): "val5" }

# some test data
arr = np.random.choice(lookupdict.keys(), 1000000)

# create a list of words looked up
res = [ lookupdict[k] for k in arr ]

这个时间提高到了177毫秒。虽然这是一个不小的提升，但仍然远远落后于76毫秒。

VI - 将数组转换为使用int对象

import numpy as np

# dictionary to use as the lookup dictionary
lookupdict = {
     1: "val1",
     2: "val2",
    27: "val3",
    35: "val4",
    59: "val5" }

# some test data
arr = np.array([ random.choice(lookupdict.keys()) for _ in range(1000000) ], 
               dtype='object')

# create a list of words looked up
res = [ lookupdict[k] for k in arr ]

这个结果给出了86毫秒，已经非常接近本机Python的76毫秒。

结果概述

1. 使用int作为字典键，在使用int进行索引时（本地Python）：76毫秒。
2. 使用int作为字典键，在使用int对象进行索引时（NumPy）：86毫秒。
3. 使用np.int32作为字典键，在使用np.int32进行索引时：177毫秒。
4. 使用int作为字典键，在使用np.int32进行索引时：758毫秒。

问题

为什么？我该怎么做才能使字典查找尽可能快？我的输入数据是一个NumPy数组，因此到目前为止最好的选择（最快但难看）是将字典键转换为np.int32。（不幸的是，字典键可能分布在大范围的数字上，因此逐个数组进行索引不是可行的替代方案。虽然它会很快，只需要10毫秒。）

- DrV

Answer 1

正如你所怀疑的那样，这里有问题的是 int32.__hash__，它比 int.__hash__ 慢了 11 倍：

%timeit hash(5)
10000000 loops, best of 3: 39.2 ns per loop
%timeit hash(np.int32(5))
1000000 loops, best of 3: 444 ns per loop

（int32类型是用C实现的。如果你真的很好奇，可以挖掘源代码并找出它在那里执行什么操作，这需要花费很长时间。）

---

编辑：

第二个减慢速度的部分是字典查找中隐式的==比较：

a = np.int32(5)
b = np.int32(5)
%timeit a == b  # comparing two int32's
10000000 loops, best of 3: 61.9 ns per loop
%timeit a == 5  # comparing int32 against int -- much slower
100000 loops, best of 3: 2.62 us per loop

这就解释了为什么你的V比I和IV快得多。当然，坚持使用全int的解决方案会更快。

---

据我所见，你有两个选择：

1. 坚持使用纯的 int 类型，或在查找字典之前转换为 int。
2. 如果最大的代码值不太大，和/或内存不是问题，可以将字典查找换成列表索引，这不需要 hash。

E.g.:

lookuplist = [None] * (max(lookupdict.keys()) + 1)
for k,v in lookupdict.items():
    lookuplist[k] = v

res = [ lookuplist[k] for k in arr ] # using list indexing

（编辑：您可能还想尝试在此处使用{{link1：np.choose}}进行实验）

- shx2

---

谢谢你的见解！如果我能给予超过+1的奖励，我会这么做。我正在深入研究NumPy源代码，但在其中找到自己的路并不容易...（对于最后一个解决方案，是的，特别是使用NumPy数组索引，速度要快得多，但正如我在帖子中所说，查找数组可能会变得太大。） - DrV

Answer 2

这很有趣，我可能找到了我的问题的答案。

第三种选择是将数组转换为列表。如果做得正确，似乎这提供了非常好的结果。 这样做：

res = [ lookupdict[k] for k in list(arr) ]

时钟 778 毫秒。

但是这个：

res = [ lookupdict[k] for k in arr.tolist() ]

时钟86毫秒。

这背后的技术解释是，arr.tolist将数组转换为int对象，而list(arr)创建一个由np.int32对象组成的列表。

- DrV

Answer 3

在我的测试中，你的II - Without NumPy比I要慢得多。

In [11]: timeit [lookupdict[k] for k in np.random.choice(lookupdict.keys(),1000000)]
1 loops, best of 3: 658 ms per loop

In [12]: timeit [lookupdict[k] for k in [np.random.choice(lookupdict.keys()) for _ in range(1000000)]]
1 loops, best of 3: 8.04 s per loop

然而，如果通过对值进行选择而跳过查找过程，您将节省更多时间。

In [34]: timeit np.random.choice(lookupdict.values(),1000000)
10 loops, best of 3: 85.3 ms per loop

---

好的，让我们专注于查找：

In [26]: arr =np.random.choice(lookupdict.keys(),1000000)

In [27]: arrlist=arr.tolist()

In [28]: timeit res = [lookupdict[k] for k in arr]
1 loops, best of 3: 583 ms per loop

In [29]: timeit res = [lookupdict[k] for k in arrlist]
10 loops, best of 3: 120 ms per loop

In [30]: timeit res = [lookupdict[k] for k in list(arr)]
1 loops, best of 3: 675 ms per loop

In [31]: timeit res = [lookupdict[k] for k in arr.tolist()]
10 loops, best of 3: 156 ms per loop

In [32]: timeit res = [k for k in arr]
1 loops, best of 3: 215 ms per loop

In [33]: timeit res = [k for k in arrlist]
10 loops, best of 3: 51.4 ms per loop

In [42]: timeit arr.tolist()
10 loops, best of 3: 33.6 ms per loop

In [43]: timeit list(arr)
1 loops, best of 3: 264 ms per loop

第一个观察结果 - 迭代np.array比迭代等效的列表慢。

第二个观察结果 - list(arr)比arr.tolist()慢。 list()似乎有两个问题。 它本身较慢，并且项目是np.int32。

- hpaulj

---

啊，我没有表达清楚我只对查找部分计时感兴趣。当然，使用列表推导式的随机部分会慢得多。我将编辑我的问题以反映这一点。 - DrV

似乎您最新的编辑得到的答案和我所做的完全一样，确切地讲是在同一分钟内，因此我认为授予您答案分数是公平的。 - DrV

Answer 4

以下是使用Pandas的解决方案，可以提高五倍效率：

import numpy as np
import pandas as pd

# dictionary to use as the lookup dictionary
lookupdict = {
 1: "val1",
 2: "val2",
27: "val3",
35: "val4",
59: "val5" }

# some test data
arr = np.random.choice(lookupdict.keys(), 1000000)

# create a list of words looked up
%timeit res = [ lookupdict[k] for k in arr ]
%timeit res_pd = pd.Series(lookupdict).reindex(arr).values
print all(res == res_pd)

10 loops, best of 3: 192 ms per loop
10 loops, best of 3: 35.3 ms per loop
True

这是每个元素平均需要35ns，因此在本机Python中几乎不可能被超越。如果您不熟悉Pandas，Series对象类似于OrderedDict或索引数组，可以从标准Python字典构造。reindex方法提供了非常快的查找；我不确定原理是什么，因为我不是非常有经验的程序员，但它可能是用C或Cython编写的。也许您可以查看源代码，并想出一个更快的定制解决方案来解决问题。最后，values属性只返回Series下面的数组。

编辑： 下面是纯numpy解决方案，几乎与Pandas一样好：

keys = np.array(lookupdict.keys())
strings = np.array(lookupdict.values())
%timeit res_np = strings[(np.atleast_2d(arr).T == keys).argmax(axis=1)]
10 loops, best of 3: 44.6 ms per loop

print all(res == res_np)
True

- JoeCondron

Answer 5

这似乎是处理整数键时最快的方法，只需使用数组索引。

import numpy as np

# dictionary to use as the lookup dictionary
lookupdict = {
     1: "val1",
     2: "val2",
    27: "val3",
    35: "val4",
    59: "val5" }

# some test data
arr = np.random.choice(lookupdict.keys(), 1000000)
lookup_array=np.array([None if i not in lookupdict else lookupdict[i] for i in range(60)])
%timeit lookup_array[arr]

"5.3毫秒±73.7微秒每个循环（平均值±7次运行的标准差，每个循环100次）"

- philn

Answer 6

你在问题中没有考虑的一个选项，尽管这是一种受限制的选项，在某些情况下可能行不通，那就是将你的lookupdict转换为数组。 在我的机器上，对于像你的例子这样小的字典，它非常快。

import numpy as np

# dictionary to use as the lookup dictionary
lookupdict = {
     1: "val1",
     2: "val2",
    27: "val3",
    35: "val4",
    59: "val5" }

# some test data
arr = np.random.choice(lookupdict.keys(), 1000000)

table = np.empty(max(lookupdict.keys()) + 1, dtype='S4')
for key, value in lookupdict.items():
    table[key] = value

res = table[arr]

- Bi Rico

---

我没有写得够清楚，但我尝试了一下（在我的问题末尾的括号中）。是的，10毫秒。它非常快，但不幸的是，我的键可能分散在int32空间的各个位置。 - DrV