pandas.DataFrame内のlist型を維持してファイルに吐き出したい

MLモデル開発中にデータをファイルで保存したい

MLモデルの開発では、前処理や特徴量エンジニアリングなどの工程があり、
前処理を終えたデータを一時的にそのままの形でファイルに吐き出したい時などがある。
通常の場合、pandasのto_csv 関数でcsv出力すれば問題はない。しかし、DataFrame内のlist型が存在する時に型を保存できなくて困り、調べた。
(list型をDataFrameに入れることはあまりないが、複数選択可の項目で処理のために配列にしたい時があった。)

僕の場合の結論は、parquet で保存することが一番好ましかった。
調べたら色々と方法があったので、それぞれの方法をまとめる。

DataFrameの準備

今回はDataFrame内のlist型の挙動を確認したいので、list型のデータを作る。

.py

import pandas as pd


# ファイルの保存先のパス
file_path = './filepath'
# S3用
bucket = 'bucket_name'
prefix = 'prefix_path'

df = pd.DataFrame(data={'int_col': [1, 2], 'list_col': [[1, 2], [3, 4]]})

print(type(df['list_col'][0])) # <class 'list'>

csvで読み書き

DataFrameをcsv出力させ、そのファイルを改めて読み込むとlist型はstring型になってしまう。

.py

# 書き込み
df.to_csv(f'{file_path}/sample.csv.gz', index=False, compression='gzip')
# 読み込み
df = pd.read_csv(f'{file_path}/sample.csv.gz')
# 型確認
print(type(df['list_col'][0])) # <class 'str'>

以下のようにすることで、list型に変換できる。

.py

df['list_col'] = df['list_col'].apply(eval)
print(type(df['list_col'][0])) # <class 'list'>

pickleで読み書き

pickle は python に標準的に組み込まれており、オブジェクトをファイルとして保存することができる。
DataFrame内のlist型もしっかり扱うことができる。

.py

# 書き込み
df.to_pickle(f'{file_path}/sample.pickle.gz', compression='gzip')
# 読み込み
df = pd.read_pickle(f'{file_path}/sample.pickle.gz')
# 型確認
print(type(df['list_col'][0])) # <class 'list'>

pythonのオブジェクトを読み込む必要があるので、pythonでしか扱えないのがネック。

parquetで読み書き

こちらが今回の本命。
恥ずかしながら今回調べるまでは「Athenaで扱うときのフォーマットでしょ？」くらいの認識だった。
Column-orientedフォーマットというものにカテゴリ分けされるらしく、 列方向にデータが格納されている。
そのため、圧縮アルゴリズムも効きやすく、列を指定しての読み込みもできてしまう。

pythonで扱うためには、fastparquet か pyarrow というライブラリをインストールする必要がある。どちらも pip で簡単にインストールできる。
今回は fastparquet をインストールして試してみる。

.py

# 書き込み
df.to_parquet('sample.parquet.gz', compression='gzip')
# 読み込み(columnsでカラムを指定できる)
df = pd.read_parquet('sample.parquet.gz', columns=['list_col'])
# 型確認
print(type(df['list_col'][0])) # <class 'list'>

pickleと同様にlist型も保存されている。

S3に読み書き

fileをS3に置く場合も調べたので書いておく。
僕の調べた限りだと to_parquet 関数では送れなさそうだった。
以下のようにすればできた。

.py

# 書き込み
import s3fs
from fastparquet import write

s3 = s3fs.S3FileSystem()
myopen = s3.open

write(
    f's3://{bucket}/{prefix}/sample.parquet.gz', 
    df, 
    compression='GZIP',
    open_with=myopen)

.py

# 読み込み
import s3fs
from fastparquet import ParquetFile

s3 = s3fs.S3FileSystem()
myopen = s3.open

pf = ParquetFile(
    f's3://{bucket}/{prefix}/sample.parquet.gz',
    open_with=myopen)
df = pf.to_pandas()

まとめ

csv(型が保存できない) → pickle(pythonでしか使えないし...) → Parquetという話だったのだが、
色々調べていたら列志向とデータ分析の相性の良さを改めて感じさせられた。
csvでもpickleでもparquetでも、扱うことはできるので使い分けが必要。

個人的には、データをpython(pandas)で扱う際に、カラムを指定して読み込みたい時は多々あるので、
今後はParquetフォーマットを使っていきたいと思う。

参考

https://blog.amedama.jp/entry/2017/10/10/135331
http://nagix.hatenablog.com/entry/2015/12/08/235512
https://note.nkmk.me/python-pandas-to-pickle-read-pickle/
https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.read_parquet.html