一个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行代码，不包括函数的定义就只剩下2行代码就搞定了。这是我见过最优雅的分块方法。然后我就花一点时间搞明白代码是如何工作的。

那么这个分块功能是如何实现的呢？主要有两个知识点：迭代器切片islice+迭代器生成函数iter。通过这两个函数的配合，完成了分块功能。下面我详细介绍这两个方法的使用。

二、islice

islice是python内置模块itertool中的一个函数，功能是对迭代器切片，传入一个迭代器，返回从迭代器中的start位置到stop位置的元素，可缺省起始位置。

函数定义如下：

islice(iterable, [start, ] stop [, step])

• iterable 可迭代对象
• start 切片开始位置
• stop 切片结束位置
• step 步长

2.1示例

from itertools import islice
from collections import Iterator

iter_list = iter([1,2,3,4,5,6,7])
slice = islice(iter_list, 0, 7, 2)

print(slice)
>>>
<itertools.islice object at 0x7fc864e5aef8>

print(isinstance(slice, Iterator))
>>> 
True

print(list(slice))
>>>
[1, 3, 5, 7]

指定start为0，stop为7，step2，得到一个新的迭代器，元素是从1开始的步长为2取到的数据。

2.2只指定步长

islice可以只传入步长参数，当没有start和stop时，默认从start为起点，stop为终点。

from itertools import islice

iter_list = iter([1,2,3,4,5,6,7])

slice = islice(iter_list, 2)
print(list(slice))

slice = islice(iter_list, 2)
print(list(slice))

slice = islice(iter_list, 2)
print(list(slice))

slice = islice(iter_list, 2)
print(list(slice))

slice = islice(iter_list, 2)
print(list(slice))

slice = islice(iter_list, 2)
print(list(slice))

slice = islice(iter_list, 2)
print(list(slice))
>>>
[1, 2]
[3, 4]
[5, 6]
[7]
[]
[]
[]

除了获得切片之外，以上代码还说明了两个非常重要的特征，是否有留意？

第一个：那就是切片能够保留位置信息，多次调用切片功能，当前取值是从上一次结尾的地方开始的。比如第一次取值1、2，结尾位置是3；第二次就从3开始取到了3、4；第三次从5开始取到5、6。原因islice是对迭代器切片，迭代器取值会记住位置信息。

第二个：当迭代完所有的元素之后，返回空数组。将原始列表迭代完之后不会报错，而是一直返回空数组。

有了上面这种使用方法就为分块提供了可能性，如果要使用islice来分块，只需要在一个死循环里调用islice取值，当取值为[]时退出循环即可。可通过如下方法实现：

from itertools import islice

def chunk(it, limit):
    it = iter(it)
    while True:
        temp = list(islice(it, limit))
        if temp == []:
            break
        yield temp

iter_list = iter([1,2,3,4,5,6,7])

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

这样就完成了使用islice就完成了分块的功能，但是看上可不是很优雅，又有while循环，又有yield关键值。

不优雅关键在于需要循环调用切片函数而且还需要判断跳出循环的条件。那么有没有一个既可以循环调用又能判断结束条件的函数呢？还真的有的，那就是iter。

三、iter

iter()方法用来创建迭代器，iter()本质上就是调用可迭代对象的__iter__方法，返回一个迭代器对象。

3.1常规使用

常见的iter的使用方法是，对一个可迭代对象调用iter方法，让其变成一个迭代器，可以通过next取值。

list = [1,2,3,4,5,6,7]
iter_list = iter(list)

print(next(iter_list))
print(next(iter_list))
print(next(iter_list))
>>>
1
2
3

3.2进阶使用

iter还有一种不常用的方法，来看iter函数的定义

iter(object[, sentinel])

• object – 支持迭代的集合对象。
• sentinel – 如果传递了第二个参数，则参数 object 必须是一个可调用的对象（如，函数），此时，iter
  创建了一个迭代器对象，每次调用这个迭代器对象的__next__()方法时，都会调用 object。

也就是说如果iter函数如果传了第二个参数，那么第一个参数就必须是一个可调用对象，每一次调用next函数时，实际上就是调用第一个参数，如果结果等于第二个参数，那就是迭代完成了。

听起来有点弯弯绕，跑一个示例就清楚了。

import random

def get_random():
    return random.randint(1,5)

demo = iter(get_random, 4)

print(next(demo))
print(next(demo))
print(next(demo))
print(next(demo))
print(next(demo))
print(next(demo))
print(next(demo))
>>>
3
2
1
2
Traceback (most recent call last):
  File "islice_demo.py", line 62, in <module>
    print(next(demo))
StopIteration

iter传入第一个参数是一个函数get_random，函数的功能是获取1-5之间的随机数，第二个参数是4，也就是说如果函数返回的数值是4，那算迭代完成。每一次调用next取值就会调用get_random函数，直到结果为4。当迭代完成之后，会抛出一个StopIteration的异常。

上面是通过next调用，如果是通过for循环调用，就不会抛出异常，for循环会捕获异常。

import random

def get_random():
    return random.randint(1,5)

demo = iter(get_random, 4)

for i in demo:
    print(i)
>>>
1
5

这个功能刚好可以实现调用某一个函数，又能判断退出条件，如果现在再把分块的代码摆上来，能否实现优雅的分块呢？

from itertools import islice

def chunk(it, limit):
    it = iter(it)
    while True:
        temp = list(islice(it, limit))
        if temp == []:
            break
        yield temp

iter_list = iter([1,2,3,4,5,6,7])

for temp_list in chunk(iter_list, 2):
    print(temp_list)

四、islice 和 iter 组合使用

islice 提供分块功能，iter 提供循环调用islice的功能和判断退出的功能，最后在两个函数的的配合使用下，完成了优雅的分块。

便于理解的示例：

from itertools import islice

def chunk_list(it, limit):
    it = iter(it)
    # 实现分块的内函数
    def iter_fun():
        return list(islice(it, limit))

    return iter(iter_fun, [])

it = [1,2,3,4,5,6,7]
chunk = chunk_list(it, 2)

print(next(chunk))
print(next(chunk))
print(next(chunk))
print(next(chunk))
print(next(chunk))
>>> 
[1, 2]
[3, 4]
[5, 6]
[7]
Traceback (most recent call last):
  File "chunk_demo.py", line 44, in <module>
    print(next(chunk))
StopIteration

最终的示例：

from itertools import islice

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

iter 第一个参数传入lambda表达式，有一个更贴合场景的叫法是无头函数。 lambda: list(islice(it, limit))。没有传入参数，函数体是islice(it, limit);

第二个参数是空列表[]，作为迭代退出的判断。

工作原理：

当使用for循环遍历分块函数时，每循环一次就通过iter调用islice一次，将分块结果list处理，然后返回。直到islice返回空列表，iter根据第二个参数判断退出循环。