panderaでUnion型を扱う方法

Python

2024.02.24

pandera (0.18.0) ではUnion型がサポートされていなさげ。そのためちょっとだけ自分でコードを書く必要がある。

例えばあるカラムが文字列or数値であることを検証するには、以下のようにすればよい。

from datetime import datetime

import pandas as pd
import pandera as pa
import pandera.extensions as extensions


@extensions.register_check_method(statistics=[], check_type="element_wise")
def element_wise_check(element):
    return isinstance(element, str) or isinstance(element, float) or isinstance(element, int)

schema = pa.DataFrameSchema({
    "col": pa.Column(object, pa.Check.element_wise_check(), coerce=True)
})

df_ok = pd.DataFrame({"col": [1, 2.0, "x"]})
schema.validate(df_ok)

df_ng = pd.DataFrame({"col": [datetime.now()]})
schema.validate(df_ng)

なお、coerce=True をつけておかないと、例えば pd.DataFrame({"col": [1]}) で int は object じゃないって怒られる。