一个Python中优雅的数据分块方法

背景

看到这个标题你可能想一个分块能有什么难度？还值得细说吗，最近确实遇到一个有意思的分块函数，写法比较巧妙优雅，所以写一个分享。

日前在做需求过程中有一个对大量数据分块处理的场景，具体来说就是几十万量级的数据，分批处理，每次处理100个。这时就需要一个分块功能的代码，刚好项目的工具库中就有一个分块的函数。拿过函数来用，发现还挺好用的，传入列表和分块大小，然后就能遍历取出分好的数据。调用方式如下：

from xxx import chunk_fun

chunk_list = chunk_fun(arr, 100) # 对数据进行分块，指定块的大小为100

for chunk in chunk_list:
    print(chunk)

然后我就对这个分块函数产生了兴趣，想看看这个小功能是如何实现的。如果让我来写一个分块函数，我知道Python中range函数可以指定步长，用这个特性就完全可以优雅的实现分块功能。

arr = [1,2,3,4,5,6,7,8,9,10]

step = 3
for i in range(0, len(arr), step):
    chunk = arr[i:i+step]
    print(chunk)
>>>
[1, 2, 3]
[4, 5, 6]
[7, 8, 9]
[10]

没想到看到源码竟然才用了3行代码就实现了分块，不仅支持列表等线性结构的分块，而且还支持集合这种非线性结构的分块。这让我感到震撼，这3行代码不是最优雅的分块方法，也是接近最优雅的分块方法了。废话不多说，先上代码：

from itertools import islice

def chunk_list(it, limit):
    it = iter(it)
    return iter(lambda: list(islice(it, limit)), [])

对于这3行代码，有多少人第一眼没看出功能的呢？反正我第一眼看的是一脸懵逼，有种不明觉厉的感觉。

首先来看一下这个分块函数的使用。

set_num = {1,2,3,4,5,6,7}

for temp_list in chunk_list(set_num, 2):
    print(temp_list)
>>>
[1, 2]
[3, 4]
[5, 6]
[7]

完全没有使用显示循环就把分块这件事安排的明明白白的，而且才用了3行代码，不包括函