ビット順反転処理を BigQuery の SQL UDF で書く

Bit reversal - ビット順反転とは

タイトルには BigQuery とあるのにいきなり Spanner の話から始まります。Spanner には BIT_REVERSE という関数があります。これはビットの並び順を反転させるもので、たとえば Spanner 上で SELECT BIT_REVERSE(3, TRUE) をクエリすると、INT64 の 3 をビット列（2進数）で見たときに、最上位の符号ビット以外のビットの並び順を反転した値である 6917529027641081856 値が返ってきます。

Spanner の BIT_REVERSE() 関数によるビット順反転のイメージ

一般に RDBMS の主キー（PK）には連番が採用されることが多いと思います。しかし連番のように単調増加するキーというのは、Spanner のようなキーの範囲で自動シャーディングを行う DB では、INSERT 時に特定のシャードに更新負荷が偏ってしまいがちになるため、単調増加ではなく分散して採番される値がよいとされています。

そのため Spanner には連番を ビット順反転をして 採番するシーケンスの仕組みがあり、これを主キーに設定することができます。例えば 1 → 2 → 3 ... → 100 → 101 という連番は、（符号ビット残しで）ビット順反転させると4611686018427387904 → 2305843009213693952 → 6917529027641081856 ... → 1369094286720630784 → 5980780305148018688 という値になります。なぜビット順を反転させると値が分散するのかについてはこちらの記事をご覧ください。

なぜ BigQuery でビット順反転処理をしたいのか

Spanner へデータを流し込む方法はいくつもありますが、その際に BigQuery からデータを流し込みたい場合もあります。最近では BigQuery から Spanner へ直接リバース ETL ができるようになったこともあり、Spanner のデータを BigQuery 側でバッチ処理を行って Spanner に書き戻すのも容易になりました。その際にもちろん Spanner 側の BIT_REVERSE 関数を使ってもいいのですが、バッチ処理が得意な BigQuery 側でやってしまいたい場合もあります。

他にも MySQL など他の DB からデータ マイグレーションを行う際に、もともと MySQL 側で連番となっている値を、ビット順反転して Spanner に持ち込みたいこともあるでしょう。逆に Spanner 側のデータを MySQL に書き戻したい場合に、ビット順を再反転させてもとの連番に戻したいことがあるかもしれません。

もちろんこのようなデータ パイプラインを組む方法はいくつもありますが、BigQuery を経由することでこのようなパイプラインをシンプルに組める場合も多いです。ところが BigQuery にはこの BIT_REVERSE 関数相当のものがありません。しかしながら一般的なビット演算は一通り可能なので、これらを使って SQL UDF として実装してみました。

SQL UDF による自前 BIT_REVERSE 関数

Spanner の BIT_REVERSE 関数は INT64 と BOOL を引数に取ります。BIT_REVERSE(num, FALSE) のときは num を単純にビット順反転しますが、BIT_REVERSE(num, TRUE) のときは符号ビット（最上位ビット、MSB）を保持します、つまり MSB 以外の 63 bit でビット順反転をします。最近　Ruby で Spanner の BIT_REVERSE を実装した記事も出ていたので参考になると思います。その記事でも触れられていますが、具体的なビット演算の方法は、Spanner Emulator の実装をみると分かります。

https://github.com/GoogleCloudPlatform/cloud-spanner-emulator/blob/v1.5.28/common/bit_reverse.cc#L23-L38

ではこの処理を参考にしながら BigQuery のビット演算だけで書いてみましょう。こんな感じです。

-- 32 bit の bit reverse を行う UDF
CREATE OR REPLACE FUNCTION <your_dataset>.bit_reverse32(value INT64)
RETURNS INT64
AS
((
  WITH `step1` AS ( SELECT ((value >> 1) & 0x55555555) | ((value & 0x55555555) << 1) AS n),
  `step2` AS ( SELECT ((n >> 2) & 0x33333333) | ((n & 0x33333333) << 2) AS n FROM `step1`),
  `step3` AS ( SELECT ((n >> 4) & 0x0F0F0F0F) | ((n & 0x0F0F0F0F) << 4) AS n FROM `step2`),
  `step4` AS ( SELECT ((n >> 8) & 0x00FF00FF) | ((n & 0x00FF00FF) << 8) AS n FROM `step3`),
  `step5` AS (SELECT (n >> 16)  | (n << 16) AS n FROM `step4`)
  SELECT CAST(n & 0xFFFFFFFF AS INT64)  FROM `step5`
));

