Python: pandas の DataFrame, Series, Index を拡張する

Python でデータ分析をするときに、ほぼ必ずといって良いほど使われるパッケージとして pandas がある。 そのままでも便利な pandas だけど、代表的なオブジェクトの DataFrame, Series, Index には実は独自の拡張を加えることもできる。 これがなかなか面白いので、今回はその機能について紹介してみる。

ただし、あらかじめ断っておくと注意点もある。 独自の拡張を加えると、本来は存在しないメソッドやプロパティがオブジェクトに生えることになる。 そのため、便利だからといってこの機能を使いすぎると、コードの可読性が低下する恐れもある。 使うなら、後から別の人がコードを読むときにも困らないようにしたい。 具体的には、使用するにしても最小限に留めたり、あるいはパッケージ化やドキュメント化をしておくことが挙げられる。

今回使った環境は次の通り。

$ sw_vers
ProductName:    Mac OS X
ProductVersion: 10.13.5
BuildVersion:   17F77
$ python -V
Python 3.6.5
$ pip list --format=columns | grep -i pandas
pandas          0.23.3

下準備

まずは pandas をインストールしておこう。

$ pip install pandas

ここからは Python の REPL を使って解説していく。

$ python

DataFrame を拡張する

まずは一番よく使うであろう DataFrame の拡張方法から。 あんまり実用的な例じゃないけど、ひとまず API がどんな感じになってるかを紹介したいので。

pandas のオブジェクトを拡張するときは、基本的に pandas.api.extensions 以下にある API を用いる。 例えば DataFrame を拡張するなら @pandas.api.extensions.register_dataframe_accessor() デコレータでクラスを修飾する。 次のサンプルコードでは DataFrame に helloworld という名前空間で greet() メソッドと length プロパティを追加している。

>>> import pandas as pd
>>> 
>>> # "helloworld" という名前空間で DataFrame を拡張する
... @pd.api.extensions.register_dataframe_accessor('helloworld')
... class HelloWorldDataFrameAccessor(object):
...     """DataFrameを拡張するためのクラス"""
...     def __init__(self, df):
...         self._df = df
...     # DataFrame#helloworld に greet() メソッドを追加する
...     def greet(self):
...         """標準出力にメッセージを出す"""
...         print('Hello, World!')
...     # DataFrame#helloworld に length プロパティを追加する
...     @property
...     def length(self):
...         """DataFrameの長さを返す"""
...         return len(self._df)
... 

これだけで DataFrame の拡張ができる。

実際に DataFrame のインスタンスを作って、上記の動作を確認してみよう。

>>> df = pd.DataFrame(list(range(1, 11)), columns=['n'])
>>> df
    n
0   1
1   2
2   3
3   4
4   5
5   6
6   7
7   8
8   9
9  10

特に意味はないけど DataFrame#helloworld#greet() メソッドを実行すると標準出力にメッセージが出るようになる。

>>> df.helloworld.greet()
Hello, World!

あとは DataFrame#helloworld#length プロパティを参照すると DataFrame の長さが得られるようになる。

>>> df.helloworld.length
10

たしかに DataFrame に自分で拡張したメソッドやプロパティを生やすことができた。

Series を拡張する

続いては Series の拡張方法を紹介する。 基本的にやることは先ほどと同じなので、次はもうちょっと実用的な例を紹介してみる。

例えば Series をマルチプロセスで並列に処理したい、というシチュエーションを考えてみよう。 使ってるマシンの CPU コアがたくさんあって、扱うデータセットが大きいときは結構やりたくなるんじゃないかな？ 典型的には、次のような高階関数を用意するはず。

>>> import multiprocessing as mp
>>> import numpy as np
>>> 
>>> def parallelize(f, data, n_jobs=None):
...     """関数の適用をマルチプロセスで処理する"""
...     if n_jobs is None:
...         # 並列度の指定がなければ CPU のコア数を用いる
...         n_jobs = mp.cpu_count()
...     # データを並列度の数で分割する
...     split_data = np.array_split(data, n_jobs)
...     # プロセスプールを用意する
...     with mp.Pool(n_jobs) as pool:
...         # 各プロセスで関数を適用した結果を結合して返す
...         return pd.concat(pool.map(f, split_data))
... 

続いて、マルチプロセスで適用したい関数を適当に用意する。

>>> def square(x):
...     return x * x
... 

そして、こんな感じで使う。

>>> parallelize(square, df.n)
0      1
1      4
2      9
3     16
4     25
5     36
6     49
7     64
8     81
9    100
Name: n, dtype: int64

先ほどの使い方でも構わないんだけど pandas のオブジェクトから直接呼び出せると便利そうなので拡張してみよう。 次のようにして Series に parallel という名前空間で apply() メソッドを追加する。

>>> # "parallel" という名前空間で Series を拡張する
... @pd.api.extensions.register_series_accessor('parallel')
... class ParallelSeriesAccessor(object):
...     """Seriesを拡張するためのクラス"""
...     def __init__(self, s):
...         self._s = s
...     # Series#parallel に apply() というメソッドを定義する
...     def apply(self, f):
...         """Series に対して関数を並列で適用する"""
...         return parallelize(f, self._s)
... 

すると Series#parallel#apply() メソッドが使えるようになる。

>>> df.n.parallel.apply(square)
0      1
1      4
2      9
3     16
4     25
5     36
6     49
7     64
8     81
9    100
Name: n, dtype: int64

呼び出し方が違うだけで、やっていることは先ほどと変わらない。

Index を拡張する

続いては Index を拡張してみよう。

以下のサンプルコードでは Index が整数という前提で偶数・奇数だけを抜き出す機能を追加している。 また、あんまり実用性がない例になっちゃった。

>>> # "sampling" という名前空間で Index を拡張する
... @pd.api.extensions.register_index_accessor('sampling')
... class SamplingIndexAccessor(object):
...     """Indexを拡張するためのクラス"""
...     def __init__(self, idx):
...         self._idx = idx
...     # Index#sampling に even というプロパティを定義する
...     @property
...     def even(self):
...         return self._idx[self._idx % 2 == 0]
...     # Index#sampling に odd というプロパティを定義する
...     @property
...     def odd(self):
...         return self._idx[self._idx % 2 != 0]
... 

早速試してみよう。

>>> df.index.sampling.even
Int64Index([0, 2, 4, 6, 8], dtype='int64')
>>> df.index.sampling.odd
Int64Index([1, 3, 5, 7, 9], dtype='int64')

ちゃんと偶数・奇数だけ取り出すことができた。

めでたしめでたし。

参考

Extending Pandas — pandas 0.23.3 documentation

Pythonによるデータ分析入門 第2版 ―NumPy、pandasを使ったデータ処理

スマートPythonプログラミング: Pythonのより良い書き方を学ぶ