itertools.product在无限生成器上会导致内存溢出并崩溃
我正在尝试使用 itertools.product 在两条数据流之间生成组合。一条是标准列表,另一条是从实时API拉取数据的无限生成器。
由于 itertools 是为“高效循环”和惰性求值而设计的,我以为它能够在不将所有数据加载到内存中的情况下处理流。然而,一旦我开始循环,脚本的内存使用量就会立即暴涨,直到发生内存溢出(OOM)崩溃。
下面是会导致进程崩溃的最小代码:
import itertools
def infinite_stream():
i = 0
while True:
yield f"data_{i}"
i += 1
# This line causes an immediate OOM crash
for item in itertools.product([1, 2], infinite_stream()):
print(item)
为什么 itertools.product 在无限生成器上会崩溃,尽管它是一个“迭代器”,以及我该如何在不将整个流加载到内存中的情况下处理这些组合?
解决方案
Why does
itertools.productcrash on an infinite generator even though it is an "iterator," and how do I process these combinations without loading the entire stream into RAM?
考虑到可迭代对象也可能是生成器,对它进行浅拷贝不会产生一个从头开始生成值的新可迭代对象。同样,请注意,生成器不能进行深拷贝。请考虑以下内容:
def stream():
yield "a"
yield "b"
yield "c"
it = stream()
for i in itertools.product([0, 1], it):
print(i)
以上将产生输出 (0, 'a')、(0, 'b')、(0, 'c')、(1, 'a')、(1, 'b') 和 (1, 'c')。请注意,作为迭代器的 it 已经在 0 和 1 下产生了两条有效的输出流。在某些时刻,迭代器产生的所有值可能需要在后续运行中重新产生。鉴于对迭代器进行浅克隆不会从头重新生成输出,因此需要分配内存来存储由提供的生成器产生的所有值。product 的底层C 实现简单地在产出前将来自每个输入可迭代对象的所有值收集到它们的对应元组中 然后再产出乘积,即在对无限生成器尚未产生一个值时就触发内存溢出(OOM),而不是在跟踪来自无限生成器的值直到内存耗尽时再产出值。
如果你需要使用无限生成器来产生某些值子集,只需使用嵌套的for循环来产出值。例如,如果你有以下内容:
def demo():
for i in itertools.product(infinite_stream(), stream()):
yield i
改写为:
def demo():
for a in infinite_stream():
for b in stream():
yield (a, b)
这样做的原因是,每次外层循环迭代时,stream() 都会产生一个全新的迭代器,在内层循环中产出值,从而产生新值,并且不会让额外的缓存值在某处吞噬内存。
当然,如果 product 接受产生可迭代对象的可调用对象作为参数,那么就不需要缓存值,因为每个内层循环都会从该可调用对象产生一个新的可迭代对象。