Pandas DataFrame の任意の場所に列を追加したい場合は pandas.DataFrame.insert を使うと実現できるが、いくつか不満がある。

immutable でない (DataFrame が直接書き換えられてしまう / 破壊的メソッドである / inplace オプションがない)
追加場所をインデックスの数値で指定しなければいけない
- わかりにくい
- 「col1 の後に追加」のような指定が難しい
Series を入力しても name が無視されるため、別途指定する必要がある
DataFrame を入力できない

これらを解決するような便利関数を書いた。

実装

from typing import Union, Optional
import pandas as pd


def insert_columns(
        df: pd.DataFrame,
        data: Union[pd.Series, pd.DataFrame],
        *,
        before: Optional[str] = None,
        after: Optional[str] = None,
        allow_duplicates: bool = False,
        inplace: bool = False,
    ) -> pd.DataFrame:

    if not inplace:
        df = df.copy()

    if not (after is None) ^ (before is None):
        raise ValueError('Specify only "before" or "after"')

    if before:
        loc = df.columns.get_loc(before)
    else:
        loc = df.columns.get_loc(after) + 1

    if type(data) is pd.Series:
        df.insert(loc, data.name, data, allow_duplicates)
    elif type(data) is pd.DataFrame:
        for column in data.columns[::-1]:
            df.insert(loc, column, data[column], allow_duplicates)

    return df

before または after にはカラム名を指定する
デフォルトでは入力された DataFrame を直接書き換えない
- inplace=True を指定すれば直接書き換える

使い方

サンプルデータとして sklearn の iris データセットを使う。

from sklearn import datasets

iris = datasets.load_iris()

df = pd.DataFrame(iris.data, columns=iris.feature_names)
target = pd.Series(iris.target_names[iris.target], name='target')

df.head()

target.head()

df の sepal width (cm) の後ろに target を追加してみる。

insert_columns(df, target, after='sepal width (cm)')

この例では Series を追加しているが、DataFrame も指定できる。