pandas.read_parquet()でpandas.Intervalクラスが読めたり読めなかったりする
pd.read_parquet()ってfastparquetとpyarrowのときで読み込まれるデータフレームが異なる場合があるのか。fastparquetだとInterval型が読めないので勝手にhoge.leftとhoge.rightって列に展開される。
— Kiichi (@Ki_chi) June 30, 2022
読み込みエンジンを変えると挙動が変わるっぽいという話です。
実際に下記のコードを試してみましょう。動作確認用Dockerfileからビルドしたdockerコンテナ上で確認しています。
各ソフトウェアのバージョンは下記の通り:
- Python 3.10.5
- pandas 1.4.3
- pyarrow 8.0.0
- fastparquet 0.8.1
まずは下準備として、pandas.Interval クラスのオブジェクトを値として持つデータフレームを用意し、外部ファイルとして保存します:
import pandas as pd
df = pd.DataFrame({"a": [0,1,2], "b": [pd.Interval(0,1), pd.Interval(1,2), pd.Interval(2,3)]})
print(df)
# a b
# 0 0 (0, 1]
# 1 1 (1, 2]
# 2 2 (2, 3]
df.info()
# <class 'pandas.core.frame.DataFrame'>
# RangeIndex: 3 entries, 0 to 2
# Data columns (total 2 columns):
# # Column Non-Null Count Dtype
# --- ------ -------------- -----
# 0 a 3 non-null int64
# 1 b 3 non-null interval[int64, right]
# dtypes: int64(1), interval(1)
# memory usage: 200.0 bytes
df.to_parquet("./temp.parquet")
上記で保存した temp.parquet をエンジン指定してそれぞれ読み込みます。まずはpyarrowを指定して読んでみます。
df_pyarrow = pd.read_parquet("./temp.parquet", engine="pyarrow")
print(df_pyarrow)
# a b
# 0 0 (0, 1]
# 1 1 (1, 2]
# 2 2 (2, 3]
df_pyarrow.info()
# <class 'pandas.core.frame.DataFrame'>
# RangeIndex: 3 entries, 0 to 2
# Data columns (total 2 columns):
# # Column Non-Null Count Dtype
# --- ------ -------------- -----
# 0 a 3 non-null int64
# 1 b 3 non-null interval[int64, right]
# dtypes: int64(1), interval(1)
# memory usage: 200.0 bytes
無事、保存したデータが想定通り読めています。次にfastparquetでも読んでみましょう。
df_fastparquet = pd.read_parquet("./temp.parquet", engine="fastparquet")
print(df_fastparquet)
# a b.left b.right
# 0 0 0 1
# 1 1 1 2
# 2 2 2 3
df_fastparquet.info()
# <class 'pandas.core.frame.DataFrame'>
# RangeIndex: 3 entries, 0 to 2
# Data columns (total 3 columns):
# # Column Non-Null Count Dtype
# --- ------ -------------- -----
# 0 a 3 non-null int64
# 1 b.left 3 non-null int64
# 2 b.right 3 non-null int64
# dtypes: int64(3)
# memory usage: 200.0 bytes
pyarrowを使ったときと異なり、保存時と違う形式のデータフレームが読み込まれているようです。 pd.Interval オブジェクトだったはずの"b"列が、b.left, b.rightの2つの列に分裂してしまいました。
どうやら、pyarrowのみが拡張タイプとしてpandas.Intervalクラスをサポートしているために起きるようです。（公式マニュアルによればpyarrowではシリアライズ/デシリアライズをユーザー定義してあげることで、Pythonクラスをpyarrowで扱えるようにできるとのこと。pandas.Intervalクラスはすでにライブラリ側で定義されているようです。）
fastparquetでもメタデータ内にb列が interval[int64, right] であることは読めるんですが、これを pandas.Interval クラスとして変換することができないようです。 逆の変換も然りで、最初にParquetファイルを書き出す時点でfastparquetを指定すると下記のようなエラーとなります。
df.to_parquet("./temp.parquet", engine="fastparquet")
# Traceback (most recent call last):
# File "<stdin>", line 1, in <module>
# File "/usr/local/lib/python3.10/site-packages/pandas/util/_decorators.py", line 207, in wrapper
# return func(*args, **kwargs)
# File "/usr/local/lib/python3.10/site-packages/pandas/core/frame.py", line 2835, in to_parquet
# return to_parquet(
# File "/usr/local/lib/python3.10/site-packages/pandas/io/parquet.py", line 420, in to_parquet
# impl.write(
# File "/usr/local/lib/python3.10/site-packages/pandas/io/parquet.py", line 301, in write
# self.api.write(
# File "/usr/local/lib/python3.10/site-packages/fastparquet/writer.py", line 1214, in write
# fmd = make_metadata(data, has_nulls=has_nulls, ignore_columns=ignore,
# File "/usr/local/lib/python3.10/site-packages/fastparquet/writer.py", line 824, in make_metadata
# se, type = find_type(data[column], fixed_text=fixed,
# File "/usr/local/lib/python3.10/site-packages/fastparquet/writer.py", line 225, in find_type
# raise ValueError("Don't know how to convert data type: %s" % dtype)
# ValueError: Don't know how to convert data type: interval[int64, right]df.to_parquet("./temp.parquet", engine="fastparquet")
# Traceback (most recent call last):
# File "<stdin>", line 1, in <module>
# File "/usr/local/lib/python3.10/site-packages/pandas/util/_decorators.py", line 207, in wrapper
# return func(*args, **kwargs)
# File "/usr/local/lib/python3.10/site-packages/pandas/core/frame.py", line 2835, in to_parquet
# return to_parquet(
# File "/usr/local/lib/python3.10/site-packages/pandas/io/parquet.py", line 420, in to_parquet
# impl.write(
# File "/usr/local/lib/python3.10/site-packages/pandas/io/parquet.py", line 301, in write
# self.api.write(
# File "/usr/local/lib/python3.10/site-packages/fastparquet/writer.py", line 1214, in write
# fmd = make_metadata(data, has_nulls=has_nulls, ignore_columns=ignore,
# File "/usr/local/lib/python3.10/site-packages/fastparquet/writer.py", line 824, in make_metadata
# se, type = find_type(data[column], fixed_text=fixed,
# File "/usr/local/lib/python3.10/site-packages/fastparquet/writer.py", line 225, in find_type
# raise ValueError("Don't know how to convert data type: %s" % dtype)
# ValueError: Don't know how to convert data type: interval[int64, right]
ちなみに、冒頭のdf.to_parquet("./temp.parquet")というコードのようにエンジンを指定しない場合はpyarrow, fastparquetの順で有効なエンジンが自動で選択されます。上記の例では裏でpyarrowが呼ばれていたため、pandas.Intervalクラスも難なく保存できたというわけですね。
ParquetフォーマットはCSVに比べて容量も小さくなるし、大量のデータについて読み書きも速いのでつい気軽に使ってしまいがちですが、必ずしも中身がデシリアライズできるとは限りません。 異なる環境をまたぐ可能性があるときは、可能な限り各言語での実装状況を参考に、Apache Arrowでサポートされているような型を使うのがよいかと思います。