pandasで複数カラムを参照して高速に1行1行値を調整する際のメモ

TL;DR

pandasでデータフレームを扱っていて、ローカルでのちょっとした集計などで、たまに1行１行（複数のカラムの）値を参照して集計や補正をしたいことがあります。

その際、許容できるくらいの計算時間に収めるための方法のメモです。（色々やり方はあると思いますが、一例として）

500万行程度の、そこまで大きくないデータを想定します。

この記事で使うもの

以下のものだけで対応することを想定します。

• Pythonのビルドインのものやpandas、NumPyなどの基本的なもの

この記事で扱わないもの

• NumbaだったりCythonなど。（機会があれば将来別の記事で・・）

ダミーデータの準備

記事内で扱うための、適当な500万行のデータを用意します。

import pandas as pd
import numpy as np

df = pd.DataFrame(
    columns=['apple_price', 'orange_price', 'melon_price'],
    index=np.arange(0, 5000000))

df.apple_price = np.random.randint(low=60, high=160, size=5000000)

df.orange_price = np.random.randint(low=70, high=140, size=5000000)

df.melon_price = np.random.randint(low=120, high=340, size=5000000)

df[:3]

	apple_price	orange_price	melon_price
0	119	106	181
1	75	117	243
2	113	120	172

len(df)

5000000

上記のようなデータを使って、適当ですが例えば以下のような条件で複数のカラムを参照して、fruit_typeというカラムに値を設定することを想定します。

• もしapple_priceが120以上且つorange_priceが130以上 -> fruit_type = 1
• 上記以外で、もしapple_priceが130以下且つmelon_priceが200以上 -> fruit_type = 2
• 上記以外で、もしorange_priceが100以下且つmelon_priceが300以下 -> fruit_type = 3
• それ以外 -> fruit_type = 4

とりあえずループは遅い

Pythonで計算をする際によく言われることで、ループを書いたりすると大分遅くなります。数千件程度であればさくっとループで対応したりすることも多いですが、今回のような7桁件数だと結構苦しくなります。

とりあえず、比較としてデータフレームのiterrows関数でループを回して、1行1行値を設定してみます。もちろん、7桁件数など流すと終わらなくなってしまうので、100件程度に絞って実施してみます。

sliced_df = df[:100]

%%timeit
for index, sr in sliced_df.iterrows():
    if sr.apple_price >= 120 and sr.orange_price >= 130:
        df.loc[index, 'fruit_type'] = 1
        continue

    if sr.apple_price <= 130 and sr.melon_price >= 200:
        df.loc[index, 'fruit_type'] = 2
        continue

    if sr.orange_price <= 100 and sr.melon_price <= 300:
        df.loc[index, 'fruit_type'] = 3
        continue

    df.loc[index, 'fruit_type'] = 4

7.9 s ± 635 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

ループの処理が遅いだけでなく、シリーズやデータフレームへの個別のアクセスは大分遅いので、件数を少なくしているにも関わらずとても時間がかかっています。

※ms = 1/1000秒、µs = 1/1000/1000秒、ns = 1/1000/1000/1000秒

# シリーズの値にアクセスする場合の速度確認。
%timeit apple_price = sr.apple_price

27.5 µs ± 2.62 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

# データフレームの個別の個所に値を設定する場合の速度確認。
%timeit df.loc[0, 'fruit_type'] = 1

85.3 ms ± 20 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

比較として、Pythonの辞書のデータで、値の参照や設定をしてみると、個別のシリーズやデータフレームへのアクセスが遅いことがよくわかります。

sample_dict = {'apple_price': 100}

%timeit apple_price = sample_dict['apple_price']

66.8 ns ± 4.6 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)

%timeit sample_dict['fruit_type'] = 1

138 ns ± 10.5 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)

apply関数を利用する

基本的には、pandas側で色々関数が用意されているので、それらを使って全行一気に処理したり、スライスやベクトル演算的なことをして対応ができれば、それがシンプルで計算も早く終わります。

今回は、それらでは対応ができないと仮定して、pandasのapply関数を使って、各行ごとに任意の関数を適用していってみます。

このapply関数ですが、指定したい関数を用意しつつ、以下のように書きます。

def get_multiplied_price(apple_price):
    # サンプルのため、適当な記述にしてあります。
    print('apple_price :', apple_price)
    return apple_price * 2

sliced_df = df[:3]

sliced_df

	apple_price	orange_price	melon_price	fruit_type
0	119	106	181	1.0
1	75	117	243	2.0
2	113	120	172	4.0

sliced_df.apple_price.apply(get_multiplied_price)

apple_price : 119
apple_price : 75
apple_price : 113


0    238
1    150
2    226
Name: apple_price, dtype: int64

データフレームのapple_priceの列に対してapplyを実行し、引数に反映する関数（get_multiplied_price）を指定します。そうすると、apple_priceの各行の値がget_multiplied_price関数の第一引数に渡されて、且つ指定した行数のシリーズが返却されます。

これを利用して各行の値に応じて、既存のカラムの値を任意の関数を反映して更新したり、新しいカラムを追加したりすることができます。

しかしながら、上記のようなコードを見て分かる通り、引数にデータフレーム内の1つのカラムの値（apple_price）しか指定できていません。今回目的とする値では、3カラム分関数内で必要になります。

