BigQuery UDF(JS)で色々やってみる

はじめに

今回はBigQuery UDFについて、少し調べてみたのでそれについて書いてみます。
cloud.google.com

BigQuery UDF

UDF？

要はBigQueryのSQL上で使える関数を自分で作れる機能です。一般的な言語でいえば、関数定義だと思えば良いと思います。
BigQueryでは便利な関数は大概用意されています。
しかし、例えばあるプロジェクトの文脈でよく使う変換処理があったときに、それをいちいち書くと仮にバグが混入したときに全部チェックしなくてはいけません。
そこをUDF化しておくことで見通しよく管理できる感じですね。まぁ本当に関数定義です。

BigQueryのUDFはSQLもしくはJavaScriptで書くことができます。
せっかく作るなら、SQLだと書きにくいことを書けるJavaScriptだなというところで、JavaScript UDFをどのように作っていくのかを書こうと思います。

シンプルなUDF

チュートリアルなんかにもあるレベルですが、例えば足し算を行うUDFは次のようにかけます。

CREATE TEMP FUNCTION add(a FLOAT64, b FLOAT64) 
RETURNS INT64
LANGUAGE js
AS """
return a + b
""";
SELECT add(1,2)

見た目から単純ですね。一点注意するとすると

Standard SQL user-defined functions  |  BigQuery  |  Google Cloud

JavaScript で 64 ビット整数型はサポートされていないため、INT64 は入力データ型としてサポートされません。代わりに、FLOAT64 を使用して整数値を数値として表すか、STRING を使用して整数値を文字列として表します。

BigQuery では、JavaScript UDF の戻り型として INT64 がサポートされています。この場合、JavaScript 関数の本体で JavaScript の数値や文字列を返すことができます。この際、BigQuery によってこれらいずれかの型が INT64 に変換されます。

という制約があるので、型には気をつけましょう。

値ではなく、データのまとまりを返すUDF

先程の例では数値2つを受け取って1つの数値を返す形でしたが、特定の構造を持つ列構造を渡して、それぞれに決められた数字を足す事を考えると次のようにかけます。

CREATE TEMP FUNCTION add_a_b (x STRUCT<a FLOAT64, b FLOAT64>, addition FLOAT64) 
RETURNS STRUCT<a INT64, b INT64>
LANGUAGE js
AS """
return {
    a: x.a + addition,
    b: x.b + addition
};
""";
SELECT add_a_b(x, 1)
FROM (
    SELECT 1 AS a, 2 AS b
    UNION ALL
    SELECT 3 AS a, 4 AS c
    UNION ALL
    SELECT 5 AS a, 6 AS c
) x

STRUCTと書いたらJavaScriptのオブジェクトを思い浮かべればOKです。

結果をテーブル構造として返すUDF

例えば、横持ちのテーブルを縦持ちに変換するといったようなときは、入力の行と出力の行は1対1でなく、1対Nになります。
FROM句の中で元データっぽいものをUNION ALLで生成していますが、イメージ的にはid=店番号, date=日付単位で6日前までのデータを横持ちしている感じです。
よくExcelとかで見そうな感じのものですね（笑）

元ネタのSQL版は以下を参考にさせていただきました。
note.com

CREATE TEMP FUNCTION unpivot(json_string STRING, col_regex STRING)
RETURNS ARRAY<STRUCT<original STRING, value INT64>>
LANGUAGE js
AS """
row = JSON.parse(json_string)
const struct = []
const regex = new RegExp(col_regex)

for (key of Object.keys(row)) {
    if (key.match(regex)) {
        struct.push({
          original: key,
          value: row[key]
        })
    }
}

return struct;
""";
SELECT
  id,
  date,
  unpivot.*
FROM (
  SELECT 1 AS id, "2020-01-01" AS date, 1 AS c_day1, 2 AS c_day2, 3 AS c_day3, 4 AS c_day4, 5 AS c_day5, 6 AS c_day6
  UNION ALL
  SELECT 1 AS id, "2020-01-02" AS date, 2 AS c_day1, 1 AS c_day2, 2 AS c_day3, 3 AS c_day4, 4 AS c_day5, 5 AS c_day6
  UNION ALL
  SELECT 1 AS id, "2020-01-03" AS date, 5 AS c_day1, 2 AS c_day2, 1 AS c_day3, 2 AS c_day4, 3 AS c_day5, 4 AS c_day6
  UNION ALL
  SELECT 2 AS id, "2020-01-01" AS date, 6 AS c_day1, 5 AS c_day2, 4 AS c_day3, 3 AS c_day4, 2 AS c_day5, 1 AS c_day6
  UNION ALL
  SELECT 2 AS id, "2020-01-02" AS date, 3 AS c_day1, 6 AS c_day2, 5 AS c_day3, 4 AS c_day4, 3 AS c_day5, 2 AS c_day6
  UNION ALL
  SELECT 2 AS id, "2020-01-03" AS date, 4 AS c_day1, 3 AS c_day2, 6 AS c_day3, 5 AS c_day4, 4 AS c_day5, 3 AS c_day6
) x,
  UNNEST(unpivot(TO_JSON_STRING(x), 'c_')) unpivot

1点注意ですが、unpivot関数の1つ目の入力をSTRUCTではなくSTRINGにして、TO_JSON_STRINGにしていちいち行のデータをJSON化して渡しています。
これは、STRUCTで渡すこともできると思うのですが、そうしてしまうことで関数がその構造のテーブル以外受け付けられなくなってしまうためこのようにしています。
（STRUCTは後ろの<>内で指定した以外の列があるとエラーになります）

SQLは別に苦手意識はないのですが、ちょっと複雑な構造になると読むのが少し辛くなるので、この手の処理だとJavaScriptで書けるとなんとなく見やすい気がします。

おわりに

今日はBigQueryのJavaScript UDFについて浅くやったことをまとめてみました。
今回は全部TEMP FUNCTIONですが、TEMPにしないこともできるかつ速度面もメリットありらしいので、今度はそのへんも試してみたいですね。