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


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

ということで、今年のちょっとした締めくくりとして僭越ながらアドベントカレンダーに参加させていただきました。
[BrainPad AdventCalendar 2021](https://qiita.com/advent-calendar/2021/brainpad)の14日目を担当させていただきます。


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

それらの改善策の一つとして、ここでは[**pandera**](https://pandera.readthedocs.io/en/stable/)というライブラリを紹介したいと思います。

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

panderaにおけるvalidationは

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

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

## DataFrameSchemaによるvalidation
例示のために適当なDataFrameを用意します。

```python
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`の部分がどのような値であるべきかを記述している部分になっています。

```python
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することになります。
あとは

```python
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`)](https://pandera.readthedocs.io/en/latest/reference/generated/pandera.checks.Check.html)

| メソッド         | チェック内容                         |
| ---------------- | ------------------------------------ |
| `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`を用いて実装します。

```python
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はデコレータを付けた関数の実行時に行われます。

```python
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`が記述されていることで性別カラムを付け加えたことがよくわかります。

```python
# 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の使い方を紹介させていただきました。この記事を書いている時点ではあまり日本語で紹介されているページが見つからなかったため、必要としている方のお役に立てれば幸いです。

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

## 参考
- [公式ドキュメント](https://pandera.readthedocs.io/en/latest/index.html)
- [panderaでDataFrameをバリデーションする](https://ohke.hateblo.jp/entry/2020/09/12/230000)