苦しみポイント

先日、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_onやright_onの中に列名を入れる代わりに、pandas_bj.Betweenを入れればOKです。

left_onやright_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`の引数

区間の小さい方の区切りの列名
区間の大きい方の区切りの列名
左閉区間か左開区間か
右閉区間か右開区間か

例：

引数	区間
`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

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

ランダムにデータフレームXとYを生成する。
X、Yのどちらも、区間(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']を使う。