-- 64 bit の bit reverse を行う UDF、内部で bit_reverse32 を利用している
CREATE OR REPLACE FUNCTION <your_dataset>.bit_reverse64(value INT64)
RETURNS INT64
AS (
  <your_dataset>.bit_reverse32(value >> 32) | <your_dataset>.bit_reverse32(value & 0xffffffff) << 32
);

-- Spanner の BIT_REVERSE 関数と同じ入出力となる spanner_bit_reverse UDF
CREATE OR REPLACE FUNCTION <your_dataset>.spanner_bit_reverse(value INT64, preserve_sign BOOL)
RETURNS INT64
AS (
  CASE preserve_sign
    WHEN TRUE THEN (<your_dataset>.bit_reverse64(value) >> 1) | (value & -0x8000000000000000)
    WHEN FALSE THEN <your_dataset>.bit_reverse64(value)
  END
);

3 つの SQL UDF で構成されています。bit_reverse32(value INT64) は INT64 を引数にとり、そのうち下位32 bit についてビット順反転をします。bit_reverse64(value INT64) は INT64 を引数にとり、内部で bit_reverse32() を呼び出し 64bit 全体をビット順反転させます。そして、spanner_bit_reverse(value INT64, preserve_sign BOOL) は Spanner の BIT_REVERSE 関数と同じ挙動になっており、preserve_sign=TRUE で実行すると、符号ビットを残す処理を行います。

ビット演算の解説が趣旨ではないので詳細な説明は省きますが、いくつかこの記事を読んだ人が気になりそうなポイントを補足します。bit_reverse32(value INT64) では、最初は隣り合った 2 bit ずつ、次に 4 bit ずつ... とビットを入れ替えて、最終的に 32 bit のビット順反転をしています。最後に 0xFFFFFFFF で CAST しているのは、BigQuery では実際には INT64 で処理しているので、下位 32 bit の反転結果だけを返すためです。spanner_bit_reverse(value INT64, preserve_sign BOOL) で preserve_sign=TRUE の場合の処理は、もともとあった符号ビットを戻す処理をしています。細かい部分では -0x8000000000000000 は INT64 の最小値であり、ビット列的には 64 bit のうち最上位の符号ビットのみが立っており他は全部 0 の値です。つまりこれと元の値との AND を取ることで符号ビットだけを取り出せます。これは本来のビット表現的には頭につけてるマイナス符号は不要で 16 進数で 0x8000000000000000 と書きたいところですが、マイナス符号をつけずにリテラルを記述すると BigQuery 的にはこの 16 進数は扱える範囲外としてみなされてしまい Invalid hex integer literal とエラーになってしまうため、しかたなくマイナス符号をつけています。

実行

さて早速使ってみましょう。いくつかの整数値をビット順反転させてみます。

SELECT
  num,
  <your_dataset>.spanner_bit_reverse(num, TRUE) AS bit_reversal_bq,
FROM
  UNNEST([
    0, 1, 3, 5, 5764607523034234880, 6917529027641081856, 4611686018427387904,
    -1, -9223372036854775808,9223372036854775807]) AS num;

BigQuery でビット順反転

続いて Spanner の BIT_REVERSE 関数を連携クエリ経由で実行し、自前の spanner_bit_reverse 関数と結果を比較してみます。

SELECT
  num,
  bit_reverse_spanner,
  <your_dataset>.spanner_bit_reverse(num, TRUE) AS bit_reverse_bq,
FROM EXTERNAL_QUERY(
  "<your_spanner_connection>",
  '''
  SELECT num, BIT_REVERSE(num, TRUE) AS bit_reverse_spanner
  FROM UNNEST([
    0, 1, 3, 5, 5764607523034234880, 6917529027641081856, 4611686018427387904,
    -1, -9223372036854775808,9223372036854775807]) AS num;
  '''
);

Spanner の BIT_REVERSE との結果比較

気になる速度ですが、1 億レコードの連番を生成してそれに対して自前 spanner_bit_reverse 関数を通してビット順反転させてみました。ちゃんとは計測していませんが 1 億レコードの生成（GENERATE_ARRAY を再帰 CTE で回しています） 含めて 1分半程度で終わりました。