対象読者

日々のデータ分析でpythonを活用している人
データ分析や機械学習の初学者
データの前処理に時間がかかっている人

記事を書いたきっかけ

皆さんはテーブルデータの分析や前処理をどのような方法で行っていますか？

Pythonを使ったデータ分析の3種の神器として、
numpy/pandas/matplotlib を思い浮かべる方も多いのではないでしょうか？

実務でもPandasはテーブルデータの処理で非常によく使います。Pandasを使っていて感じるのが「コードが冗長になりがち」だということ。

近年ではPolarsというよりPandasより高性能のライブラリが登場しているものの、
他のライブラリとの依存関係や学び直しのコストなどを考えるとまだまだPandasから離れられない方が多いのではないでしょうか？

そんなPandsについて、できるだけ見やすく、冗長になりにくくいコードを書くコツに関してお話を聴く機会があったので一番重要なassignメソッドに関してまとめました。

assignメソッドを使いこなそう

Pandasでは、行列を持つテーブルデータをデータフレームクラスとして定義して取り扱います。
データフレームとして定義されたオブジェクトは .メソッド名() によりさまざまな処理を行うことが可能です。

メソッドの1つである.assignはデータフレームに対して、新たな行や列を追加することが可能です。
MLやデータ分析の学習でよく使われるirisデータセットを使い検証してみましょう。

import pandas as pd
from sklearn.datasets import load_iris

# irisデータセットをpd.DataFrameクラスとして定義
iris = load_iris()
df_iris = pd.DataFrame(
    data=iris.data,
    columns=[
        'sepal_length',
        'sepal_width',
        'petal_length',
        'petal_width']
    )
print(df_iris)

# 出力
     sepal_length  sepal_width  petal_length  petal_width
0             5.1          3.5           1.4          0.2
1             4.9          3.0           1.4          0.2
2             4.7          3.2           1.3          0.2
3             4.6          3.1           1.5          0.2
4             5.0          3.6           1.4          0.2
..            ...          ...           ...          ...
145           6.7          3.0           5.2          2.3
146           6.3          2.5           5.0          1.9
147           6.5          3.0           5.2          2.0
148           6.2          3.4           5.4          2.3
149           5.9          3.0           5.1          1.8

[150 rows x 4 columns]

df_irisとしてデータフレームオブジェクトを定義しました。
sepal_lengthの値を使い、0.5ごとに階級分けしたカテゴリカル変数を作成してみます。

この時、条件として5.0よりも小さいものは1つのカテゴリとしてまとめるようにしましょう。

まずはassignメソッドを使わない例で見てみます。

# 0.5で階級分け
df_iris_new = df_iris.copy()
df_iris_new['sepal_length_cat'] = (df_iris['sepal_length'] / 0.5).astype(int)

# maskメソッドを使って、値が5.0より小さいものを1つの階級に
df_iris_new['sepal_length_cat'] = df_iris_new['sepal_length_cat'].mask(df_iris_new['sepal_length_cat'] <= 9, 9)
print(df_iris_new)

# 出力
     sepal_length  sepal_width  petal_length  petal_width  sepal_length_cat
0             5.1          3.5           1.4          0.2                10
1             4.9          3.0           1.4          0.2                 9
2             4.7          3.2           1.3          0.2                 9
3             4.6          3.1           1.5          0.2                 9
4             5.0          3.6           1.4          0.2                10
..            ...          ...           ...          ...               ...
145           6.7          3.0           5.2          2.3                13
146           6.3          2.5           5.0          1.9                12
147           6.5          3.0           5.2          2.0                13
148           6.2          3.4           5.4          2.3                12
149           5.9          3.0           5.1          1.8                11

[150 rows x 5 columns]

df_irisが一文に何度も出てきており、横長で冗長なコードになっています。
また、再帰代入と呼ばれる
df[カラム名] = df[カラム名].メソッド名()
のコードが登場しており、こうしたコードはバグの原因にもなるため書かない方が良いです。

ではassignメソッドを使うとどうでしょうか？

df_iris_new = (
    df_iris
        .assign(sepal_length_cat=lambda df: 
            (df.sepal_length / 0.5)
            .astype(int)
            .mask(lambda s: s <= 9, 9)
            )
)
print(df_iris_new)

# 出力
     sepal_length  sepal_width  petal_length  petal_width  sepal_length_cat
0             5.1          3.5           1.4          0.2                10
1             4.9          3.0           1.4          0.2                 9
2             4.7          3.2           1.3          0.2                 9
3             4.6          3.1           1.5          0.2                 9
4             5.0          3.6           1.4          0.2                10
..            ...          ...           ...          ...               ...
145           6.7          3.0           5.2          2.3                13
146           6.3          2.5           5.0          1.9                12
147           6.5          3.0           5.2          2.0                13
148           6.2          3.4           5.4          2.3                12
149           5.9          3.0           5.1          1.8                11

[150 rows x 5 columns]

出力結果は全く一緒ですがインデントも整理されており、コードが見やすくなっています。
assignを使った書き方では、.methodを繋げて書くメソッドチェーンと呼ばれる書き方を使っています。これによって各メソッドの処理範囲がわかりやすく、またSQLライクな書き方にすることが可能です。

また、再帰代入も行っておらずシンプルでわかりやすいコードとなっています。

終わりに

このようにPandasでは、assignメソッドを活用することで再帰代入や冗長なコードを減らすことができます。また、メソッドの中でも処理時間が長いapplyやfor文を使ったデータフレーム処理もassignでシンプルに描ける場合が多いです。

assignメソッドはlambda関数と組み合わせることでさまざまな処理が可能となるため、
Pandasを使ってデータ処理される方々はぜひ使い方をマスターしましょう！

参考文献

改訂新版前処理大全〜SQL/pandas/Polars実践テクニック(https://gihyo.jp/book/2024/978-4-297-14138-7)