バッチコミット作成例(配列を刻みながら処理する方法をいくつか)とFirestoreへのバッチ書き込みのサンプル


大量のデータをDBへ格納する際など、バッチオブジェクトに一定量の書き込み処理を蓄積し、定期的にコミットするという処理を行うことがあります。

上記を行うための、配列を一定間隔で刻みながら処理する方法のPythonのコードによるメモです。
もっと良い方法をご存知でしたらご指摘いただけると幸いです。

最後にGCPのFirestoreへのバッチ書き込みのサンプルコードを載せています。

更新:指摘を頂き「方法4(スライスを利用する)」を追加(@shiracamusさんありがとうございます)

方法1(専用のカウンタを使う)

data = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h']

q = []
batch_size = 3
batch_count = 0
for d in data:
    print("{}".format(d))
    q.append(d)
    batch_count += 1
    if batch_count == batch_size:
        print("commit {}".format(q))
        q = []
        batch_count = 0

print("commit {}".format(q))

> python sample1.py
a
b
c
commit ['a', 'b', 'c']
d
e
f
commit ['d', 'e', 'f']
g
h
commit ['g', 'h']

方法2(インデックスと剰余を使う)

data = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h']

q = []
batch_size = 3
for i, d in enumerate(data):
    print(d)
    q.append(d)
    if (i + 1) % batch_size == 0:
        print("commit {}".format(q))
        q = []

print("commit {}".format(q))

> python sample2.py
a
b
c
commit ['a', 'b', 'c']
d
e
f
commit ['d', 'e', 'f']
g
h
commit ['g', 'h']

方法3(最後のコミットを工夫する)

最後のcommitの記述をなくせてコードの見通しが良い方法です。

予めデータサイズを取得する必要があるのでイテラブルなデータには使えないのと、判定処理の効率が悪いのが弱点です。

data = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h']

q = []
last = len(data)
batch_size = 3
for i, d in enumerate(data):
    print(d)
    q.append(d)
    if ((i + 1) % batch_size == 0) | ((i + 1) == last):
        print("commit {}".format(q))
        q = []

> python sample3.py
a
b
c
commit ['a', 'b', 'c']
d
e
f
commit ['d', 'e', 'f']
g
h
commit ['g', 'h']

方法3+アルファ

Pythonのenumerateは第二引数で開始番号を指定できるので、方法3はもう少しシンプルにできます。

data = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h']

q = []
last = len(data)
batch_size = 3
for n, d in enumerate(data, 1):
    print(d)
    q.append(d)
    if (n % batch_size == 0) | (n == last):
        print("commit {}".format(q))
        q = []

> python sample3.py
a
b
c
commit ['a', 'b', 'c']
d
e
f
commit ['d', 'e', 'f']
g
h
commit ['g', 'h']

方法4(スライスを利用する)

オシャレな方法(@shiracamusさん、ありがとうございます)
バッチオブジェクトのqに配列で渡せない場合はfor inなどでバラす。

data = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h']

batch_size = 3
for i in range(0, len(data), batch_size):
    q = data[i:i+batch_size]
    print("commit", q)

> python sample3.py
commit ['a', 'b', 'c']
commit ['d', 'e', 'f']
commit ['g', 'h']

参考

方法4で、GCPのfirestoreにバッチインサートの上限の500件づつデータを追加するサンプルです。

    db = firestore.Client()
    collection = db.collection("<COLLECTION NAME>")

    batch_size = 500
    batch = db.batch()
    for i in range(0, len(data), batch_size):
        for row in data[i:i + batch_size]:
            batch.set(collection.document(), row)
        print('committing...')
        batch.commit()
        batch = db.batch()