【Pandas】2つのDataFrameが一致していることをテストする（assert_frame_equel）
DataFrameを返す関数のテストを書く時に、期待されるDataFrameと返り値のDataFrameをどう比較したものか頭を悩ませたことはありませんか？
私は悩んだ結果、for文で１つ１つの要素を比較するというなんとも面倒なことをした経験があります。
↓ こんな感じ
# a_df と b_df の各要素が一致することをテスト for index, row in a_df.iterrows(): for column in a_df.columns: assert row[column] == b_df.at[index, column]
その後、2つのDataFrameを良い感じに比較してくれる assert_frame_equel を偶然見つけてからはだいぶテストを書くのが楽になりました。
この記事ではその assert_frame_equel を紹介していきます。
- データ準備
- assert_frame_equelの使い方
- [chack_dtype=False] 型の比較をしない
- [check_like=True] インデックス・カラムの順序を無視する
- その他のオプション
- Index, Seriesを比較する
動作環境 MacOS Catalina 10.15.7 Python 3.8.6 pandas 1.2.1 VSCode 1.52.1
VSCode拡張機能のJupyterを使用します。 詳細はこちらを参照ください。
データ準備
他の記事で良く使っているデータをDataFrameにして使おうと思います。
2020年のa店舗、b店舗、c店舗の日次来客数のデータです。
import pandas as pd a_df = pd.read_csv("./csv/store_visits_2020.csv")
assert_frame_equelの使い方
from pandas.testing import assert_frame_equal b_df = a_df.copy() assert_frame_equal(a_df, b_df)
一致している場合 → None
print(assert_frame_equal(a_df, b_df)) > None
一致していない場合 → 例外が発生
# 列を増やす b_df = a_df.copy() b_df["hoo"] = "bar" assert_frame_equal(a_df, b_df)
--------------------------------------------------------------------------- AssertionError Traceback (most recent call last) <ipython-input-11-e8e4f3008e7b> in <module> 2 b_df["hoo"] = "bar" 3 ----> 4 assert_frame_equal(a_df, b_df) [... skipping hidden 1 frame] ~/.pyenv/versions/3.8.2/lib/python3.8/site-packages/pandas/_testing.py in raise_assert_detail(obj, message, left, right, diff, index_values) 1071 msg += f"\n[diff]: {diff}" 1072 -> 1073 raise AssertionError(msg) 1074 1075 AssertionError: DataFrame are different DataFrame shape mismatch [left]: (1098, 3) [right]: (1098, 4)
どこが違うのか教えてくれるので、デバッグにも便利です。
上のエラーではDataFrameのshapeが違うと言ってますね。
a_df、上でいう [left] は 1098行3列に対して、
b_df、上でいう [right] は1098行4列
left, rightは引数で与えられたDataFrameのうちどちらに該当しているかを示しています。
assert_frame_equelはTrue, Falseを返すものではないので、if文などで使うことはできません。
基本的にテストで使うものだと思われます。
また、さすがにこの間違え方をするのは私くらいだと思いますが、
# × assert assert_frame_equal(a_df, b_df) is None # ○ assert_frame_equal(a_df, b_df)
です。
（まあ上でも問題はないと思いますが...）
[chack_dtype=False] 型の比較をしない
デフォルトはTrue
# 型を変える b_df = a_df.copy() b_df["visit_num"] = b_df["visit_num"].astype(float) assert_frame_equal(a_df, b_df)
エラーが発生します
AssertionError: Attributes of DataFrame.iloc[:, 2] (column name="visit_num") are different Attribute "dtype" are different [left]: int64 [right]: float64
型の違いまでチェックしなくて良い場合、 check_dtype=False を設定してあげましょう。
ただし、日付の型で同じことをやりたい場合、 check_datetimelike_compat の指定をすることでうまくできるかもしれません。
# 型を変える b_df = a_df.copy() b_df["visit_num"] = b_df["visit_num"].astype(float) assert_frame_equal(a_df, b_df, check_dtype=False)
[check_like=True] インデックス・カラムの順序を無視する
デフォルトはFalse
# カラムの順番を入れ替える b_df = a_df.copy() b_df = b_df[["date", "visit_num", "store_id"]] assert_frame_equal(a_df, b_df)
カラムの順番が違うと怒られます
AssertionError: DataFrame.columns are different DataFrame.columns values are different (66.66667 %) [left]: Index(['date', 'store_id', 'visit_num'], dtype='object') [right]: Index(['date', 'visit_num', 'store_id'], dtype='object')
カラムの順番を合わせてから比較することもできますが、
check_like=Trueを指定することで順序を無視することが可能です。
# カラムの順番を入れ替える b_df = a_df.copy() b_df = b_df[["date", "visit_num", "store_id"]] assert_frame_equal(a_df, b_df, check_like=True)
インデックスの場合も同様です。
# インデックスの順番を入れ替える b_df = a_df.copy() b_df = b_df.sort_values("visit_num") assert_frame_equal(a_df, b_df, check_like=True)
ただし、インデックスを振り直してしまうとうまくいかないので注意が必要です。
# インデックスの順番を入れ替える b_df = a_df.copy() b_df = ( b_df.sort_values("visit_num") .reset_index(drop=True) ) assert_frame_equal(a_df, b_df)
AssertionError: DataFrame.iloc[:, 0] (column name="date") are different DataFrame.iloc[:, 0] (column name="date") values are different (99.6357 %) 以下略
その他のオプション
他のオプションに関しては、使ったことがない or 使い方がわからないので、公式のリファレンスを翻訳したのを載せておきます。
Index, Seriesを比較する
Indexを比較する場合、assert_index_equel
Seriesを比較する場合、assert_series_equel
が使えます。
