calculate NDCG in BigQuery

ndcg-bq-udf.sql

create temp function ndcg(rels array<int64>, k int64) returns float64 as (
    (
      with gain as (
        select
          (pow(2, rel) - 1) / log(i + 2) as g,
          (pow(2, rel_sorted) - 1) / log(i + 2) as g_ideal
        from
          unnest(rels) as rel with offset as i
          left join unnest((
            select
              array_agg(rel order by rel desc)
            from
              unnest(rels) as rel
          )) as rel_sorted with offset as i
          using(i)
        where
          i + 1 <= k
      )
      select
        round(safe_divide(sum(g), sum(g_ideal)), 5)  /* DCG / IDCG */
      from
        gain
    )
  );
select ndcg([0, 1, 0], 3), ndcg([1, 0, 0], 3), ndcg([1, 0, 1, 0, 0], 5), ndcg([0, 0, 1, 1, 1], 5)

ref:

• https://github.com/apache/spark/blob/v3.2.0-rc7/mllib/src/main/scala/org/apache/spark/mllib/evaluation/RankingMetrics.scala#L143
• https://en.wikipedia.org/wiki/Discounted_cumulative_gain