优化 Python 中的数据结构访问

此工作得到了 Anaconda Inc 和 Moore Foundation 的数据驱动发现计划的支持。

上周在 PyCon DE 大会上，我有幸结识了 Stefan Behnel，他是 Cython 的核心开发者之一。我们一起合作优化了一个小型基准测试，该测试代表了 Dask 的中央任务调度器，这是一个纯 Python 应用，主要受数据结构访问限制。

我们的基准测试是一个玩具问题，它创建了三个相互索引的数据结构，使用了字典、列表和集合，然后进行一些简单的算术运算。（你不需要深入理解这个基准测试来阅读本文。）

import random
import time

nA = 100
nB = 100
nC = 100

A = {'A-%d' % i: ['B-%d' % random.randint(0, nB - 1)
                  for i in range(random.randint(0, 5))]
     for i in range(nA)}

B = {'B-%d' % i: {'C-%d' % random.randint(0, nC - 1)
                  for i in range(random.randint(1, 3))}
     for i in range(nB)}

C = {'C-%d' % i: i for i in range(nC)}

data = ['A-%d' % i for i in range(nA)]


def f(A, B, C, data):
    for a_key in data:
        b_keys = A[a_key]
        for b_key in b_keys:
            for c_key in B[b_key]:
                C[c_key] += 1


start = time.time()

for i in range(10000):
    f(A, B, C, data)

end = time.time()

print("Duration: %0.3f seconds" % (end - start))

$ python benchmark.py
Duration: 1.12 seconds

这对于 Python 优化来说是一个非典型的问题，因为它主要受数据结构访问（字典、列表、集合）限制，而不是通常由 Cython 优化的数值运算（浮点算术的嵌套 for 循环）。Python 在这方面已经相当快，通常与 Java 或 C++ 等编译语言相比，速度差距在 2-5 倍以内，但我们仍然希望在可能的情况下改进这一点。

在这篇文章中，我们结合两种不同的方法来优化受数据结构限制的工作负载

1. 使用 Cython 编译 Python 代码，不添加其他注释
2. 字符串驻留以加快字典查找速度

最后，在文章末尾，我们还在 PyPy 下运行基准测试以比较性能。

Cython

首先，我们使用 Cython 编译 Python 代码。通常在使用 Cython 时，我们使用类型注释变量，以便给编译器足够的信息完全避免使用 Python。然而，在我们的例子中，数值运算不多，而且无论如何都会使用 Python 数据结构，所以这不会有太大帮助。我们编译原始的 Python 代码，不进行任何修改。

cythonize -i benchmark.py

然后运行

$ python -c "import benchmark"
Duration: 0.73 seconds

这给我们带来了相当可观的加速，从 1.1 秒降至 0.73 秒。与典型的 Cython 加速（通常是 10-100 倍）相比，这不算巨大，但这对于我们的调度器来说将是一个非常受欢迎的变化，因为我们一直在追逐 5% 的优化很久了。

字符串驻留

我们的第二个技巧是字符串驻留。这意味着我们尽量始终只保留每个字符串的一个副本。这样做可以提高字典查找的性能，原因如下：

1. Python 只计算一次字符串的哈希值（字符串计算后会缓存其哈希值）
2. Python 在进行值相等性检查（慢）之前，会先进行对象身份检查（快）

或者，从经验上看，在 Python 中 text is text 比 text == text 更快。如果你确保每个字符串只有一个副本，那么只需要进行身份比较，例如 text is text。

因此，如果在我们的程序中任何时候看到的字符串 "abc" 都与所有其他 "abc" 字符串是完全相同的对象，那么字符串-字典查找只需要进行指针/整数相等性检查，而无需进行完整的字符串比较。

向我们的基准测试添加字符串驻留如下所示

inter = {}

def intern(x):
    try:
        return inter[x]
    except KeyError:
        inter[x] = x
        return x

A = {intern('A-%d' % i): [intern('B-%d' % random.randint(0, nB - 1))
                  for i in range(random.randint(0, 5))]
     for i in range(nA)}

B = {intern('B-%d' % i): {intern('C-%d' % random.randint(0, nC - 1))
                  for i in range(random.randint(1, 3))}
     for i in range(nB)}

C = {intern('C-%d' % i): i for i in range(nC)}

data = [intern('A-%d' % i) for i in range(nA)]

# The rest of the benchmark is as before

这使得我们的持续时间从 1.1 秒降至 0.75 秒。请注意，这不包含上面描述的单独的 Cython 改进。

Cython + 驻留

我们可以结合这两种优化。这使得时间降至约 0.45 秒，相对于原始时间提升了 2-3 倍。

cythonize -i benchmark2.py

$ python -c "import benchmark2"
Duration: 0.46 seconds

PyPy

或者，我们可以直接在 PyPy 中运行一切。

$ pypy3 benchmark1.py  # original
Duration: 0.25 seconds

$ pypy3 benchmark2.py  # includes interning
Duraiton: 0.20 seconds

因此，PyPy 在这类代码上可以比 Cython 快很多（这并不令人意外）。驻留也有一些帮助，但不如在 CPython 中那么显著。

这是相当令人鼓舞的。即使 Dask 客户端和工作进程在正常的 CPython 下运行（以便使用完整的 PyData 技术栈），Dask 调度器也可以在 PyPy 下运行。

Dask 基准测试的初步结果

我们开始这个实验时假设我们的玩具基准测试在性能特征上某种程度上代表了 Dask 的调度器。这个假设当然是错误的。Dask 调度器要复杂得多，很难构建一个简单的玩具示例来代表其性能。

当我们在一个稍微复杂一点的基准测试上尝试这些技巧时，该测试实际使用了 Dask 调度器，我们发现以下结果：

• Cython：几乎没有效果
• 字符串驻留：几乎没有效果
• PyPy：几乎没有效果

然而，我在这方面只花了很短的时间（二十分钟？），所以我希望这里缺乏性能提升是因为投入的精力不够。

如果有人对此感兴趣，我希望这篇博文提供了足够的信息，以便他们能够进一步研究。