Skip to content

Instantly share code, notes, and snippets.

@vwrs

vwrs/ndcg-bq-udf.sql

Last active May 7, 2024 19:34
Show Gist options
  • Save vwrs/37188fa064dad79414fd4db28929a6cc to your computer and use it in GitHub Desktop.
Save vwrs/37188fa064dad79414fd4db28929a6cc to your computer and use it in GitHub Desktop.
Download ZIP
calculate NDCG in BigQuery
Raw
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)
@vwrs
Copy link
Author

vwrs commented Oct 8, 2021

ref:

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment