一個Python中優雅的資料分塊方法

金色旭光發表於2022-05-09
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，需要學習理解的還有很多。

相關文章