pandasのapplyの進捗をtqdmで表示

概要

pandasで大きめのDataFrameに対してapplyすると、いつまで待てば結果が返るのか不安になることがあるのでtqdmで進捗を表示したい。

バージョン情報

• Python 3.6.8
• tqdm==4.31.1
• pandas==0.23.4

サンプルデータ

ランダムな数値で1万行のDataFrameを作成しておく。

import numpy as np
import pandas as pd

# 適当なDataFrameを生成する
df = pd.DataFrame(np.random.random([10000, 2]), columns=['a', 'b'])
df.head()

実行結果

          a         b
0  0.748046  0.457551
1  0.657718  0.628164
2  0.395166  0.623360
3  0.610704  0.236926
4  0.576350  0.902701

やり方

pandasでの進捗表示の仕方はtqdmのexamplesの配下にサンプルコードが提示されている。

https://github.com/tqdm/tqdm/blob/master/examples/pandas_progress_apply.py

tqdmにはprogress_applyという進捗を出しながらpandasでapplyする為の機能が用意されている。

tqdm.pandas() を実行した後に、普段applyと書いているところをprogress_applyに変更すれば良いようだ。

import time
from tqdm import tqdm

tqdm.pandas()

def sleep(a):
    time.sleep(0.01)
    return a
df['c'] = df.a.progress_apply(sleep)

実行結果。applyの進捗が表示されている。

実行中に一度下記のようなエラーに遭遇した。

ImportError: cannot import name 'PanelGroupBy'

少し古いバージョンを使っていたのが原因で、4.23.3以前で起きていたらしい。

tqdm.pandas() ImportError: cannot import 'PanelGroupBy' · Issue #555 · tqdm/tqdm · GitHub

tqdmを最新に更新したら発生しなくなった。

progress_applyの仕組み

progress_applyはpandasが持っている機能ではないので、tqdm.pandas() を実行してからでないと呼び出そうとしても「DataFrameはそんな属性持ってないよ」と言われる。

# tqdm.pandas()せずに実行した場合
df.progress_apply(lambda x: x)
    #=> AttributeError: 'DataFrame' object has no attribute 'progress_apply'

tqdm.pandas() を実行するとDataFrameにprogress_applyが挿し込まれる。

下記はバージョン4.31.1の tqdm.pandas() のコード。

https://github.com/tqdm/tqdm/blob/v4.31.1/tqdm/_tqdm.py#L570

# progress_applyを定義しているところのコード
DataFrame.progress_apply = inner_generator()
DataFrameGroupBy.progress_apply = inner_generator()
DataFrame.progress_applymap = inner_generator('applymap')

DataFrameだけでなく下記にも適用されているので、SeriesでもDataFrameでも意識せずに使うことができる。

• DataFrameGroupBy
• Series
• SeriesGroupBy
• Panel
• PanelGroupBy
• GroupBy

またmap、applymap、aggregate、transformに対してもそれぞれ当該関数が用意されている。

下記のようにaggregateもtransformも progress_ を付けるだけで普段書いているコードそのままに実行することができる。

df.groupby('a').progress_aggregate(np.sum)

df.groupby('a').progress_transform(np.sum)

改定履歴

Author: Masato Watanabe, Date: 2019-03-06, 記事投稿