enhance: [3.0] sort merged rerank results without per-comparison map lookups

Project: milvus-io/milvus | Language: Go | Concepts: algorithmic-optimization, data-structures, testing-strategy | Difficulty: mid
Link: https://github.com/milvus-io/milvus/pull/52885

The problem

The function that merges and sorts reranked search results (sortAndExtractResults) was inefficient: it sorted a slice of IDs, but each comparison during sorting looked up the corresponding score and location in a map. This led to many unnecessary map lookups, which are slow in Go and become a bottleneck for large result sets.

Why it's non-trivial

A naive fix might just try to cache some lookups or tweak the sort comparator, but that would not eliminate the repeated map accesses inside the sort. It’s important to maintain the exact sort order (including tie-breaking and handling of special values like NaN), and to keep the output slices (IDs, scores, locations) aligned after sorting. Any change must be thoroughly tested to ensure it is a drop-in replacement.

The approach

The PR introduces a new struct to carry all sorting information together and rewrites the sorting and extraction logic to avoid per-comparison map lookups. It also adds comprehensive tests and a benchmark.

1. Define a struct to bundle ID, score, and location:

+// scoredID carries the score and location alongside the id so that sorting does
+// not have to look them up. Sorting a slice of these with slices.SortStableFunc
+// also avoids sort.SliceStable's reflect-based swapper.
+type scoredID struct {
+	id    any
+	score float32
+	loc   idLocation
+}

This struct allows all necessary data for sorting to be stored together, so the sort function doesn't need to access maps repeatedly.

2. Build a slice of these structs instead of just IDs:

+	entries := make([]scoredID, 0, len(idScores))
+	for id, score := range idScores {
+		entries = append(entries, scoredID{id: id, score: score, loc: idLocs[id]})
+	}

This replaces the previous approach of building a slice of IDs and looking up the score/location in maps during sorting.

3. Sort using a custom comparator that matches the old logic:

+	sortScoredIDs(entries, descending)
...
+// lessScoredID is the ordering sortIDs used to express directly, kept as a
+// predicate so the three-way comparator below is equivalent to the previous
+// sort.SliceStable call by construction -- including for scores that compare
+// unequal in both directions, such as NaN.
+func lessScoredID(a, b scoredID, descending bool) bool {
+	if a.score != b.score {
+		if descending {
+			return a.score > b.score
+		}
+		return a.score < b.score
+	}
+	return compareIDs(a.id, b.id) < 0
+}
+
+// sortScoredIDs sorts by score with stable tie-breaking by ID.
+func sortScoredIDs(entries []scoredID, descending bool) {
+	slices.SortStableFunc(entries, func(a, b scoredID) int {
+		if lessScoredID(a, b, descending) {
+			return -1
+		}
+		if lessScoredID(b, a, descending) {
+			return 1
+		}
+		return 0
+	})
+}

This ensures the sort order is exactly the same as before, including tie-breaking by ID and correct handling of NaN scores.

4. Extract the sorted results into aligned slices:

+	ids := make([]any, len(entries))
+	scores := make([]float32, len(entries))
+	locs := make([]idLocation, len(entries))
+	for i, e := range entries {
+		ids[i] = e.id
+		scores[i] = e.score
+		locs[i] = e.loc
+	}

This guarantees the output slices are still aligned, as required.

5. Add comprehensive tests to verify the new logic matches the old:

+// sortAndExtractResults used to sort a []any with sort.SliceStable, looking each
+// score up in a map inside the comparator. It now sorts a slice of structs that
+// already carry the score. This pins the resulting order against the previous
+// formulation, including the tie-break by ID and the behavior for scores that
+// compare unequal in both directions (NaN).
+func TestSortAndExtractResultsMatchesPreviousOrdering(t *testing.T) {
...
+// sortAndExtractResults must keep ids, scores and locs aligned.
+func TestSortAndExtractResultsKeepsTuplesAligned(t *testing.T) {

These tests check that the new logic produces exactly the same order and alignment as the old logic, including edge cases.

6. Add a benchmark to measure performance:

+// MergeOp is the operator every rerank chain starts with, and it was the only
+// significant one without a benchmark.
+func BenchmarkMergeOp(b *testing.B) {
...
+	b.Run(bc.name, func(b *testing.B) {
+		...
+	})
+}

This provides data on the performance impact of the change.

From the review thread

No notable review discussion on this PR.

Takeaways

• Bundle related data into structs to avoid repeated map lookups during sorting. This reduces time complexity and improves cache locality.
• When refactoring core logic, write tests that pin new behavior to the old implementation, including edge cases and tie-breaking.
• Use stable sorting and explicit tie-breaking to ensure deterministic output, especially when results may have equal scores or special values like NaN.