データフレームの比較にpandasのassert_frame_equalが便利だった

はじめに

データアナリティクス事業本部のおざわ（じ）です。

今回は個人的に便利だと思ったpandasの関数をご紹介したいと思います。

先日、２つのデータフレームを比較する必要があったため、なにかよい方法がないか調べていたところ、pandasのassert_frame_equalを発見しました。主にユニットテストで使われる関数で、動きとしては名前のとおりデータフレームの内容が等しいか確認するものです。さまざまなパラメータが用意されていて、用途に合わせてこれらを調整することで、厳しめにもゆるめにもデータフレームの比較ができます。

環境

今回使った環境です。

• macOS Monterey 12.6
• Python: 3.8.10
• pandas: 1.2.0

使用例

さっそく公式ドキュメントにある例を見てみます。 この例では、df1とdf2で同じ名前のカラムが設定されており、２つのレコードが存在しています。違いとしてはbのカラムが一方ではint型で、もう一方ではfloat型になっている点です。

from pandas.testing import assert_frame_equal
df1 = pd.DataFrame({'a': [1, 2], 'b': [3, 4]})
df2 = pd.DataFrame({'a': [1, 2], 'b': [3.0, 4.0]})

# df1同士の比較なのでPass
assert_frame_equal(df1, df1)

# AssertionError
assert_frame_equal(df1, df2)

Traceback (most recent call last):
...
AssertionError: Attributes of DataFrame.iloc[:, 1] (column name="b") are different

Attribute "dtype" are different
[left]:  int64
[right]: float64

このように2つのデータフレームが等しくない場合、AssertionErrorを出してくれます。今回はbのカラムで型が異なっていますので、エラーメッセージにもそのことが書かれています。 Errorにあるleftはdf1のことで、rightがdf2です。

型チェックをスキップ

上記の例にあるデータフレームを同じものと判定したい場合はcheck_dtypeの引数にFalseを指定すると型チェックをスキップしてくれます。

# Pass
assert_frame_equal(df1, df2, check_dtype=False)

今回の私の用途では、型を無視して比較したかったのでcheck_dtypeはFalseにしました。

数値の比較

また、数値に関して小数点以下どこまでの差異を許容するのかについてもrtolやatolで調整することができます。この２つはcheck_exactがFalseであることが前提です（デフォルトはFalse）。以下の例ではrtolを変更しています。デフォルト値には1e-5が設定されているので、最初のassert_frame_equalはAssertionErrorになりません。

厳しくチェックしてAssertionErrorとしたい場合は、下ではrtolの引数に1e-6を指定することでAssertionErrorを発生させています。

df1 = pd.DataFrame({'a': [1, 2], 'b': [3.0, 4.0]})
df2 = pd.DataFrame({'a': [1, 2], 'b': [3.00001, 4.0]})

# Pass
assert_frame_equal(df1, df2)

# AssertionError
assert_frame_equal(df1, df2, rtol=1e-6)

AssertionError: DataFrame.iloc[:, 1] (column name="b") are different

DataFrame.iloc[:, 1] (column name="b") values are different (50.0 %)
[index]: [0, 1]
[left]:  [3.0, 4.0]
[right]: [3.00001, 4.0]

データが異なる場合には、どの程度の違いがあるのか示してくれるのは助かります。

最後に

以上、最近見つけた便利な関数のご紹介でした。

データフレームの比較が必要になった場合は、assert_frame_equalを使ってみていただければと思います。ソースを覗いてみると、forループの中でassert_series_equalという関数を使って各カラムを比較していましたので、特定のカラムのみ比較する場合はこの関数が使えそうです。

参考リンク