通过索引和字符串连接改善python代码的运行时

Question

我努力改善以下代码段的运行时间，这些代码段原来是我正在开发的asyncio-client程序包中的CPU瓶颈：

data = [''] * n
for i, ix in enumerate(indices):
    data[ix] = elements[i]
s = '\t'.join(data)
return s

我所做的基本上很简单：

• elements是str（每个<= 7个字符）的列表，我最终在特定位置将它们写入制表符分隔的文件中。
• indices是int的列表，给出了文件中每个elements的位置
• 如果在某个位置没有元素，则会插入一个空字符串

我最终使用aiofiles将字符串写入文本文件。

到目前为止，我尝试使用生成器动态创建数据，并使用numpy进行更快的索引编制，但是没有成功。任何使该代码运行得更快的想法都很好。这是一个带有时间的可复制示例：

import numpy as np
import timeit

n = 1_000_000  # total number of items
k = 500_000  # number of elements to insert
elements = ['ade/gua'] * k  # elements to insert, <= 7 unicode characters
indices = list(range(0, n, 2))  # indices where to insert, sorted
assert len(elements) == len(indices)

# This is where I started
def baseline():
    data = [''] * n
    for i, ix in enumerate(indices):
        data[ix] = elements[i]
    s = '\t'.join(data)
    return s

# Generate values on the fly
def generator():
    def f():
        it = iter(enumerate(indices))
        i, ix = next(it)
        for j in range(n):
            if j == ix:
                yield elements[i]
                try:
                    i, ix = next(it)
                except:
                    pass
            else:
                yield ''
    s = '\t'.join(f())  # iterating though generation seem too costly
    return s

# Harness numpy
indices_np = np.array(indices)  # indices could also be numpy array
def numpy():
    data = np.full(n, '', dtype='<U7')
    data[indices_np] = elements  # this is faster with numpy
    s = '\t'.join(data)  # much slower. array2string or savetxt does not help
    return s

assert baseline() == generator() == numpy()

timeit.timeit(baseline, number=10)  # 0.8463204780127853
timeit.timeit(generator, number=10)  # 2.048296730965376 -> great job
timeit.timeit(numpy, number=10)  # 4.486689139157534 -> life sucks

编辑1

要解决评论中提出的一些问题：

• 我写了字符串aiofiles.open(filename, mode='w') as file和file.write()

• 索引通常不能表示为范围

• 可以假定索引始终是免费排序的。

• ASCII字符足够

编辑2

根据Mad Physicist的回答，我尝试了以下代码，但没有成功。

def buffer_plumbing():
m = len(elements) # total number of data points to insert
k = 7  # each element is 7 bytes long, only ascii 
total_bytes = n - 1 + m * 7  # total number of bytes for the buffer

# find out the number of preceeding gaps for each element
gap = np.empty_like(indices_np)
gap[0] = indices_np[0]  # that many gaps a the beginning
np.subtract(indices_np[1:], indices_np[:-1], out=gap[1:])
gap[1:] -= 1  # subtract one to get the gaps (except for the first)

# pre-allocate a large enough byte buffer
s = np.full(total_bytes , '\t', dtype='S1')

# write element into the buffer
start = 0
for i, (g, e) in enumerate(zip(gap, elements)):
    start += g
    s[start: start + k].view(dtype=('S', k))[:] = e
    start += k + 1
return s.tostring().decode('utf-8')

timeit.timeit(buffer_plumbing, number=10)  # 26.82

Answer 1

在将数据转换为一对numpy数组后，您可以对其进行预排序。这将允许您操作单个预先存在的缓冲区，而不是在重新分配字符串时一遍又一遍地复制字符串。我的建议与您的尝试之间的区别在于，假设您只有ASCII字符，我们将使用ndarray.tobytes（或ndarray.tostring）。实际上，您可以直接使用ndarray.tofile来完全绕开涉及转换为bytes对象的复制操作。

如果手头有elements，则知道行的总长度将是elements和n-1制表符分隔符的总长度。因此，完整字符串中元素的开始是它的索引（它前面的选项卡数）加上它前面所有元素的累积长度。以下是使用大多数Python循环的单缓冲区填充的简单实现：

lengths = np.array([len(e) for e in elements])
indices = np.asanyarray(indices)
elements = np.array(elements, dtype='S7')
order = np.argsort(indices)

elements = elements[order]
indices = indices[order]
lengths = lengths[order]

cumulative = np.empty_like(lengths)
cumulative[0] = 0
np.cumsum(lengths[:-1], out=cumulative[1:])
cumulative += lengths

s = np.full(cumulative[-1] + n - 1, '\t', dtype='S1')
for i, l, e in zip(cumulative, lengths, elements):
    s[i:i + l].view(dtype=('S', l))[:] = e

这里有很多可能的优化方法，例如使用np.empty分配s并仅用制表符填充所需元素的可能性。这将留给读者作为消费税。

另一种可能性是避免完全将elements转换为一个numpy数组（这可能只会浪费空间和时间）。然后，您可以将for循环重写为

for i, l, o in zip(cumulative, lengths, order):
    s[i:i + l].view(dtype=('S', l))[:] = elements[o]

您可以使用

将结果转储到bytes对象中

s = s.tobytes()

OR

s = s.tostring()

您可以按原样将结果写入打开的文件中以进行二进制写入。实际上，如果您不需要bytes形式的缓冲区副本，则可以直接写入文件：

s.tofile(f)

这将为您节省一些内存和处理时间。

出于同样的考虑，您最好直接逐段直接写入文件。这样不仅节省了分配完整缓冲区的需要，还节省了累积长度。实际上，您唯一需要的方法是连续索引的diff，告诉您要插入多少个标签：

indices = np.asanyarray(indices)
order = np.argsort(indices)

indices = indices[order]

tabs = np.empty_like(indices)
tabs[0] = indices[0]
np.subtract(indices[1:], indices[:-1], out=tabs[1:])
tabs[1:] -= 1

for t, o in zip(tabs, order):
    f.write('\t' * t)
    f.write(elements[o])
f.write('\t' * (n - indices[-1] - 1))

第二种方法除了减少计算量外，还有两个主要优点。第一个是它可与Unicode字符一起使用，而不仅仅是ASCII。第二个问题是，除了制表符字符串之外，它不分配任何缓冲区，这应该非常快。

在两种情况下，将elements和indices按索引升序排列将大大加快处理速度。第一种情况减少为

lengths = np.array([len(e) for e in elements])
indices = np.asanyarray(indices)

cumulative = np.empty_like(lengths)
cumulative[0] = 0
np.cumsum(lengths[:-1], out=cumulative[1:])
cumulative += lengths

s = np.full(cumulative[-1] + n - 1, '\t', dtype='S1')
for i, l, e in zip(cumulative, lengths, elements):
    s[i:i + l].view(dtype=('S', l))[:] = e

第二个就变成

indices = np.asanyarray(indices)

tabs = np.empty_like(indices)
tabs[0] = indices[0]
np.subtract(indices[1:], indices[:-1], out=tabs[1:])
tabs[1:] -= 1

for t, e in zip(tabs, elements):
    f.write('\t' * t)
    f.write(e)
f.write('\t' * (n - indices[-1] - 1))