【pandera】pandasでも型をしっかりつけたい！ #Python

はじめに

こんにちは。今年の春からブレインパッドでデータサイエンティストをしている者です。
ようやく日々の業務に慣れてきたと思っていたら、月日が流れるのは早いもので今年ももうあと2週間ほどですね。来年に向けていろいろと（物理的にも情報的にも）整理をしていこうと思います。

ということで、今年のちょっとした締めくくりとして僭越ながらアドベントカレンダーに参加させていただきました。
BrainPad AdventCalendar 2021の14日目を担当させていただきます。

さてpythonでデータ分析と言えばやはりpandasが主流かと思います。書き方に多少の癖はありますが、DataFrameなどを使うことでデータの読み込み、処理、書き出しまでを非常に幅広くかつ手軽に行うことができます。
一方でその柔軟さゆえに、思った通り処理が行われているか分かりづらいことや、想定していない操作が行われてしまっているなどということがしばしば起こります。また、アドホックに利用されるコードの場合は問題ありませんが、将来的に自分で使い回すコードや他人も利用するコードの場合、DataFrameの中身は実行してみないと分からないためコードが読みにくいというデメリットがあると思います。

それらの改善策の一つとして、ここではpanderaというライブラリを紹介したいと思います。

panderaとは

panderaは「型が正しいか」「数値の正負は想定通りか」「特定文字列を含んでいるか」「nullはないか」などのDataFrameの値のvalidationを行うライブラリです。タイトルでは型を管理すると言ってみましたが、機能としてはもう少し広いです。またpandasだけでなくdask、modinやkoalasなどにも対応しているようです。

panderaにおけるvalidationは

pandera.schemas.DataFrameSchemaで関数として
pandera.model.SchemaModelで主にデコレータとして

使う2通り方法があり、機能としてはどちらもほぼ同じことが実装できます。実際に簡単な例を紹介したいと思います。

DataFrameSchemaによるvalidation

例示のために適当なDataFrameを用意します。

import pandas as pd


df = pd.DataFrame({
    "id": ['id_000', 'id_001', 'id_002'],
    "name": ['Alice', 'Bob', 'Michel'],
    "age": [23, 35, 12],
})

# 中身
#        id    name  age
# 0  id_000   Alice   23
# 1  id_001     Bob   35
# 2  id_002  Michel   12

DataFrameSchemaでは以下のようにしてschemaを定義します。checksの部分がどのような値であるべきかを記述している部分になっています。

import pandera as pa
from pandera import Column, DataFrameSchema, Check, Index


test_schema = DataFrameSchema(
    {
        "id": Column(str, checks=[
            Check.str_startswith("id_"),
            pa.Check(lambda s: s.str.split("_", expand=True).shape[1] == 2)
        ]),
        "name": Column(str),
        "age": Column(int, checks=Check.ge(0)),
    },
    index=Index(int),
    strict=True,
    coerce=True,
)

この例では

idカラムがstring型か
idカラムが'id_'の文字列から開始しているか
idカラムが'_'で分離したときに二つに分割されるか
nameカラムがstring型か
ageカラムがint型か
ageカラムが0以上か

をvalidateすることになります。
あとは

schema.validate(df) # もしくは単にschema(df)でも同じ

とすることで簡単にvalidationを行うことができます。正しく通った場合は基本的に何も出力されません。問題がある場合にエラーやwarningを出力します。また、明示されていませんがnullが含まれているかもチェックしています。例えばidの一つをnullにすると以下のようなエラーが返ってきます。nullを許容したい場合はnullable=Trueを指定します。

SchemaError: non-nullable series 'id' contains null values:

schemaは定義されていないカラムに関してはvalidationを行わないため、記述されていないカラムがあっても通常はチェックされません。定義したカラムのみ持っていることをチェックしたい場合はstrict=Trueとします。実際、dfからageのカラムを削除すると以下のようなエラーが返ってきます。

SchemaError: column 'age' not in DataFrameSchema {'id': <Schema Column(name=id, type=DataType(str))>, 'name': <Schema Column(name=name, type=DataType(str))>}

またcoerce=Trueの場合はチェックを行う前にschemaで定義した型に変換してくれます。ただ強制的に変換させてしまうため、想定外の挙動を防ぐという目的ならばTrueにせずエラーを返してくれた方がいいのではないかと個人的には思っています。

他にもデフォルトで実装されているcheckの例を以下に挙げておきます。(参考: pandera.checks.Check)

