ブログ・ア・ラ・クレーム

技術的なメモとかライフログとか。

BigQuery SQL UDF の挙動を色々確認する

BigQuery ではユーザ定義関数(UDF) を作ることができる。 これを使って、よく使われる式や関数呼び出しの組み合わせを名前付けして再利用できる。

cloud.google.com

UDF は SQL と JavaScript の二種類の言語による記述が可能で、後者は色々なトリッキーな利用例も世に出ているので知っている人も多いかも知れない。 例えば WebAssembly のコードを実行するとかである。

medium.com

本記事では 地味な方 素直な SQL での記述による UDF について色々挙動を確認してみる。

戻り値の型を推論してくれる

UDF は例えばシンプルなものだと以下のように記述できる。

CREATE TEMP FUNCTION
  zero()
  RETURNS INT64 AS (0);
SELECT
  zero()

ここで戻り値の型は省略することができる（勝手に解決してくれる）

CREATE TEMP FUNCTION
  zero()
  AS (0);
SELECT
  zero()

STRUCT 型の戻り値も推論してくれる。が、もちろん各フィールド名の情報はないのでダミーの値が補完されてしまう。

CREATE TEMP FUNCTION
  user()
  AS ((1, "taro"));
SELECT
  user()

[
  {
    "f0_": {
      "_field_1": "1",
      "_field_2": "taro"
    }
  }
]

RETURNS で型を明示するとそこに記述されたフィールド名を拾ってくれる。

CREATE TEMP FUNCTION
  user()
  RETURNS STRUCT<id INT64, name STRING>
  AS ((1, "taro"));
SELECT
  user()

[
  {
    "f0_": {
      "id": "1",
      "name": "taro"
    }
  }
]

UDF から UDF を呼ぶ際はすでに定義済みである必要がある

UDF から別の UDF を呼び出すこともできる。

CREATE TEMP FUNCTION
  callee() AS (42);
CREATE TEMP FUNCTION
  caller() AS (callee());
SELECT
  caller()

この定義順序を逆にするとクエリの実行が通らなくなる。 Function not found: callee; failed to parse CREATE [TEMP] FUNCTION statement at [5:16] というエラーが出てしまう。

CREATE TEMP FUNCTION
  caller() AS (callee());  -- <- Function not found :(
CREATE TEMP FUNCTION
  callee() AS (42);
SELECT
  caller()

UDF の再帰呼び出しをすることもできない。この例では Function not found: fib; failed to parse CREATE [TEMP] FUNCTION statement at [10:7] と怒られてしまう。

CREATE TEMP FUNCTION
  fib(n INT64) AS (
  IF
    ( n > 2,
      fib(n - 1) + fib(n - 2),
      1 ) );
SELECT
  fib(2)

だったら一度ダミーの UDF を定義してみよう！ということで一時 UDF じゃなく永続 UDF を、一度ダミーのものを定義したあと REPLACE してみる。 このクエリはそれ自体は valid と言われ実行可能である。

-- dummy
CREATE OR REPLACE FUNCTION
  udfs.fib(n INT64) AS (1);
 
CREATE OR REPLACE FUNCTION
  udfs.fib(n INT64) AS (
  IF
    ( n > 2,
      udfs.fib(n - 1) + udfs.fib(n - 2),
      1 ) );
SELECT
  udfs.fib(2)

が、これは実行してみると Query error: Too many nested views/persistent user-defined functions or possible circular reference of views/persistent user-defined functions referenced in query. Only 16 levels of nested views/persistent user-defined functions are allowed. at [11:1] と怒られてしまう。 UDF では 16 チェーンまでしか UDF 呼び出しをできない制限がありそれに引っかかっているように見えるエラーメッセージである。この例では udfs.fib(2) を呼び出してもそもそも再帰呼び出しされない気がするが・・・。

cloud.google.com

頑張ればループ処理を書ける

UDF では前述の通り再帰呼び出しできず、ループするような制御構文もサポートされていない。 が、 ARRAY 型の値と ARRAY 型関連関数を使うとループ、リストの各要素に対する繰り返し処理的なことができる。 ミソとなるのがみんな大好き UNNEST() でテーブルを手にいれることができることと、 ARRAY() の引数はサブクエリを取れる辺りにある！

CREATE TEMP FUNCTION
  array_twice(n INT64) AS (ARRAY(
    SELECT
      AS STRUCT (
      SELECT
        ARRAY_AGG(vv)
      FROM (
        SELECT
          v * 2  -- <- like a map function
        FROM
          UNNEST(GENERATE_ARRAY(1, n)) AS v) AS vv)));
SELECT
  array_twice(100)

テーブルを参照することもできる

上記より SELECT 文が打てることを確認できたので FROM 句に実テーブルを渡してみる。 このクエリも合法で実行することができる！

CREATE TEMP FUNCTION
  do() AS (ARRAY(
    SELECT
      AS STRUCT (
      SELECT
        ARRAY_AGG(vv)
      FROM (
        SELECT
          LENGTH(repo_name)  -- <- like a map function
        FROM
          `bigquery-public-data.github_repos.licenses` -- <- refer to a table
        LIMIT
          10) AS vv)));
SELECT
  do()

ドキュメントにテーブル参照に関する記述があるので、この挙動は一応想定されたもののはず。（もしかしたら他のまっとうな参照方法があるのかも）

一意の UDF とテーブル参照を合わせたクエリあたりの最大数 - 1,000。完全な展開後に、UDF ごとに一意のテーブルと UDF を合わせて 1,000 個まで参照できます。

cloud.google.com