What you will build

Retrieving the right documents is only part of a RAG pipeline. When a prompt can include just the first few results, their order matters. Calculate normalized discounted cumulative gain at a cutoff using hand-labeled document fixtures. This evaluates ranking quality independently of generation and performs no retrieval service or model call.

We choose linear gain: a grade of three contributes three times the undiscounted gain of a grade of one. Some systems use exponential gain instead. Record the convention alongside every metric so scores from different implementations are not compared accidentally.

Prerequisites

Python 3.12 or later, with no packages. Prepare a dictionary of document IDs to integer relevance grades from zero through three and an ordered list of retrieved IDs. Save this as ndcg.py.

from math import log2, isclose

def ndcg(ranking, judgments, k):
    if type(k) is not int or k < 1:
        raise ValueError('positive integer k required')
    if len(set(ranking)) != len(ranking):
        raise ValueError('duplicate document')
    if any(doc not in judgments for doc in ranking):
        raise ValueError('unjudged document')
    if any(type(g) is not int or not 0 <= g <= 3 for g in judgments.values()):
        raise ValueError('grades must be integers from 0 to 3')
    def dcg(grades):
        return sum(g / log2(rank + 2) for rank, g in enumerate(grades))
    ideal = dcg(sorted(judgments.values(), reverse=True)[:k])
    if ideal == 0:
        return None
    return dcg([judgments[d] for d in ranking[:k]]) / ideal

judgments = {'guide': 3, 'faq': 2, 'old': 0}
score = ndcg(['old', 'faq', 'guide'], judgments, 3)
assert isclose(score, (2 / log2(3) + 3 / 2) / (3 + 2 / log2(3)))
assert ndcg(['guide', 'faq'], judgments, 2) == 1.0
assert ndcg([], judgments, 2) == 0.0
assert ndcg(['x'], {'x': 0}, 2) is None
for ranking, grades, k in [(['guide', 'guide'], judgments, 2),
                            (['missing'], judgments, 2),
                            ([], {'x': -1}, 2), ([], judgments, True)]:
    try:
        ndcg(ranking, grades, k)
    except ValueError:
        pass
    else:
        raise AssertionError('invalid fixture accepted')
print(f'linear_ndcg@3={score:.4f}')
print('fixtures passed')

Run and expected output

python3 ndcg.py
linear_ndcg@3=0.6480
fixtures passed

How it works

Discounted cumulative gain divides each relevance grade by a base-two logarithm of its rank plus one. The first document receives full credit; subsequent results contribute less. The ideal score sorts every judged grade from highest to lowest and applies the same cutoff. Dividing observed gain by ideal gain makes a perfect ranking score one.

The scikit-learn nDCG reference describes this linear-score formulation. Our function takes an already ordered list rather than prediction scores. It does not implement score-tie averaging. The upstream retriever must provide a deterministic order if repeatable evaluation is required.

With the irrelevant old document first, the fixture scores about 0.6480. Moving guide and faq into the first two positions gives one at cutoff two. Ideal gain comes from the complete judgment dictionary, not just retrieved results. Otherwise a retriever could receive perfect credit after omitting the best document.

Edge cases and scope

Duplicate IDs are rejected because repeating one strong document should not earn multiple contributions. Unjudged returned IDs are also rejected; they are not automatically treated as irrelevant. If evaluation uses a judged subset, document how that subset was created and monitor how many results lack judgments.

An empty ranking with positive ideal gain returns zero. A query with no positive judged grade returns None, meaning its metric is undefined. When aggregating queries, count and disclose these exclusions. The cutoff must be a positive integer; booleans are rejected even though Python treats them as integer subclasses.

Troubleshooting

If your result differs from another library, check linear versus exponential gain first, followed by cutoff, judgment universe and tie handling. Do not round document gains before summing. A high ranking score does not demonstrate that a generated answer is grounded or correct; test those stages separately. Keep judgment revisions under version control because changed labels can move the metric without a retriever change.

Primary references

Continue with AI tutorials and all Java Samples categories.