Ryota Kondo

Ryota Kondo

2022/12/17

AWS SDK for pandas (AWS Data Wrangler)でS3にヘッダーのみのcsvファイルを出力するシンプルなコード

問題点

現時点(2022/12)のAWS SDK for pandas(ver 2.18.0)でデータがない状態でcsvをS3へ出力しようとすると下のエラーが発生し、ファイルを出力することはできません。

awswrangler exceptions EmptyDataFrame

エラーが発生するサンプルコード

import pandas as pd
import awswrangler as wr

# データがないDataFrameを仮に作成
df = pd.DataFrame([], columns=["col1", "col2"])

wr.s3.to_csv(
    df,
    path="s3://bucket/folder01/test.csv",
    index=False,
)

通常はデータがあるのでこのエラーが問題となることはありませんが、例えば「処理の結果としてデータが無くなった」ということを出力したcsvから判断したいとなった時には問題が発生します。

この問題を解消するためにはデータが無い場合はヘッダーのみのcsvファイルを出力するようにすればよいので、それを実現するなるべくシンプルなコードを考えました。

データが無い場合はヘッダーのみのcsvファイルを出力するコード

import pandas as pd
import awswrangler as wr

# データがないDataFrameを仮に作成
df = pd.DataFrame([], columns=["col1", "col2"])

params = {
    "path": "s3://bucket/folder01/test.csv",
    "index": False,
}

# S3へcsv出力
if len(df) == 0:
    # データが無い場合はヘッダーのみのcsvを出力する
    wr.s3.to_csv(
        pd.DataFrame([df.columns.values]),
        **params,
        header=False# ポイント
    )

else:
    # データがある場合は通常出力する
    wr.s3.to_csv(
        df,
        **params,
    )

処理の流れとするとデータが無い場合は、もともとのDataFrameのヘッダーをデータとして持つDataFrameを新規作成してヘッダー無しで出力するようにしています。

また、to_csvの共通の引数はparamsでまとめるようにしました。

おわりに

本来であればAWS SDK for pandas側で対応するようにするのがいいのですが、やり方が見つからなかったので自前で実装しました。同じ問題に遭遇した方の助けになれば幸いです。

【2023/12/31 追記】

AWS SDK for pandasのver 3.0.0でデータが無い状態でcsvをS3へ出力したときのエラーが解消されたので、上の考慮は不要になりました。

参考|Release AWS SDK for pandas 3.0.0 · aws/aws-sdk-pandas

関連タグの記事

Ryota Kondo
Ryota Kondo

システムエンジニア・プログラマー|このブログサイトの運営もしており、思いついたことをまとめて記事を書いています💡|Twitterのフォローはお気軽に