Zenn
🐥

BigQueryで（擬似的に）集計関数や分析関数のUDFを作成する

2021/03/31に公開
SQL
BigQuery
UDF
tech

ユーザー定義関数

BigQueryではユーザー定義関数（UDF）を作成することができます。

UDFは基本的には1つの行に対して1つの値を返すような関数しか作成できないので、複数行に対して1つの値を返すような集計関数や、ウィンドウフレーム内の他の行の情報を使いつつ1つの行に対して1つの値を返すようなウィンドウ関数（分析関数）を作ることはできません。

しかし、UDFの引数にarray_agg関数の返り値を使うことで、擬似的に集計関数や分析関数のUDFを作成することができます。

集計関数を作る

BigQueryには平均値を計算するavg関数はありますが、中央値を計算する関数が用意されていません。
以下の記事では、中央値を計算するUDFの紹介と実装が載っています。

これを参考にして以前作成したジニ係数を計算するUDFの実装を紹介します。

CREATE OR REPLACE FUNCTION `project.dataset.gini`(arr ANY TYPE) AS (
(
with
    ranked as (
      select
          x
        , row_number() over (order by x desc) as rank
      from
        unnest(arr) as x
    )

select
  1 - 2 * sum(x * (rank - 1/2)) / count(*) / sum(x)
from
  ranked
)
);

ポイントは

です。通常の集計関数とは使い方も異なります。

select
    -- 通常の集計関数
    avg(x)
    -- UDFの集計関数
  , `dataset.gini`(array_agg(x))
from
  unnest(generate_array(0,100)) as x

このように、一度array_agg()で包んでからUDFに渡してあげることで機能します。

分析関数を作る

集計関数の作り方を応用して、分析関数のUDFを作成することもできます。
最近作成した、「同じ値が連続する間は同じIDで採番し、値が変わった時にIDをインクリメントする」という分析関数を紹介します。

CREATE OR REPLACE FUNCTION `project.dataset.context_id`(pref INT64, arr ANY TYPE) AS (
(
with
    wndw as (
      select 
          x
        , p
        , lag(x) over (order by p) as lagx
      from
        unnest(arr) as x
        with offset as p  -- position
    )
  , add_id as (
      select 
          *
        , countif(x != lagx) over (order by p) as id
      from
        wndw
    )
select
    id
from
  add_id
where 1 = 1
  and p = pref
)
);

集計関数の時と同様に、ARRAY型のデータを引数として受け取り、unnest(arr)で展開して扱っています。
集計関数の時の違いは最後の部分です。途中までは分析関数だけを使っているので複数行あり、最後のクエリで1行に絞って返しています。（集計関数の場合は自然と1行になる。）ここで目的の行に絞るために、arrの中の位置インデックスをprefとして引数に渡してもらうことを想定しています。

通常の分析関数との使い方の違いは、array_agg() over (...)で包んでからUDFに渡すことです。OVER句をつけることでarray_aggは分析集計関数になります。

select
    -- 通常の分析関数
    avg(x) over wndw
    -- UDFの分析関数
    -- row_numberは1-basedなので0-basedに直す
  , `dataset.context_id`(
      row_number() over wndw - 1,
      array_agg(x) over wndw
    )
from
  unnest([1,2,2,3,3,3,4,5,2,3]) as x
  with offset as p
window
  wndw as (order by p)

最後に

もっと賢いやり方を知っている方がいらっしゃれば是非コメントで教えてください！

Discussion