巨大なjsonファイルをサイズ基準で分割するpythonスクリプト

状況

• jsonファイルをRedshift serverlessのtableに書き込みたい
• RedshiftのCOPY文を実行して書きこみたいが、書き込み時はSUPER型の列に書き込むため、jsonのサイズは最大1MB以下に収める必要がある。
• 書き込み予定のjsonファイルの生データのサイズは150MB程度である。
• jsonファイルは10,000要素の配列。

やりたいこと

jsonファイルを適度に分割し、redshiftへの書き込みができる程度に小さくする。

スクリプト

json配列を先頭から読み、二分探索によりある上限サイズ未満に収まるように区切る。

import ndjson
from sys import getsizeof

def split_json(data: list[dict], max_size: int) -> list[str]:
    """
    json配列を指定のサイズを超えないサイズに区切る
    :param data: json配列
    :param max_size: 一要素の許容最大サイズ
    :return: 区切ったjsonをndjson形式でdumpした値の配列
    """

    length = len(data)

    output = []

    si = 0  # 区間の始点

    # 始点が最終要素になるまで繰り返す
    while si < length - 1:
        # 区切り位置を二分探索で決定する
        li = bisect_size(data, si, max_size)

        # 一要素のサイズがmax_sizeを超える場合は、si==liとなる
        # この場合はmax_sizeを超えた一要素のみを含むndjsonをoutputに追加する
        if si == li:
            li += 1

        output.append(ndjson.dumps(data[si:li]))

        si = li

    return output


def bisect_size(array: list, s: int, max_size: int) -> int:
    """
    配列の分割点を二分探索で決定する
    :param array: 配列
    :param s: 配列の始点
    :param max_size: 許容最大サイズ(Byte)
    :return:
    """
    left = s
    right = len(array)

    while right - left > 1:
        tmp = (left + right) // 2
        tmp_json = ndjson.dumps(array[s:tmp])

        if getsizeof(tmp_json) < max_size:
            left = tmp
        else:
            right = tmp

    return left

最大4MBで区切りたい場合は、max_size=1_000_000とする形で呼ぶ。

自分の場合は、一要素で許容上限サイズを超えてしまうものがありましたが、適当に対応しました。