pandas.DataFrame.applyを使って、DataFrameに複数列を一度に追加したい

8512 ワード
Python pandas Python テキストリンク

TL;DR

  • pandas.DataFrame.applyのオプショナル引数に、result_type="expand"を指定する
  • pandas.DataFrame.applyの引数の関数(ラムダ式)は、タプルまたはリストを返すようにする
  • 代入式の左辺では、追加する列名をリストで指定する
def get_values(value0):
    # some calculation
    return value1, value2

df[["column1", "column2"]] = df.apply(
    lambda r: get_values(r["column0"]), 
    axis=1, 
    result_type="expand")

解説

適当なpandas.DataFrameがあるとします。

import pandas as pd

df = pd.DataFrame()
df["x"] = [1, 2, 3, 4, 5]
print(df)

   x
0  1
1  2
2  3
3  4
4  5

このxという名前の列を使って計算を行い、新たな列を作成したいとします。ベクトルのまま計算できる場合は簡単です。ここでは2乗を計算する場合を例に取ってみます。

df["x2"] = df["x"] ** 2
print(df)

   x  x2
0  1   1
1  2   4
2  3   9
3  4  16
4  5  25

もし新たな列の求める計算が複雑な関数で定義されており、ベクトルのまま計算できない場合には、pandas.DataFrame.applyを用いることができます。

def square(x):
    return x ** 2

df["x2"] = df.apply(lambda r: square(r["x"]), axis=1)
print(df)

もちろん結果は前と同じです。

さて、以下のような複数の値を返す関数を使って、複数列を同時に追加する場合はどうなるでしょうか？

def get_values(x):
    return x ** 2, x ** 3

次のように書くとエラーになります。

df[["x2","x3"]] = df.apply(lambda r: get_values(r["x"]), axis=1)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-19-3039b5c1e701> in <module>
      1 def get_values(x):
      2     return x**2, x**3
----> 3 df[["x2", "x3"]] = df.apply(lambda r: get_values(r["x"]), axis=1)

(中略)

ValueError: Must have equal len keys and value when setting with an iterable

このエラーが起きている理由は、右辺の値をprintすることで分かります。

print(df.apply(lambda r: get_values(r["x"]), axis=1))

0       (1.0, 1.0)
1       (4.0, 8.0)
2      (9.0, 27.0)
3     (16.0, 64.0)
4    (25.0, 125.0)
dtype: object

2個の要素を持つタプルが入った(1列の)Seriesになってしまっています。
2列5行のDataFrameにするには、result_type="expand"を指定します。

df[["x2","x3"]] = df.apply(lambda r: get_values(r["x"]), axis=1, result_type="expand")
print(df)

   x    x2     x3
0  1   1.0    1.0
1  2   4.0    8.0
2  3   9.0   27.0
3  4  16.0   64.0
4  5  25.0  125.0

参考にしたリンク

xmlファイルへのプロファイルの読み込み
Spring data japのopenEntityManagerInViewFilterとhibernateのopenSessionInViewFilter