使用 NumPy 数据类型的 Python 字典查找速度答案

【问题标题】：Python dictionary lookup speed with NumPy datatypes使用 NumPy 数据类型的 Python 字典查找速度
【发布时间】：2014-07-11 20:02:51
【问题描述】：

背景

我在 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 与 int），我尝试即时转换类型。这有点愚蠢，可能 dict 也是如此：

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

但是，这似乎做了一些有趣的事情，因为时间下降到 266 毫秒。似乎几乎但不完全一样的数据类型在字典查找方面玩了讨厌的把戏，而且字典代码在转换方面效率不高。

V - 字典键转换为np.int32

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

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 毫秒。

结果总结

dict 键int，使用int（本机Python）进行索引：76 毫秒
dict 键 int，使用 int 对象 (NumPy) 进行索引：86 毫秒
dict 键np.int32，使用np.int32 进行索引：177 毫秒
dict 键int，使用np.int32 进行索引：758 毫秒

问题

为什么？我能做些什么来尽可能快地进行字典查找？我的输入数据是一个 NumPy 数组，所以到目前为止最好的（最快但丑陋的）是将 dict 键转换为np.int32。（不幸的是，dict 键可能分布在很宽的数字范围内，因此逐个数组索引不是一个可行的选择。不过很快，10 毫秒。）

【问题讨论】：

标签： python python-2.7 numpy

【解决方案1】：

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

替代方案 III 是将数组转换为列表。 如果方法正确，这似乎提供了非常好的结果。这个：

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 对象的列表。

【讨论】：

【解决方案2】：

在我的时间里，你的 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

但如果通过在值上设置choice 来跳过查找，您将获得更多时间

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。

【讨论】：

啊，我说得不够清楚，我只对查找部分的时间感兴趣。自然，列表理解的随机部分要慢得多。我将编辑我的问题以反映这一点。
您的最新编辑似乎与我在同一分钟所做的和发现的答案完全相同，所以我认为奖励您答案是公平的。

【解决方案3】：

这是一个使用 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 dict 构造。 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

【讨论】：

【解决方案4】：

正如你所怀疑的，这是 int32.__hash__ 的错，它的 x11 和 int.__hash__ 一样慢：

%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 中实现的。如果你真的很古玩，你可以在源代码中挖掘并找出它在那里做了什么，这需要很长时间。

编辑：

第二个让事情变慢的部分是隐含的== dict 查找比较：

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 解决方案会更快。

所以在我看来，你有两个选择：

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

例如：

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

（编辑：您可能还想在这里尝试np.choose）

【讨论】：

感谢您的洞察！如果我可以奖励超过+1，我会这样做。我只是在深入研究 NumPy 源代码，但是在其中找到自己的方式并不容易......（对于最后一个解决方案，是的，它更快，尤其是使用 NumPy 数组索引，但正如我所说在我的帖子中，查找数组可能会变得太大。）

【解决方案5】：

您在问题中没有考虑的一个选项是将您的 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]

【讨论】：

我写的不够清楚，但我试过了（它在我问题末尾的括号中）。是的，10 毫秒。它很快，但不幸的是我的密钥可能遍布int32 空间。

【解决方案6】：

当您的键是整数时，这似乎是最快的方法。只做数组索引。

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 ms ± 73.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

【讨论】：