Problem 123 · Evaluation

Retrieval Metrics: Recall@k, MRR and nDCG@k

MediumAI engineering

01Problem

Given a ranked list of retrieved doc ids and graded relevance labels `qrels: dict[doc, int]` (0 = not relevant), implement `recall_at_k`, `mrr` (reciprocal rank of the first relevant doc, 0 if none) and `ndcg_at_k` with gain `2**rel - 1` and discount `log2(rank + 1)`. Duplicate ids in the ranking count only once. Then `evaluate(runs, qrels, k)` averages each metric over queries, **skipping** queries with no relevant documents and reporting how many were skipped.

02Key insight

Each metric answers a different question: recall@k asks 'did the evidence reach the reader model at all' (the one that matters most for RAG), MRR asks 'how far down is the first hit' (matters for single-answer UX), and nDCG rewards putting the most relevant documents highest with graded labels. Queries with no relevant docs make recall and nDCG undefined; counting them as 0 or 1 silently shifts the average, so skip and report them. Deduplicate retrieved ids, or a retriever returning the same chunk twice gets double credit.

03Solution

solution.pypython51 lines

import math


def _dedupe(ranking: list[str]) -> list[str]:
    seen: set[str] = set()
    return [d for d in ranking if not (d in seen or seen.add(d))]


def recall_at_k(ranking: list[str], qrels: dict[str, int], k: int) -> float | None:
    relevant = {d for d, r in qrels.items() if r > 0}
    if not relevant:
        return None
    return len(set(_dedupe(ranking)[:k]) & relevant) / len(relevant)


def mrr(ranking: list[str], qrels: dict[str, int]) -> float:
    for rank, d in enumerate(_dedupe(ranking), start=1):
        if qrels.get(d, 0) > 0:
            return 1.0 / rank
    return 0.0


def ndcg_at_k(ranking: list[str], qrels: dict[str, int], k: int) -> float | None:
    ideal = sorted((r for r in qrels.values() if r > 0), reverse=True)[:k]
    if not ideal:
        return None
    dcg = sum((2 ** qrels.get(d, 0) - 1) / math.log2(i + 2) for i, d in enumerate(_dedupe(ranking)[:k]))
    idcg = sum((2 ** r - 1) / math.log2(i + 2) for i, r in enumerate(ideal))
    return dcg / idcg


def evaluate(runs: dict[str, list[str]], qrels: dict[str, dict[str, int]], k: int) -> dict[str, float]:
    rows = []
    skipped = 0
    for qid, ranking in runs.items():
        rel = qrels.get(qid, {})
        r = recall_at_k(ranking, rel, k)
        if r is None:
            skipped += 1
            continue
        rows.append((r, mrr(ranking, rel), ndcg_at_k(ranking, rel, k)))
    n = len(rows)
    if n == 0:
        return {"recall": 0.0, "mrr": 0.0, "ndcg": 0.0, "queries": 0, "skipped": skipped}
    return {
        "recall": sum(r[0] for r in rows) / n,
        "mrr": sum(r[1] for r in rows) / n,
        "ndcg": sum(r[2] for r in rows) / n,
        "queries": n,
        "skipped": skipped,
    }

Try it first.

TimeO(n + |qrels| log |qrels|) per query

SpaceO(n)

tests.py · 10 passing

tests.pypython15 lines

ranking = ["d1", "d2", "d3", "d4"]
rel = {"d2": 3, "d4": 1, "d9": 2, "d5": 0}
assert recall_at_k(ranking, rel, 3) == 1 / 3
assert recall_at_k(ranking, rel, 4) == 2 / 3
assert mrr(ranking, rel) == 0.5
L = math.log2(3)
assert abs(ndcg_at_k(ranking, rel, 3) - (7 / L) / (7 + 3 / L + 1 / 2)) < 1e-12
assert ndcg_at_k(["d2", "d9", "d4"], rel, 3) == 1.0
assert recall_at_k(["d2", "d2", "d2"], rel, 3) == 1 / 3
assert mrr(["x", "x", "d4"], rel) == 0.5
assert recall_at_k(ranking, {"d1": 0}, 3) is None and ndcg_at_k(ranking, {}, 3) is None
res = evaluate({"q1": ranking, "q2": ["d9"], "q3": ["zz"]},
               {"q1": rel, "q2": {"d9": 1}, "q3": {}}, 3)
assert res["queries"] == 2 and res["skipped"] == 1
assert abs(res["recall"] - (1 / 3 + 1) / 2) < 1e-12 and res["mrr"] == 0.75

04Staff-level follow-up

For RAG, recall@k at the k you actually pass to the model is the headline metric, but labels are incomplete: unjudged documents are treated as irrelevant, which penalises a new retriever that finds good unlabelled docs. Pool results from several systems for judging, or use an LLM judge on unjudged hits and audit a sample. Then connect retrieval metrics to answer quality, since a recall gain that does not move end-to-end accuracy is not worth its latency.