Pandas DataFrame を .csv.gz として Amazon S3 に保存する

Pandas DataFrame を gzip 圧縮しつつ CSV ファイルとして Amazon S3 バケットに保存しようとしたときに少しハマったので備忘録。

import gzip
from io import BytesIO
import pandas as pd
import boto3

def save_to_s3(df: pd.DataFrame, bucket: str, key: str):
    """Pandas DataFrame を .csv.gz として Amazon S3 に保存する"""
    buf = BytesIO()
    with gzip.open(buf, mode="wt") as f:
        df.to_csv(f)
    s3 = boto3.client("s3")
    s3.put_object(Bucket=bucket, Key=key, Body=buf.getvalue())

ポイントとしては以下。

• gzip.open の第一引数は gzip フォーマットを表す file-like オブジェクトなので BytesIO() を入力する
• pandas.DataFrame.to_csv の出力は文字列なので、gzip.open の mode は「テキスト書き込み (wt)」を指定する

最初 pandas.DataFrame.to_csv に compression="gzip" を指定すれば明示的に圧縮しなくてもいけるかと思って試したが、 to_csv に file-like オブジェクトを入力した場合は compression オプションは無視されるらしく、使えなかった。