S3上の数GBファイルをstreamingでZIP化する【smart-open】

TL;DR

• boto3のStreamingBody.read()を使うと、S3オブジェクト全体をメモリへ展開してしまう
• boto3のStreamingBody.iter_chunks()を使うことで、chunk単位で読み込める
• smart_openを使うことで、ZIPの出力先をそのままS3へ向けられる
• メモリへ全展開せず、streamingしながらZIP化できる

課題

S3上にある数GBのファイルをZIP化したいケース、ありますよね

例えば...

SELECT INTO OUTFILE S3でDBから直接S3にCSV出力した！
けど、ZIP化に対応していない、、
ユーザにダウンロードさせるにはZIPにしておきたい、、

ただZIP化したいだけだったんですが、相手が数GBのファイルで苦戦したって話です

NGパターン

まず、ふつーにZIP化しようとするとこんな感じになると思います

import io
import boto3
import zipfile

s3_client = boto3.client("s3")

s3_object = s3_client.get_object(
    Bucket=source_bucket,
    Key=source_key
)

s3_object_bytes = s3_object["Body"].read()

zip_buffer = io.BytesIO()

with zipfile.ZipFile(zip_buffer, "w") as zip_file:
    zip_file.writestr("large-file.csv", s3_object_bytes)

zip_binary = zip_buffer.getvalue()

s3_client.put_object(
    Bucket=dest_bucket,
    Key="output.zip",
    Body=zip_binary
)

一見シンプルですが、相手のファイルサイズによっては問題があります

s3_object["Body"].read()

の時点で、S3オブジェクト全体をメモリへ展開しています

さらに、

zip_buffer = io.BytesIO()

にもZIP全体が載るため、

元ファイル + ZIPファイル

の両方をメモリ保持することになります

数GBのファイル相手だと厳しい戦いとなります

実装

import boto3
import zipfile
import smart_open

CHUNK_SIZE = 8 * 1024 * 1024

s3_client = boto3.client("s3")

s3_object = s3_client.get_object(
    Bucket=source_bucket,
    Key=source_key
)

with smart_open.open(
    "s3://dest-bucket/output.zip",
    "wb"
) as s3_output:

    with zipfile.ZipFile(
        s3_output,
        mode="w",
        compression=zipfile.ZIP_DEFLATED,
        allowZip64=True,
    ) as zip_file:

        with zip_file.open(
            "large-file.csv",
            "w",
            force_zip64=True,
        ) as zip_entry:

            for chunk in s3_object["Body"].iter_chunks(
                chunk_size=CHUNK_SIZE
            ):
                zip_entry.write(chunk)

それぞれの役割

今回の処理は以下のような役割分担になっています

[boto3]
S3からchunk単位で読む
        ↓
[zipfile]
ZIP圧縮する
        ↓
[smart_open]
S3へstreaming uploadする

boto3

response["Body"].iter_chunks()

を使うことで、S3オブジェクトをchunk単位で読み込めます

つまり、

8MB読む
↓
zipへ書く
↓
次の8MB読む

を繰り返すことで、ファイル全体をメモリへ載せず処理が可能

smart_open

今回の救世主はsmart_openでした

smart_open.open(
    "s3://bucket/output.zip",
    "wb"
)

smart_openを使うと、S3上のオブジェクトを書き込み可能な file-like object として扱えます

つまり、ZIPファイルを一度 BytesIO にすべて溜めてから put_object() するのではなく、

zipfile が書き込む
↓
smart_open が受け取る
↓
S3へ順次書き込む

という流れにできます

そのため、巨大ファイルでもメモリ使用量を chunk サイズ程度に抑えながらZIP化できます

救世主はこちら↓

まとめ

S3上の巨大ファイルをZIP化する際、単純に .read() してしまうとS3オブジェクト全体をメモリへ展開してしまい、メモリ不足の原因になります

今回は、

• boto3 の StreamingBody.iter_chunks()
• smart_open
• zipfile

を組み合わせることで、

S3 → streaming read → ZIP化 → streaming upload → S3

という形で、メモリに全展開せずZIP化できました
使用メモリもほぼチャンクサイズ程度で抑えられるため、もう巨大ファイルも怖くありません

【参考】