itertools.groupby 的反面？

Posted 2023-03-15

【中文标题】itertools.groupby 的反面？

【英文标题】：reverse of itertools.groupby?

【发布时间】：2021-12-03 10:51:46

【问题描述】：

我正在组合生成器以进行一些数据处理。我首先批处理数据生成器，以便在 API 调用中进行线程化，例如：

from itertools import groupby, count
def batch(data: List[Any], size=4):
    c = count()
    for _, g in groupby(data, lambda _: next(c)//size):
        yield g  

然后我将其提供给线程以进行 API 调用

from concurrent.futures import ThreadPoolExecutor
def thread(data: Iterable, func: Callable, n=4):
    with ThreadPoolExecutor(max_workers=n) as executor:
        for batch in data:
            yield executor.map(func, batch) 

现在我正在尝试将批次合并回列表/生成器中，以便在生成器管道的下游使用。我试过这个

from itertools import chain
def flat_map(batches: Iterable):
    for i in list(chain(batches)):
        yield i

但i 似乎仍然是生成器，而不是列表中的项目？

【问题讨论】：

如果可以复制&粘贴&测试，我可以测试我认为可以解决的问题...

chain(*batches)?

【参考方案1】：

所以我最终将三个函数浓缩为一个：

from itertools import chain, groupby
from concurrent.futures import ThreadPoolExecutor

def spread(data: Iterable, func: Callable, n=4):
    """ Combines `batch`, `thread` and `flat_map`"""
    c = count()
    with ThreadPoolExecutor(max_workers=n) as executor:
        for _, batch in groupby(data, lambda _: next(c)//n):
            yield from executor.map(func, batch)

所以我只需要yield from 就可以让它工作。 感谢@ShadowRanger！

【参考方案2】：

您想要chain(*batches) 或chain.from_iterable(batches)。 chain(batches) 基本上只是产生与直接使用 batches 相同的值，它只是增加了一层包装。所以正确的代码（没有listifying，这里几乎肯定是错误的）只是：

from itertools import chain
def flat_map(batches: Iterable):
    return chain.from_iterable(batches)  # chain(*batches) would also work, but if batches is an iterator itself, it would be forced to eagerly run to completion first; chain.from_iterable can begin work when the first batch is ready

你甚至不需要yield，因为迭代器已经在产生你想要的东西了。如果您需要它成为真正的生成器，只需将 return 替换为 yield from 即可获得类似的结果。

另请注意：您可能只需更改即可完全避免对该功能的需要：

yield executor.map(func, batch) 

到：

yield from executor.map(func, batch) 

所以thread 一开始就变平了。

【讨论】：

flat_map = chain.from_iterable 怎么样？

@don'ttalkjustcode：当然，为什么不呢？如果 typing 注释匹配，则它是等效的。有些人喜欢命名包装器，但无论哪种方式都可以。