ABEJA Tech Blog

中の人の興味のある情報を発信していきます

pandas DataFrameをもっと自由にJOINしよう!

f:id:archibay:20180220204119j:plainデータエンジニアの千葉です。

今回は、データ分析の必須ライブラリ「pandas」について、 データフレーム結合機能 mergeと 自作の拡張ライブラリの話を書いています。


苦しみポイント

先日、2つのpandas.DataFrameを、こんな感じで結合する必要がありました。

SELECT
    *
FROM
    df_a
JOIN
    df_b
ON
    df_a.xxx_id = df_b.xxx_id AND
    df_a.yyy_id = df_b.yyy_id AND
    df_b.ccc BETWEEN df_a.aaa AND df_b.bbb
;

BETWEENを使ったJOINです。

実はこの処理、pandas上では簡単にできません。

なぜかというと、pandasは結合機能は、

1. 標準の結合条件がすごくしょぼい!

2. カスタムするとすごく面倒くさい!

からです。

標準では、df_a.xxx_id = df_b.xxx_id相当の比較しかサポートされていません。 また、内部的に複雑な結合処理が動くため、どの関数をどのように実装すれば正しい評価ができるのか理解が難しいです。 (ちなみにこの結合処理は、バージョンアップで結構変更される部分なので、深入り注意です。)

カスタムクラスの__hash__関数や__eq__関数を適切にカスタマイズすることで、ある程度まではいけますが、 「複数の結合キーの整合性」「ソートによる最適化」なども考え始めると、本当に面倒くさいです。

今回は、

  • BETWEENを含む複数の結合条件を扱いたい
  • 1000万件級の中規模データに使いたい
  • ぱぱっと計算が終わってほしい

という苦しみから解放されるために、ライブラリを実装したので紹介します。


自作ライブラリについて

今回実装したライブラリはこちらです。

pandas BETWEEN JOIN

インストール

pip install pandas-bj

使い方

ほとんど、pandas.mergeと同じ要領で使えます。

left_onright_onの中に列名を入れる代わりに、pandas_bj.Betweenを入れればOKです。

left_onright_onが複数ある場合、全てAND結合として評価されます。

# インポートして
import pandas_bj
import pandas

# データの準備
df1 = pandas.DataFrame({
    'id1': [1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3],
    'id2': [1, 1, 1, 2, 2, 1, 1, 2, 2, 2, 1, 1, 1, 2, 2],
    's': [1, 2, 3, 4, 5, 2, 3, 4, 5, 6, 3, 4, 5, 6, 7],
    'e': [5, 6, 7, 8, 9, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]}
)

df2 = pandas.DataFrame({
    'id3': [1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3],
    'id4': [1, 1, 1, 2, 2, 1, 1, 2, 2, 2, 1, 1, 1, 2, 2],
    'v': [1, 5, 2, 6, 3, 7, 4, 8, 5, 9, 6, 10, 7, 11, 8]}
)

# こういう条件で結合したい
# ON xxx.id1 = yyy.id3 AND
#       xxx.id2 = yyy.id4 AND
#       yyy.v BETWEEN xxx.s AND xxx.e
result = pandas_bj.merge(
    left=df1, right=df2,
    left_on=['id1', 'id2', pandas_bj.Between('s', 'e', True, True)], 
    right_on=['id3', 'id4', 'v'],
    how='inner', sort=True
)

print(result)
       e  id1  id2    s  id3  id4     v
0    5.0  1.0  1.0  1.0  1.0  1.0   2.0
1    6.0  1.0  1.0  2.0  1.0  1.0   5.0
2    7.0  1.0  1.0  3.0  1.0  1.0   5.0
3    8.0  1.0  2.0  4.0  1.0  2.0   6.0
4    9.0  1.0  2.0  5.0  1.0  2.0   6.0
5    5.0  2.0  1.0  2.0  2.0  1.0   4.0
6    6.0  2.0  1.0  3.0  2.0  1.0   4.0
7    7.0  2.0  2.0  4.0  2.0  2.0   5.0
8    9.0  2.0  2.0  6.0  2.0  2.0   8.0
9   10.0  3.0  1.0  3.0  3.0  1.0   6.0
10  10.0  3.0  1.0  3.0  3.0  1.0   7.0
11  11.0  3.0  1.0  4.0  3.0  1.0   6.0
12  11.0  3.0  1.0  4.0  3.0  1.0  10.0
13  11.0  3.0  1.0  4.0  3.0  1.0   7.0
14  12.0  3.0  1.0  5.0  3.0  1.0   6.0
15  12.0  3.0  1.0  5.0  3.0  1.0  10.0
16  12.0  3.0  1.0  5.0  3.0  1.0   7.0
17  13.0  3.0  2.0  6.0  3.0  2.0  11.0
18  13.0  3.0  2.0  6.0  3.0  2.0   8.0
19  14.0  3.0  2.0  7.0  3.0  2.0  11.0
20  14.0  3.0  2.0  7.0  3.0  2.0   8.0

