pandasのMultiIndexで階層取り出しでミスった

pandasのMultiIndexは便利なのだが、単純にIndexの多次元版として扱っていたらハマったのでメモしておく。

ハマったこと

例として次のようなテーブルがhoge.csvとして存在することとする。

		val
1	a	b
2	c	d
3	a	d
4	b	c
5	a	b

hoge.csvのval以外の列をindexとして読みだしてみるとMultiIndexのDataFrameとして読みだされる

>>> import pandas as pd
>>> df = pd.read_csv("hoge.csv", index_col=[0, 1])

>>> df
    val
1 a   b
2 c   d
3 a   d
4 b   c
5 a   b

この適当なDataFrameをvalでフィルタリングしてみる

>>> tmp_df = df.query("val=='b'")

>>> tmp_df.index
MultiIndex([(1, 'a'),
            (5, 'a')],
           )

全5要素のDataFrameから2要素が取り出された。

更に、フィルタをかけた結果に対してlevelsプロパティの0階層目を取得すると、1,5が取り出されるかと思いきや…

>>> tmp_df.index.levels[0]

Int64Index([1, 2, 3, 4, 5], dtype='int64')

フィルタに無関係でオリジナルのDataFrameの0階層目の要素群が取り出される
テーブルの値に条件を設けてフィルタ後の状態で、各階層の値を取り出したいことはあるので、これでは困る。

解決

levelsはあくまで、各階層に含まれる要素を格納しているリストで、各level間の関係を定義することで組み合わせて実現しているみたい。

そこで、MultiIndexを解除して、最終的に取り出したい階層を残したシングルのIndexにしてからフィルタをかける。

>>> df.reset_index(level=1)
  level_1 val
1       a   b
2       c   d
3       a   d
4       b   c
5       a   b
>>> tmp_df = df.reset_index(level=1).query("val=='b'")
>>> tmp_df.index
Int64Index([1, 5], dtype='int64')

こうすると、そのままIndexがフィルタの要素通りになるので、フィルタしてからある階層を取り出したい時は、上記のようにreset_indexで対応しないといけない。

reset_indexではMultiIndexのカラム名が降られていればその名称を、降られていなければlevel=の引数に解除したい階層の番号を指定する。