Pandas DataFrame に期待する列が存在するかどうかをチェックする

例えば以下のような「DataFrame を加工する関数」があったとする。

import pandas as pd

def preprocess(df: pd.DataFrame) -> pd.DataFrame:
    df["full_name"] = df["first_name"] + " " + df["last_name"]
    return df

この関数の引数の DataFrame には first_name と last_name という列が含まれることが期待されるが、これを関数のはじめでチェックしたいことがある。

これは set 型 ¹ の演算を使うと簡単に書ける。

import pandas as pd

def preprocess(df: pd.DataFrame) -> pd.DataFrame:
    required_columns = {"first_name", "last_name"}
    if not required_columns <= set(df.columns):
        raise ValueError(f"missing columns: {required_columns - set(df.columns)}")
    df["full_name"] = df["first_name"] + " " + df["last_name"]
    return df

このように書いておくと、必要な列が抜けていた場合に ValueError を投げてくれる。

df = pd.DataFrame([{"first_name": "John", "age": 30}])  # 'last_name' 列が抜けている DataFrame
preprocess(df)  #=> ValueError: missing columns: {'last_name'}

---

1. https://docs.python.org/ja/3/library/stdtypes.html#set ↩