大量のデータをDBへ格納する際など、バッチオブジェクトに一定量の書き込み処理を蓄積し、定期的にコミットするという処理を行うことがあります。
上記を行うための、配列を一定間隔で刻みながら処理する方法のPythonのコードによるメモです。
もっと良い方法をご存知でしたらご指摘いただけると幸いです。
最後にGCPのFirestoreへのバッチ書き込みのサンプルコードを載せています。
更新：指摘を頂き「方法４（スライスを利用する）」を追加（@shiracamusさんありがとうございます）
方法１（専用のカウンタを使う）
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']
方法２（インデックスと剰余を使う）
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']
方法３（最後のコミットを工夫する）
最後の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']
方法３＋アルファ
Pythonのenumerateは第二引数で開始番号を指定できるので、方法３はもう少しシンプルにできます。
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']
方法４（スライスを利用する）
オシャレな方法（@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']
参考
方法４で、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()