AWS SDK for pandas (AWS Data Wrangler)でS3にヘッダーのみのcsvファイルを出力するシンプルなコード
問題点
現時点(2022/12)のAWS SDK for pandas(ver 2.18.0)でデータがない状態でcsvをS3へ出力しようとすると下のエラーが発生し、ファイルを出力することはできません。
エラーが発生するサンプルコード
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へ出力したときのエラーが解消されたので、上の考慮は不要になりました。
関連タグの記事
