Skip to content

Instantly share code, notes, and snippets.

@AugustKarlstedt
Forked from vwrs/ndcg-bq-udf.sql
Created May 7, 2024 19:34
Show Gist options
  • Save AugustKarlstedt/7fe733f5813403c2430a9cf0440230d8 to your computer and use it in GitHub Desktop.
Save AugustKarlstedt/7fe733f5813403c2430a9cf0440230d8 to your computer and use it in GitHub Desktop.
calculate NDCG in BigQuery
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment