DataFrame に新しい列を追加するとき、直接代入する方法がもっとも一般的ですが、assign メソッドを使うとメソッドチェーンの中で列を追加でき、コードの可読性が上がります。
直接代入との違い
まずは従来の直接代入を確認しておきます。
import pandas as pd
df = pd.DataFrame({
"name": ["Alice", "Bob", "Charlie"],
"price": [1000, 2000, 1500],
"quantity": [3, 1, 4]
})
df["total"] = df["price"] * df["quantity"]
この書き方は元の DataFrame を直接変更(破壊的操作)します。一方、assign は新しい DataFrame を返すため、元のデータは変わりません。
df2 = df.assign(total=df["price"] * df["quantity"])
元の DataFrame を直接変更する。シンプルだが副作用がある
新しい DataFrame を返す。元のデータを変えずにチェーンできる
ラムダ式を使う
assign にはラムダ式を渡すこともできます。ラムダ式の引数にはその時点の DataFrame が渡されるため、直前のチェーンで加工した結果を参照できます。
df2 = (
df
.assign(total=lambda x: x["price"] * x["quantity"])
.assign(tax=lambda x: x["total"] * 0.1)
.assign(grand_total=lambda x: x["total"] + x["tax"])
)
print(df2)
ラムダ式を使うメリットは、途中で作った列をすぐ次の assign で参照できる点にあります。上の例では total を先に定義し、それを使って tax と grand_total を順番に計算しています。
複数列を一度に追加する
assign は複数のキーワード引数を受け取れるため、1 回の呼び出しで複数列を追加できます。
df2 = df.assign(
total=lambda x: x["price"] * x["quantity"],
avg_price=lambda x: x["price"].mean()
)
ただし同じ assign 内では、先に定義した列を後の列から参照できない点に注意してください。依存関係がある場合は assign を分けてチェーンする必要があります。
列名に変数を使う
列名を動的に指定したい場合は、辞書のアンパック(**)を使います。
col_name = "total"
df2 = df.assign(**{col_name: df["price"] * df["quantity"]})
プログラム的に列名を生成するケースや、ループで複数列を追加するケースで役立つテクニックです。
assign はメソッドチェーンの中に自然に組み込めるため、データの加工手順を上から下へ一連の流れとして記述できます。pipe と組み合わせればさらに整理されたコードが書けるでしょう。