pandas_bj.Betweenの引数

  1. 区間の小さい方の区切りの列名
  2. 区間の大きい方の区切りの列名
  3. 左閉区間か左開区間か
  4. 右閉区間か右開区間か

例:

引数 区間
pandas_bj.Between('a', 'b', False, False)  a \leq x \leq b
pandas_bj.Between('start', 'end', True, False)  start \lt x \leq end
pandas_bj.Between('from', 'to', True, True)  from \lt x \lt to

上のコードブロックでは、 s \lt x \lt eを使用しています。

pandas_bj.mergeの引数

sort

sort=Trueの箇所がありますが、すごく大事です。

sortを使用しないと、パフォーマンスが激落ちしますので、特別な理由がない限りはTrueにしましょう。 どうしてもソートできない列を扱う場合にも、部分的にソートを行うことで、全探索の範囲を絞り込むことができます。

how

how='inner'の箇所がありますが、これはINNER JOIN相当です。

how 挙動
left LEFT OUTER JOIN
right RIGHT OUTER JOIN
outer FULL OUTER JOIN
inner INNER JOIN

パフォーマンス

肝心のパフォーマンスですが、まずまずの結果になりました。

Xのレコード数 Yのレコード数 ソート有無 Time [s] 平均結合レコード数
100 1,000 False 0.1499 1.0
100 1,000 True 0.0614 1.0
1,000 10,000 False 8.1311 1.4669
1,000 10,000 True 0.3608 1.4669
10,000 100,000 True 3.843 6.0406
10,000 1,000,000 True 28.5253 51.8505

パフォーマンステストは、下記の様に実施しました。

  • ランダムにデータフレームXYを生成する。
  • XYのどちらも、区間(0, 100]をとる整数の列id1と 区間(0,50]をとる整数の列id2を持つ。
  • Xは、区間(0, 1000]をとる実数の列sと列eを持つ。常にe \geq s
  • Yは、区間(0, 1000]をとる実数の列vを持つ。
  • 結合条件として、 left_on=['id1', 'id2', pandas_bj.Between('s', 'e', False, False)], right_on=['id1', 'id2', 'v']を使う。

パフォーマンスの補足として、

  • 日時型の比較は、数値型の比較よりも遥かに時間がかかる
  • 範囲が広いなどの理由で結合結果が多くなると、時間がかかる

といった点から、測定時と実用時が大きくズレることがあります。


結局、手元のデータに対して、10分弱くらい処理時間がかかりましたが、今回は十分でした。

ソートなしの1,000件 VS 10,000件 に8秒かかっているところを見ると、 工夫しないでコードを書いていたら、相当な計算時間だったと思います。

まとめ

pandas BETWEEN JOINを使えば、

  • そこそこ高速なJOIN ON BETWEEN結合
  • \geq\gt\leq\ltもできる
  • 複数のAND条件を使える

もし、pandasでJOIN ON BETWEENをやりたい!そう思ったら、是非試してみてください。 需要があれば、実装の詳細についてもそのうち紹介しようと思います。

最後までありがとうございました!

宣伝

ABEJAでは、データ分析・高速処理大好きなエンジニアを募集しています! データ分析チームのPMも大絶賛募集中です!! www.wantedly.com

ABEJAが発信する最新テクノロジーに興味がある方は、是非ともブログの読者に!

ABEJAという会社に興味が湧いてきた方はWantedlyで会社、事業、人の情報を発信しているので、是非ともフォローを!! www.wantedly.com

ABEJAの中の人と話ししたい!オフィス見学してみたいも随時受け付けておりますので、気軽にポチッとどうぞ↓↓