Pandasのjson_normalizeを詳しく調べてみる

普段、仕事で入れ子になっているデータを扱う必要のある時がぼちぼちあります（特に、スキーマレスのMongoDBのログであったりなどなど）。

BigQueryなどに入れてしまえば、WITHINなどを使ったクエリでいい感じに行列の形に変えてくれたりして簡単ではありますが、そうではない場合はPandasで大規模なデータを扱う際には結構苦労することがあります。

その辺りいい感じにやってくれるライブラリでも書くか・・と思ったところ、調べていたらPandasにjson_normalizeというAPIがあるようです。使ったことがなかったので、色々動かしつつ調べてみます。

早速動かしてみる。

まずはimport。

import numpy as np
import pandas as pd
from pandas.io.json import json_normalize

仮に以下のような構造で、量が膨大なデータがあったとして、

log_data_list = [{
    'id': 1,
    'name': 'john',
    'equipment_status': [{
        'equipment_id': 1,
        'attack': 1000,
        'defense': 1200,
        'speed': 800,
        'luck': 500,
    }, {
        'equipment_id': 2,
        'attack': 1100,
        'defense': 1300,
        'speed': 900,
        'luck': 600,
    }]
}, {
    'id': 2,
    'name': 'henry',
    'equipment_status': [{
        'equipment_id': 3,
        'attack': 1200,
        'defense': 1400,
        'speed': 1000,
        'luck': 700,
    }, {
        'equipment_id': 4,
        'attack': 1300,
        'defense': 1400,
        'speed': 1000,
        'luck': 700,
    }, {
        'equipment_id': 5,
        'attack': 1400,
        'defense': 1500,
        'speed': 1100,
        'luck': 800,
    }]
}]

普通にデータフレームに放り込むとequipment_statusカラム内の値に対する集計などが扱いづらい。

df = pd.DataFrame(data=log_data_list)
df

	equipment_status	id	name
0	[{'equipment_id': 1, 'attack': 1000, 'defense'...	1	john
1	[{'equipment_id': 3, 'attack': 1200, 'defense'...	2	henry

json_normalize 関数で record_path 引数のところに、リストや辞書になっているカラム名を指定すると、それらを展開したデータフレームが生成されます。

df = json_normalize(data=log_data_list, record_path='equipment_status')
df

	attack	defense	equipment_id	luck	speed
0	1000	1200	1	500	800
1	1100	1300	2	600	900
2	1200	1400	3	700	1000
3	1300	1400	4	700	1000
4	1400	1500	5	800	1100

とりあえず行列にはなったものの、これだとどのユーザーのログなのかが分からなくなります。そのため、 meta 引数に、残したいカラムのリストを指定します。（今回はidとname両方残しましたが、idだけ残したければidだけ指定といった具合に）

df = json_normalize(
    data=log_data_list, record_path='equipment_status', meta=['id', 'name'])

	attack	defense	equipment_id	luck	speed	id	name
0	1000	1200	1	500	800	1	john
1	1100	1300	2	600	900	1	john
2	1200	1400	3	700	1000	2	henry
3	1300	1400	4	700	1000	2	henry
4	1400	1500	5	800	1100	2	henry

BigQueryでWITHINを使ったときのような結果になりました！Pandasでもシンプルですね・・（もっと早めに知っておけば、仕事で楽ができたのに・・）

ここまで変換すれば、あとはスライスなり集計なりをさくっと行えます。

パフォーマンス確認

適当なデータを用意して、雰囲気を確認してみます。
まずは10万行で、各行に5件ずつのリストが入っていることを想定します。（結果が50万行になる程度）

log_data_list = []

for i in range(100000):
    data_dict = {
        'id': i + 1,
        'name': 'john',
        'equipment_status': [{
            'equipment_id': i + 1,
            'attack': i * 100,
            'defense': i * 110,
            'speed': i * 120,
            'luck': i * 130,
        }, {
            'equipment_id': i + 2,
            'attack': i * 200,
            'defense': i * 210,
            'speed': i * 220,
            'luck': i * 230,
        }, {
            'equipment_id': i + 3,
            'attack': i * 300,
            'defense': i * 310,
            'speed': i * 320,
            'luck': i * 330,
        }, {
            'equipment_id': i + 4,
            'attack': i * 400,
            'defense': i * 410,
            'speed': i * 420,
            'luck': i * 430,
        }, {
            'equipment_id': i + 5,
            'attack': i * 500,
            'defense': i * 510,
            'speed': i * 520,
            'luck': i * 530,
        }]
    }
    log_data_list.append(data_dict)

%%timeit
df = json_normalize(
    data=log_data_list, record_path='equipment_status', meta=['id', 'name'])

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

1秒くらい。全然問題ないレベルです。

次に、100万行（変換後は500万行想定）でやってみます。

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

データ量に応じて、そのまま10倍くらいになっているようです。
日次で、ちょこちょことETLのcronを流したりする分には十分満足な速度ですね
（変換結果が億の桁数になってくると、メモリなどが辛くなってきそうですが、今担当している仕事での1日単位の1回の処理でそこまでいくことは今のところしばらく無さそうなので、安心です）