python関数の引数でpandas dataframeを指定した時、指定したdataframeによって条件分岐させたい

Question

前提・実現したいこと
初歩的な内容で恐縮です。調べるキーワードが見つけられず、質問させて頂きます。
python関数の引数でpandas dataframeを指定した時、指定したdataframeによって条件分岐させたい

発生している問題・エラーメッセージ
The truth value of a DataFrame is ambiguous. Use a.empty, a.bool(), a.item(), a.any() or a.all().

該当のソースコード

#もちろんこれが実際のコードではありませんが、やりたい方向がこういう内容です。
def function_a(df):
  if df == train_df:
    a = 1
  if df == test_df:
    a = 2
  return a

function_a(train_df)

自分で調べたことや試したこと
if df == 'train_df':
と文字列にしてみても、ダメでした。
if df.iloc[0][2] == 1:
などとすれば出来ないことはないのですが、もっとスマートな書き方はありますか？

使っているツールのバージョンなど補足情報
Google Colaboratory上

Answer 1

import pandas as pd
import numpy as np

train_df = pd.DataFrame({
    'col1': ['a', 'a', 'b', 'b', 'a'],
    'col2': [1.0, 2.0, 3.0, np.nan, 5.0],
    'col3': [1.0, 2.0, 3.0, 4.0, 5.0]
  }, columns=['col1', 'col2', 'col3'])

test_df = pd.DataFrame({
    'col1': ['b', 'b', 'a', 'a', 'b'],
    'col2': [np.nan, 1.0, 2.0, 3.0, 4.0],
    'col3': [5.0, 4.0, 3.0, 2.0, 1.0]
  }, columns=['col1', 'col2', 'col3'])

def function_a(df):
  if df.equals(train_df):
    a = 1
  elif df.equals(test_df):
    a = 2
  else:
    a = 0
  return a

if __name__ == '__main__':
  df = train_df.copy()
  print(function_a(df))

  df = test_df.copy()
  print(function_a(df))

#
1
2

Comments

[coffeebar]

頂いたコードそのものを Colaboratoryに載せて、きちんと動作確認できました。
しかし、自分のコードに組み込もうとすると、今度は別の 'Can only compare identically-labeled DataFrame objects' というエラーが出てしまいました。

あまりスマートさにこだわらずに、

def function_a(df, phase):
if phase == 'train':
a = 1
if phase == 'test':
a = 2
return a
function_a(train_df, 'train')

と引数を増やすことで対応することにしました。
ご回答、ありがとうございました。

[melian]

一応、pandas.testing.assert_frame_equal というメソッドがあって、check_names=False を指定するとインデックスやカラム名の違いは無視して比較してくれます。ただ、使い方が面倒なのであまりお勧めはできません。
https://pandas.pydata.org/docs/reference/api/pandas.testing.assert_frame_equal.html

[coffeebar]

追加情報ありがとうございました。melian様のコメント、大変勉強になりました。
質問を締めさせて頂きます。

Answer 2

pandas.DataFrame.equalsでの判定を検討してみてください。

Comments

[coffeebar]

ご回答ありがとうございました。