BigQueryで範囲条件を使ったテーブル結合を行おうとして困った話

どんな話か

テーブルを結合する際、普通は以下のようにidを使いますよね。外部キーというやつでしょうか。

select ...
from A
inner join B
on A.id = B.id

一方、こういったIDによる結合ではなく、連続値を取るようなカラムと、その連続値に対してある範囲ではこれ、別の範囲ではこれ、といった風に逆関数的に別テーブルのレコードを結合させたい場合があります。
僕が遭遇したケースは重み付きの負例サンプリングです。乱数を発生させ、その値に応じて負例を結合することでランダムサンプリングを実現します。
そこで当初作成したのは以下のようなSQL。

select ...
from A
cross join B
where A.rnd >= B.lower and A.rnd < B.upper

しかしこれを実行すると、Aが10万行程度、Bが1万行程度でもリソース超過エラーとなってしまいました。
cross joinの処理が重すぎたようです。

そこで、以下のように書き換えることで解決を試みました。

• cross joinをinner joinに変更
• whereで指定していた条件をinner joinの結合条件として使用
  • BigQueryでは範囲条件を使ったテーブル結合がサポートされています。

select ...
from A
inner join B
on A.rnd >= B.lower and A.rnd < B.upper

しかし、これでもやはりリソース超過エラーが起きてしまいました。
実行の詳細を見てみると、内部的にはcross joinが実行されていました。書き方を変えても意味がなかったようです。

少し調べてみると、どうやら一致条件がないと結合による計算コストの抑制はできないみたいです。
範囲条件を満たすものを探すために、結局はAのレコード数×Bのレコード数分だけ総当たりで走査が必要となるということです。
一致条件をどうにかして導入しなければ解決しません。

解決方法

そこで以下のようにしました。
あらかじめテーブルBのlowerとupperの値を元にして整数部分との対応表Cを作成し、乱数の整数部分を使った一致条件を加えました。
これにより、テーブル全体ではなく整数部分を使ってある程度絞った上で部分的にcross joinが行われるような形となり、計算コストを大きく抑えることができました。

select ...
from A
inner join C
on cast(floor(A.rnd) as int64) = C.int
and (A.rnd >= C.lower and A.rnd < C.upper)

ただし、これはlowerやupperがある程度のスケールを持っていないといけません。極端な話、最大値が1なら先ほどと状況は変わらないです。

ともあれこれで解決できました。

おわりに

一般的なDBを使っている人には常識かもしれませんが、BigQueryの圧倒的なパワーに慣れきってしまった僕にはリソース超過でクエリが実行できないのは中々ない経験でした