分割长字符串的最快方法

Question

我在Python中有一个很长的字符串：

x = "12;14;14;14;18;12;17;19" # I only show a small part of it : there are 10 millions of ;

目标是将其转变为：

y = array([12, 14, 14, 14, 18, 12, 17, 19], dtype=int)

一种方法是使用array(x.split(";"))or numpy.fromtostring。

但两者都非常慢。

有没有更快的方法在Python中做到这一点？

非常感谢您，祝您度过愉快的一天。

Answer 1

字符串解析通常很慢。Unicode 解码通常会使速度变慢（特别是当存在非 ASCII 字符时），除非经过仔细优化（硬）。CPython 很慢，尤其是循环。Numpy 并不是真正设计用来（有效）处理字符串的。我认为 Numpy 还不能更快地做到这一点fromstring。我能想到的唯一解决方案是使用Numba、Cython甚至基本的C 扩展。最简单的解决方案是使用 Numba，最快的是使用 Cython/C 扩展。

不幸的是，到目前为止，Numba 对于字符串/字节来说非常慢（这是一个悬而未决的问题，预计不会很快得到解决）。需要一些技巧，以便 Numba 可以有效地计算：需要将字符串转换为 Numpy 数组。这意味着必须首先将其编码为字节数组，以避免任何可变大小的编码（如 UTF-8）。np.frombuffer似乎是将缓冲区转换为 Numpy 数组的最快解决方案。由于输入是只读数组（不常见，但高效），因此 Numba 签名不太容易阅读。

这是最终的解决方案：

import numpy as np
import numba as nb

@nb.njit(nb.int32[::1](nb.types.Array(nb.uint8, 1, 'C', readonly=True,)))
def compute(arr):
    sep = ord(';')
    base = ord('0')
    minus = ord('-')
    count = 1
    for c in arr:
        count += c == sep
    res = np.empty(count, np.int32)
    val = 0
    positive = True
    cur = 0
    for c in arr:
        if c != sep and c != minus:
            val = (val * 10) + c - base
        elif c == minus:
            positive = False
        else:
            res[cur] = val if positive else -val
            cur += 1
            val = 0
            positive = True
    if cur < count:
        res[cur] = val if positive else -val
    return res

x = ';'.join(np.random.randint(0, 200, 10_000_000).astype(str))
result = compute(np.frombuffer(x.encode('ascii'), np.uint8))

请注意，Numba 解决方案不会出于性能原因执行检查。它还假设数字是正数。因此，您必须确保输入有效。或者，您可以在其中执行额外的检查（但代价是代码速度较慢）。

以下是我的配备 i5-9600KF 处理器的机器（Windows 上使用 Numpy 1.22.4）的性能结果：

np.fromstring(x, dtype=np.int32, sep=';'):       8927 ms
np.array(re.split(";", x), dtype=np.int32):      1419 ms
np.array(x.split(";"), dtype=np.int32):          1196 ms
Numba implementation:                              78 ms
Numba implementation (without negative numbers):   67 ms

该解决方案比最快的解决方案（基于 ）快114 倍np.fromstring和15 倍split。请注意，取消对负数的支持会使 Numba 函数速度提高 18%。另请注意，10~12% 的时间花费在encode. 其余时间来自 Numba 函数中的主循环。更具体地说，循环中的条件是速度减慢的主要原因，因为处理器很难预测它们，并且它们妨碍了快速 (SIMD) 指令的使用。这通常就是字符串解析速度慢的原因。

一个可能的改进是使用对块进行操作的无分支实现。另一个可能的改进是使用多个线程来计算块。然而，这两种优化都很难做到，而且它们都会使生成的代码更难阅读（也更难维护）。