tqdmがpandasのapplyで便利になっていました(モンキーパッチの実戦例)

こんにちは、沙代です。
pandasのapplyでスマートに済まそうと思ったら、思ったより時間がかかって進捗がわからずドギマギすることって、ありますよね。
pandas, tqdm(プログレスバーを表示するためのパッケージ)が便利になっていたので、コードを覗いてみようと思います。

環境
MacOS Montery 12.6.2 (M1 Pro)
Python 3.11.1
pandas 2.2.3
tqdm 4.66.5

コード例

tqdm.pandas()と.progress_applyを使います。
progress_applyの使い方は、applyとほぼ同じです。

example.py

import pandas as pd
import numpy as np
import time
from tqdm import tqdm

df = pd.DataFrame(np.random.randint(0, 100, (10, 6)))

# here!
tqdm.pandas()

# slow function
def slow_func(row):
    time.sleep(1)
    return row.sum()

# Use progress_apply instead of apply
df.progress_apply(slow_func, axis=1)

applyにあたって、プログレスバーが表示されるようになりました。

(前置き)モンキーパッチとは

Wikipediaで調べてみます。

モンキーパッチ（Monkey patch）は、システムソフトウェアを補完するために、プログラムをその時その場の実行範囲内で拡張または修正するというテクニックである。モンキーパッチの影響はその時その場のプロセス（プログラムの実行インスタンス）だけに限定されて、プログラム本体には及ばない。
(中略)
当初はモンキーパッチは、ルールを無視して実行時にこっそりとコードを変更することから、ゲリラパッチと呼ばれていた。これらのパッチを複数当てると、時折直感に反するような相互作用が生まれることがあり、Zope 2では、交戦中のパッチと呼ばれていた。
ゲリラはゴリラとほぼ同音異字であり、ゲリラパッチをゴリラパッチと言う人が出てきた。そしてゴリラがより弱いモンキーとなり、戦いを思い起こさせるゲリラパッチではなく、弱そうに聞こえるモンキーパッチという言葉として使われ始めた。3)
それ以来モンキーパッチとして使用され続けているが、意味は使用するコミュニティ毎に微妙に異なっている。[3]

正確には、今回は置き換えではなく、追加ですね。

動作の概略

今の実装とは違いますが、参考[1]のページにコードのイメージが書いてあります。

def tqdm_pandas(t):
  from pandas.core.frame import DataFrame
  def inner(df, func, *args, **kwargs):
      t.total = groups.size // len(groups)
      def wrapper(*args, **kwargs):
          t.update(1)
          return func(*args, **kwargs)
      result = df.apply(wrapper, *args, **kwargs)
      t.close()
      return result
  DataFrame.progress_apply = inner

最後の行でDataFrameに、progress_applyメソッドを追加していますね。

そのメソッドは、(df, func, *args, **kwargs)(dfはselfに該当)を受けて、
内部で、df.applyに、(wrapper, *args, **kwargs)を渡していますね。
1回funcが実行される代わりに1回wrapperが実行されて、結局はt.update(1)とfuncが実行される感じですね。

df.applyに渡される関数がスマートにwrapされているのが見ることができます。(クロージャもスマートに使われていますね。)

DataFrameがimport時に参照で渡されているから(?)、メソッドをここで追加してもグローバルに有効なのでしょう。

実際のコード

では、順番に実際のコードを追いかけてみます。

__init__.py

tqdm/__init__.py

2行目のところで、

from ._tqdm_pandas import tqdm_pandas

とみえます。

tqdm/_tqdm_pandas.pyを見てみると、

def tqdm_pandas(tclass, **tqdm_kwargs):
    ...
        TqdmDeprecationWarning(
            "Please use `tqdm.pandas(...)` instead of `tqdm_pandas(tqdm(...))`.",
            fp_write=getattr(tclass.fp, 'write', sys.stderr.write))
        type(tclass).pandas(deprecated_t=tclass)

どうもこちらの方はdeprecateされるようです。よく見てみると、tqdm.pandasではなくて、tqdm_pandasでした。

tqdm/std.py

tqdm/std.py
気を取り直して、tqdmクラスの実装を見てみます。

768行目にあります。(途中適宜省略)

    @classmethod
    def pandas(cls, **tqdm_kwargs):
        ...
        from pandas.core.frame import DataFrame
        from pandas.core.series import Series

        def inner_generator(df_function='apply'):
            def inner(df, func, *args, **kwargs):
                # Precompute total iterations
                total = tqdm_kwargs.pop("total", getattr(df, 'ngroups', None))

                # Init bar
                t = cls(total=total, **tqdm_kwargs)

                # Define bar updating wrapper
                def wrapper(*args, **kwargs):
                    # update tbar correctly
                    # it seems `pandas apply` calls `func` twice
                    # on the first column/row to decide whether it can
                    # take a fast or slow code path; so stop when t.total==t.n
                    t.update(n=1 if not t.total or t.n < t.total else 0)
                    return func(*args, **kwargs)

                # Apply the provided function (in **kwargs)
                # on the df using our wrapper (which provides bar updating)
                try:
                    return getattr(df, df_function)(wrapper, **kwargs)
                finally:
                    t.close()

            return inner

        DataFrame.progress_apply = inner_generator()
        DataFrameGroupBy.progress_apply = inner_generator()
        DataFrame.progress_applymap = inner_generator('applymap')

apply以外にもapplymapなどにもほぼ同じコードで対応できるように、

DataFrame.progress_apply = innerの代わりに、
DataFrame.progress_apply = inner_generator()が用いられていますね。

新たな疑問

pandasのPanelってなんでしょう？

参考にさせていただいたページ

[1] https://github.com/tqdm/tqdm/blob/master/examples/pandas_progress_apply.py
[2] https://github.com/tqdm/tqdm/tree/v4.66.5
[3] https://ja.wikipedia.org/wiki/モンキーパッチ