🐼
pandasのDataFrameをどうテストするか(pandas.testingの話)
この記事でやったこと
pandasのDataFrame同士が同じかどうかをテストするときに、pandas.testingを使うと便利です。
なぜpandas.testingを使った方がいいのかについて実例とともに紹介します。
背景
pandasのDataFrameを使って、表形式データの加工を行いたいと思いました。
しかし、加工ロジックが複雑になってくると、テストがないと不安になりますよね。
import pandas as pd
from typing import Callable
func_type = Callable[[pd.DataFrame], pd.DataFrame]
def run(df: pd.DataFrame) -> pd.DataFrame:
df = some_process_1(df)
df = some_process_2(df)
df = some_process_3(df)
df = some_process_4(df)
df = some_process_5(df)
return df
def main():
df = pd.read_csv("sample.csv")
result_df = run(df)
# ↑ このresult_dfに正しい処理がされているか不安😢
ただ、pandasのDataFrameは直接==で比較し、Assertすることはできません。
sample_1 = pd.DataFrame([{"a": 1, "b": 2, "c": 3},{"a": 1, "b": 2, "c": 3}])
sample_2 = pd.DataFrame([{"a": 1, "b": 2, "c": 3},{"a": 1, "b": 2, "c": 3}])
assert sample_1 == sample_2
# Traceback (most recent call last):
# File "<stdin>", line 1, in <module>
# assert sample_1 == sample_2
# ^^^^^^^^^^^^^^^^^^^^
# File "/Users/*****/.venv/lib/python3.13/site-packages/pandas/core/generic.py", line 1513, in __bool__
# raise ValueError(
# ...<2 lines>...
# )
#ValueError: The truth value of a DataFrame is ambiguous. Use a.empty, a.bool(), a.item(), a.any() or a.all().
ValueErrorになってしまいますね。
これは、DataFrameに対する==は、bool値を返さずDataFrameを返し、assert は真偽値を要求するため、ValueErrorになってしまいます。
余談: DataFrame同士の == の挙動
ちなみに、DataFrame同士で==を行うと以下のような結果が返ります。便利ですね。
>>> sample_1 = pd.DataFrame([{"a": 1, "b": 2, "c": 3},{"a": 1, "b": 2, "c": 3}])
>>> sample_2 = pd.DataFrame([{"a": 1, "b": 2, "c": 3},{"a": 1, "b": 2, "c": 3}])
>>> sample_1 == sample_2
a b c
0 True True True
1 True True True
では、dictに変換して確認してみましょうか。
sample_1 = pd.DataFrame([{"a": 1, "b": 2, "c": 3},{"a": 1, "b": 2, "c": 3}])
sample_2 = pd.DataFrame([{"a": 1, "b": 2, "c": 3},{"a": 1, "b": 2, "c": 3}])
assert sample_1.to_dict() == sample_2.to_dict()
これでも同じことは確認できました。。
しかしながら、この方法だと、以下のようなTrueと1のように、python上で==で評価されたときにTrueを返してしまうようなケースだと正しく判別できません。
sample_3 = pd.DataFrame({"a": [1, 1], "b": [2, 2], "c": [True, True]})
# 違う↑↓ ^^^^^^^^^^^
sample_4 = pd.DataFrame({"a": [1, 1], "b": [2, 2], "c": [1, 1]})
print("output: ", sample_3.to_dict() == sample_4.to_dict())
# output: True
True と 1で値が異なるのにTrueとされてしまいました。
これはpython自体の仕様のため仕方がない部分もありますが、ユースケースによってはこれらを厳密に分離したいこともあるでしょう。
そこで、今回はpandas.testingを紹介します。
pandas.testing
pandas.testingは、pandasの中に含まれる、pandas.DataFrameや、pandas.Series同士が正しいことなどを確認するためのライブラリです。
公式ドキュメント: https://pandas.pydata.org/docs/reference/testing.html
このpandas.testingに含まれる、assert_frame_equalを用いることによって、今回のような列のデータ型が異なるようなケースを異なるものとしてassertすることができます。
import pandas as pd
sample_3 = pd.DataFrame({"a": [1, 1], "b": [2, 2], "c": [True, True]})
sample_4 = pd.DataFrame({"a": [1, 1], "b": [2, 2], "c": [1, 1]})
pd.testing.assert_frame_equal(sample_3, sample_4)
結果
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
assert_frame_equal(sample_3, sample_4)
~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^
File "/Users/*****/.venv/lib/python3.13/site-packages/pandas/util/_decorators.py", line 220, in wrapper
return func(*args, **kwargs)
File "/Users/*****/.venv/lib/python3.13/site-packages/pandas/_testing/asserters.py", line 1347, in assert_frame_equal
assert_series_equal(
~~~~~~~~~~~~~~~~~~~^
lcol,
^^^^^
...<12 lines>...
check_flags=False,
^^^^^^^^^^^^^^^^^^
)
^
File "/Users/*****/.venv/lib/python3.13/site-packages/pandas/util/_decorators.py", line 193, in wrapper
return func(*args, **kwargs)
File "/Users/*****/.venv/lib/python3.13/site-packages/pandas/_testing/asserters.py", line 1029, in assert_series_equal
assert_attr_equal("dtype", left, right, obj=f"Attributes of {obj}")
~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/*****/.venv/lib/python3.13/site-packages/pandas/_testing/asserters.py", line 448, in assert_attr_equal
raise_assert_detail(obj, msg, left_attr, right_attr)
~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/*****/.venv/lib/python3.13/site-packages/pandas/_testing/asserters.py", line 619, in raise_assert_detail
raise AssertionError(msg)
AssertionError: Attributes of DataFrame.iloc[:, 2] (column name="c") are different
Attribute "dtype" are different
[left]: bool
[right]: int64
assert_frame_equalを使うことで、判別ができるようになり、想定通り異なるDataFrameとして判定してくれています。
さらに、assert_frame_equalを使うことで、エラーメッセージを詳しく出してくれるようになりました。
今回は
AssertionError: Attributes of DataFrame.iloc[:, 2] (column name="c") are different
Attribute "dtype" are different
[left]: bool
[right]: int64
とあることから、cというカラムのデータ型が異なることが一目瞭然です。
他にも、assert_frame_equalには、indexを無視するかどうかや、浮動小数点の誤差をいい感じに扱ってくれたりなどの、便利機能があります。
特に、assert_frame_equalは、最初に挙げたようなデータ型についても検査してくれるという利点以外にも、assertのメッセージが丁寧なのでDataFrame同士でテストを行いたいときには、積極的に使った方がいいですね。
↓ 値が異なる例
>>> sample_5 = pd.DataFrame({"a": [1, 1], "b": [2, 3], "c": [1, 1]})
>>> sample_6 = pd.DataFrame({"a": [1, 1], "b": [2, 1], "c": [1, 1]})
>>> assert_frame_equal(sample_5, sample_6)
結果
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
assert_frame_equal(sample_5, sample_6)
~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^
File "/Users/*****/.venv/lib/python3.13/site-packages/pandas/util/_decorators.py", line 220, in wrapper
return func(*args, **kwargs)
File "/Users/*****/.venv/lib/python3.13/site-packages/pandas/_testing/asserters.py", line 1347, in assert_frame_equal
assert_series_equal(
~~~~~~~~~~~~~~~~~~~^
lcol,
^^^^^
...<12 lines>...
check_flags=False,
^^^^^^^^^^^^^^^^^^
)
^
File "/Users/*****/.venv/lib/python3.13/site-packages/pandas/util/_decorators.py", line 193, in wrapper
return func(*args, **kwargs)
File "/Users/*****/.venv/lib/python3.13/site-packages/pandas/_testing/asserters.py", line 1051, in assert_series_equal
assert_numpy_array_equal(
~~~~~~~~~~~~~~~~~~~~~~~~^
lv,
^^^
...<3 lines>...
index_values=left.index,
^^^^^^^^^^^^^^^^^^^^^^^^
)
^
File "/Users/*****/.venv/lib/python3.13/site-packages/pandas/_testing/asserters.py", line 695, in assert_numpy_array_equal
_raise(left, right, err_msg)
~~~~~~^^^^^^^^^^^^^^^^^^^^^^
File "/Users/*****/.venv/lib/python3.13/site-packages/pandas/_testing/asserters.py", line 689, in _raise
raise_assert_detail(obj, msg, left, right, index_values=index_values)
~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/*****/.venv/lib/python3.13/site-packages/pandas/_testing/asserters.py", line 619, in raise_assert_detail
raise AssertionError(msg)
AssertionError: DataFrame.iloc[:, 1] (column name="b") are different
DataFrame.iloc[:, 1] (column name="b") values are different (50.0 %)
[index]: [0, 1]
[left]: [2, 3]
[right]: [2, 1]
おしまい！
Discussion