概要

pandas.DataFrame.to_parquet() に関して、下記2点について調べたメモを残します。

出力したparquetファイルのschemaの確認方法
出力時に明示的にschemaを指定する方法

前提事項

to_parquet() の engine としてデフォルトの pyarrow を使うことを前提としています。

補足

pandas.DataFrameをparquet形式で出力する方法の1つとして、pd.DataFrameのメソッド to_parquet() があります。to_csv()など他の出力系メソッドと同様に使えて便利ですね。

to_parquet()する時にデータに対してpyarrowのschemaが定義されますが、pandasのドキュメントにはその確認方法や指定方法は記載されていません(engine側の機能なので)。

明示的にschemaを指定せずにto_parquet()で出力されたデータを読み込む際に、schemaが原因でエラーが発生したことがありました。
schemaの指定方法を調べましたが、情報がなかなか見つからなかったので今後のために忘備メモです。

実行環境

OS: Windows 10
Python: 3.9.7
pandas: 1.3.3
pyarrow: 5.0.0

基本の確認

pandas.DataFrame は Arrow で言うと Table に相当する

Apache Arrow ドキュメントのここ

The equivalent to a pandas DataFrame in Arrow is a Table. While pandas only supports flat columns, the Table also provides nested columns, thus it can represent more data than a DataFrame, so a full conversion is not always possible.

ネスト構造に関して注意点はありますが、Pandas.DataFrame は ArrowのTable に相当するのですね。

Arrow のテーブルの各列のデータ型は「schema」で定義される

Apache Arrowドキュメントのここ

it defines the column names and types in a record batch or table data structure.

schema は（record batch や）テーブルの列名と型を定義するものです。

parquetファイルのschema確認方法

まずはparquetファイルのschemaの確認方法です。

Apache Arrowドキュメントのここ

下記の例では、test.parquetのAという列がDicimal(3, 3)、Bという列がdoubleであることが分かります。

import pyarrow.parquet as pq

parquet_file = pq.ParquetFile("files/test.parquet")
print(parquet_file.schema) 

#> <pyarrow._parquet.ParquetSchema object at 0x000002B7B2188180>
#> required group field_id=-1 schema {
#>  optional fixed_len_byte_array(2) field_id=-1 A (Decimal(precision=3, scale=3));
#>  optional double field_id=-1 B;
#> }

read_metadata()を使った下記の方法でも同様に取得できます。

print(pq.read_metadata("files/test.parquet").schema)

出力時に明示的にschemaを指定する方法

to_parquet()のドキュメントでは、パラメータにschemaって記載されていないのですよね。（適用するengineに依存する部分なので）

結論を先に記載すると、schemaを定義して、to_parquet()に定義したschemaを渡してあげればOKです（**kwargs の一部となりますね）

テスト用のデータ

dfはpandas.DataFrameです。
列Aは、少数第3位に丸めたDecimalです。

print(type(df))
#> <class 'pandas.core.frame.DataFrame'>

print(df)
#>         A         B
#> 0  0.317  0.849559
#> 1  0.365  0.169753
#> 2  0.109  0.601026
#> 3  0.527  0.175449
#> 4  0.614  0.368311

print(df.dtypes)
#> A     object
#> B    float64
#> dtype: object

print(df["A"].map(lambda v: type(v)).unique())
#> [<class 'decimal.Decimal'>]

schemaを指定せずにto_parquet()

まず、schemaを指定せずにto_parquet()してみます。

df.to_parquet("files/to_parquet_auto.parquet", index=False)
print(pq.read_metadata("files/to_parquet_auto.parquet").schema)
#> <pyarrow._parquet.ParquetSchema object at 0x0000027AC80A9F40>
#> required group field_id=-1 schema {
#>   optional fixed_len_byte_array(2) field_id=-1 A (Decimal(precision=3, scale=3));
#>   optional double field_id=-1 B;
#> }

列AがDecimal(3, 3)、列Bがdouble として出力されました。
データ内容から判断してschemaが設定されるのですね。

schemaを指定してto_parquet()

次に、schemaを指定してto_parquet()してみましょう。
pyarrowのデータ型やschemaの詳細は、下記リンクを参照してください。

schema = pa.schema([
    pa.field("A", pa.decimal128(precision=4, scale=3)),
    pa.field("B", pa.float64())
])
df.to_parquet("files/to_parquet_schema.parquet", index=False, schema=schema)
print(pq.read_metadata("files/to_parquet_schema.parquet").schema)
#> <pyarrow._parquet.ParquetSchema object at 0x0000014BCA2B84C0>
#> required group field_id=-1 schema {
#>   optional fixed_len_byte_array(2) field_id=-1 A (Decimal(precision=4, scale=3));
#>   optional double field_id=-1 B;
#> }

fieldで列名とpyarrowのデータ型を指定しています。
列AにDecimal(4, 3), 列Bに元と同じfloat64()を指定するschemaを作成し、to_parquet()に渡しました。
最後にread_metadata()で出力したparquetファイルにschemaが反映されていることを確認できました。

データ依存でschemaが変わるケース

たとえば確率値は値の範囲が0以上1以下となりますが、1ってレアでデータに含まれないことも良くあると思います。
小数第3位に丸めたDecimalの場合、0.000～0.999は有効桁数3, 小数部桁数3（＝Decimal(3, 3)）で良いですが、1.000は有効桁数4が必要です。
to_parquet()するときにschema指定しないと、データから判断されるので

確率値1.000がない時に出力したparquetファイルのschemaでは、列AがDecimal(3, 3)
確率値1.000がある時に出力したparquetファイルのschemaでは、列AがDecimal(4, 3)

となります。

schemaがpyarrowに渡されていく流れのメモ

pandas.DataFrame.to_parquet() にschemaを渡す（**kwargsの一部として）
pandas.io.parquet.to_parquet() でimpl.write() にschemaが渡される（**kwargsの一部として）
pandas.io.parquet.PyArrowImpl.write() で
- **kwargsからschemaが拾われて（kwargs.pop("schema", None)）
- from_pandas_kwargs（辞書型）に格納され
- pyarrow.Table.from_pandas()にfrom_pandas_kwargsごと渡される。
  - pyarrow.Table.from_pandas()はpandas.DataFrameをpyarrow.Tableに変換するメソッド

※ schemaを指定しなかった場合は、pyarrowでfrom_pandas()に呼び出されるpandas_conpat.py > dataframe_to_arrays()関数内でschemaが生成される

[to_parquet] schemaの確認と指定方法

概要