一個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 步長

示例

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取到的數據。

只指定步長

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__方法，返回一個迭代器對象。關於iter的常規使用，可參見另一篇文章一篇文章講清楚迭代器和生成器

常規使用

常見的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

進階使用

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根據第二個參數判斷退出循環。

總結

分塊函數的優點：

1. 實現很優雅
2. 支持的分塊的數據類型豐富。不單是列表，只要能夠迭代的都可以。

分塊的實現主要有兩個思路：

1. 使用islice來完成迭代器切片，實現分塊的功能。但是需要多次調用islice直到迭代完成
2. iter 提供調用功能，並判斷迭代退出條件

有興趣的讀者可看看iter的實現，能夠明白爲什麼迭代器能記住位置，這是本文分塊的一個核心知識點。

這一個簡單的代碼讓我感受到Python的奇妙，兩個函數默契的配合，十分優雅的完成了分塊功能。同時我明白Python語言的宗旨是簡易優雅，但是簡易並不簡單，想要實現優雅需要紮實的基礎和深厚的知識儲備。追求Pythonic，需要學習理解的還有很多。