メソッド	チェック内容
`eq`	全ての要素が特定の値と等しい
`ge`	全ての要素が特定の値以上
`gt`	全ての要素が特定の値より大きい
`in_range`	全ての要素が特定の範囲内
`isin`	各要素は`isin`で指定した値のいずれか
`le`	全ての要素が特定の値以下
`lt`	全ての要素が特定の値より小さい
`ne`	どの要素も特定の値と等しくない
`notin`	各要素は`notin`で指定した値ではない
`str_contains`	特定のパターンが含まれる
`str_endswith`	特定の文字列で終わる
`str_length`	文字列が指定文字数以内
`str_matches`	指定した正規表現にマッチする
`str_startswith`	特定の文字列で始まる

(表の背景色と等幅フォント化の背景色って同じなんですね...笑)

SchemaModelによるvalidation

上記で見た通りDataFrameSchemaを使うことで非常に簡単にvalidationを行うことができます。
panderaはさらにSchemaModelを使うことでもほぼ同様のことが行えます。先ほどの例と同じことをSchemaModelを用いて実装します。

from pandera.typing import Series
from pandera import SchemaModel, Field


class TestSchema(SchemaModel):
    id: Series[str] = Field(str_startswith='id_', nullable=False, coerce=True)
    name: Series[str] = Field(nullable=False, coerce=True)
    age: Series[int] = Field(ge=0, nullable=False, coerce=True)

    @pa.check("id")
    def id_check(cls, series: Series[str]) -> Series[bool]:
        return series.str.split("_", expand=True).shape[1] == 2

    class Config:
        name = "BaseSchema"
        strict = True

TestSchema.validate(df) # もしくは単にTestSchema(df)でも同じ

SchemaModelの場合はFieldで各カラムのvalidationを定義するような形になり、自作のチェック(id_check)はクラスメソッドとして定義します。また先ほどのstrict=Trueのようにスキーマ全体に渡るようなオプションはConfigサブクラスを使って指定しています。

ここまでだと単に書き方が異なるだけですが、SchemaModelにはさらにデコレータを使えるメリットがあります。
例えば下記のようなtest_fnが定義されている場合type hintsにSchemaModelを記述することによって、入出力のvalidationを行うことができます。このvalidationはデコレータを付けた関数の実行時に行われます。

from pandera.typing import DataFrame


# 引数のdfと返り値がTestSchemaでvalidationされる
@pa.check_types
def test_fn(df: DataFrame[TestSchema]) -> DataFrame[TestSchema]:
    df['age'] = df['age'] + 1
    return df

デコレータの記述方法はいくつかあり、例えば

@pa.check_input(TestSchema)
@pa.check_output(TestSchema)
@pa.check_io(df_=TestSchema, out=TestSchema)

のようにvalidationするinputやoutputを明示することも可能です。

所感

基本的にシンプルで使いやすいと感じました。panderaの公式ドキュメントもある程度丁寧にまとまっているため、実装したいことを見つけ出すのはそれほど苦ではない印象です。紹介した通り2種類の実装がありますが、基本的にvalidationを行いたいタイミングは関数の入出力時であることや、type hintsが充実してきている昨今のpythonのアップデートを鑑みてもSchemaModelを使ったvalidationの方が扱いやすいのではないかと個人的には思っています。

また値のvalidationができるということ自体ももちろん有用ですが、個人的にはスクリプトをだけを読んでDataFrameが何を保持しているのかが分かることも大きなメリットだと思います。（もちろんコーディング量は増えますが...）
例えば下記のようなtest_fn_2だけが書かれていた場合、joinしているdf_masterから何の情報を付与したのか分かりません。しかしTestOutputSchemaが記述されていることで性別カラムを付け加えたことがよくわかります。

# TestSchemaを継承
class TestOutputSchema(TestSchema):
    gender: Series[str]


@pa.check_types
def test_fn_2(df: TestSchema, df_master) -> TestOutputSchema:
    df_joined = df.join(df_master, on='id', how='left')
    return df_joined

もちろん例示した内容に限らずより複雑なvalidationもできます。MultiIndexのDataFrameに対するvalidationも可能であったり、t検定などの簡単な手法に限定されますが統計的仮説検定も行うことができるようです。
ただ、複雑なことをしすぎてしまうとそれはそれで実装が大変だったり内容が分かりにくくなってしまったりと本末転倒になりかねないため、なるべくシンプルに使う方が効果的だとは思います。

結び

最後までお読みいただきありがとうございました。非常に簡単ではありましたがpanderaの使い方を紹介させていただきました。この記事を書いている時点ではあまり日本語で紹介されているページが見つからなかったため、必要としている方のお役に立てれば幸いです。

それでは残りのアドベントカレンダーもぜひお楽しみください。