データフレーム自体にapply関数を実行した場合、第一引数の値はシリーズとなりますが、これでは先ほど触れたとおり、シリーズの値に対して1行ごとにアクセスするため、とても時間がかかります。

def get_fruit_type(sr):
    """
    果物の値段に応じた種別値を取得する。

    Parameters
    ----------
    sr : Series
        各果物の値を格納したシリーズ。

    Returns
    -------
    fruit_type : int
        各値段に応じて、1～4までの値が設定される。
    """

    apple_price = sr.apple_price
    orange_price = sr.orange_price
    melon_price = sr.melon_price

    if apple_price >= 120 and orange_price >= 130:
        return 1

    if apple_price <= 130 and melon_price >= 200:
        return 2

    if orange_price <= 100 and melon_price <= 300:
        return 3

    return 4

# 1行ずつの値をget_fruit_type関数の第一引数に指定するため、
# axis=1を指定しています。
# 目的とする各行の値は求まるものの、処理時間がとてもかかります。
sliced_df.apply(get_fruit_type, axis=1)

0    4
1    2
2    4
dtype: int64

１行あたりどのくらい処理に時間がかかるのか、直接関数を実行してみて確認してみます。

%timeit get_fruit_type(sr=sliced_df.iloc[0])

847 µs ± 61.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

1行あたり約850µsとしましょう。500万行なので、大体71分程度かかります。

計算できないこともないですが、これだけの計算で1時間以上かかるのは大分辛い感じです。

850 * 5000000 / 1000 / 1000 / 60

70.83333333333333

ではどうするのか。

解決策の1例ですが、以下のような対応でPythonとpandasなど基本的なものだけで、比較的高速に計算することができます。

• データフレームのインデックスをユニークな値にしておきます。
• データフレームのインデックスをキーとして、各カラムの値を格納した辞書を用意します。
  • ※辞書への変換はpandasのto_dict関数を使うことで、全行まとめて1回のみで終わるため、数秒程度で終わります。
• インデックスの値を格納するカラムをデータフレームに追加します。
• apply関数の引数に、インデックスの値のカラムを指定し、そのインデックスがキーに設定されている値を各辞書から取得します。
  • 前述のとおり、Pythonの辞書へのアクセスであれば、シリーズにアクセスするのと比べてかなり高速に動作するため、1行辺りの処理が大分早くなります。

# データフレームのインデックスをユニークな値にする。
# ※基本的に、連結などをしていなければ元々ユニークな連番で
# 割り振られています。
df.reset_index(drop=True, inplace=True)

# 各カラムの辞書を用意します。キーにはデータフレームのインデックス
# が設定されます。
apple_price_dict = df.apple_price.to_dict()

orange_price_dict = df.orange_price.to_dict()

melon_price_dict = df.melon_price.to_dict()

# インデックスの値を格納するカラムをデータフレームに追加します。
df['index_val'] = df.index

# apply関数の引数に、インデックスの値のカラムを指定して、関数内で
# そのインデックスを参照して辞書から各値を取得します。

def get_fruit_type(index):
    """
    果物の値段に応じた種別値を取得する。

    Parameters
    ----------
    index : int
        対象の行のデータフレームのインデックス。

    Returns
    -------
    fruit_type : int
        各値段に応じて、1～4までの値が設定される。
    """

    apple_price = apple_price_dict[index]
    orange_price = orange_price_dict[index]
    melon_price = melon_price_dict[index]

    if apple_price >= 120 and orange_price >= 130:
        return 1

    if apple_price <= 130 and melon_price >= 200:
        return 2

    if orange_price <= 100 and melon_price <= 300:
        return 3

    return 4

試しに、直接関数を実行してみて処理時間を計ってみます。

%timeit get_fruit_type(index=0)

1.8 µs ± 150 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

変更前が大体850µsだったので、約470倍くらい早くなりました。

辞書への変換などの他の処理の分も含めても、500万行への関数反映で30秒以内に終わるのでは、というレベルになりました。

# データフレームのインデックスの値が、apply関数の引数に渡るように指定します。
df['fruit_type'] = df.index_val.apply(get_fruit_type)

df[:5]

	apple_price	orange_price	melon_price	fruit_type	index_val
0	119	106	181	4	0
1	75	117	243	2	1
2	113	120	172	4	2
3	97	118	303	2	3
4	102	110	201	2	4

df.fruit_type.unique()

array([4, 2, 3, 1])

Pythonとほぼpandaのみで、ローカルでちょこっと計算したい場合などには必要十分な速度を出すことができました。（Cythonとか、使い方をよく把握していない、といった方にもpandasくらいで対応できるので（学習面などで）楽かもしれません）

今回は適当な関数を使いましたが、実際の業務ではもっと複雑な条件などの関数が必要になってくると思います。そういった場合でも関数が適用できるというのは、いろいろ柔軟にデータの変形などを対応できるので、知っておいても損はないかもしれません。

デメリットは？

• 途中で辞書変換を挟む都合、メモリがその分余分に必要になります。

おまけ

実行環境 : Azure Notebooks

# Pythonバージョン :
!python -V

Python 3.5.4 :: Anaconda custom (64-bit)

ノートをgithubにアップしておきました。

---

他にもPythonなどを中心に色々記事を書いています。そちらもどうぞ！
今までに投稿した主な記事たち