unlock unlimited code runs for $79, limited time →
NOSSLETTER.techLIVE
back to OSSpedia
milvus-io/milvus
mid9/10
algorithmic-optimizationdata-structurestesting-strategy

enhance: 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/52276

The problem

The function responsible for sorting reranked search results, sortAndExtractResults, was inefficient. It sorted a slice of IDs by repeatedly looking up each ID's score in a map inside the sort comparator. This caused excessive CPU usage due to frequent map accesses and interface operations, dominating the runtime of rerank operations.

Why it's non-trivial

A naive fix might simply try to optimize the map access or change the map type, but that would only yield marginal improvements. The real bottleneck is the repeated map lookup and interface indirection inside the comparator, which is called many times during sorting. To optimize this properly, the data structure and sorting approach itself need to be changed, while still preserving the exact ordering and tie-breaking semantics—including for tricky cases like NaN scores and ties.

The approach

1. Introduce a struct to carry all relevant data together

Instead of sorting a slice of IDs and looking up scores and locations in a map for every comparison, the code now creates a slice of structs, each holding the ID, its score, and its location:

+type scoredID struct {
+	id    any
+	score float32
+	loc   idLocation
+}

Annotation: This struct groups together all the data needed for sorting, so each comparison has direct access to the score and location, eliminating repeated map lookups.

2. Build the slice of structs up front

The function now collects all the data into the slice of scoredID before sorting:

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

Annotation: This change means all the relevant data is together in memory, ready for efficient sorting.

3. Use a comparator that preserves previous ordering logic

The new sort uses slices.SortStableFunc with a three-way comparator that is directly derived from the old less-than predicate, ensuring identical ordering (including for 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
+}
+
+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
+	})
+}

Annotation: This preserves all previous tie-breaking and NaN handling, but now operates on fields, not map lookups.

4. Extract results after sorting

After sorting, the IDs, scores, and locations are extracted back out into separate slices, preserving their alignment:

+	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
+	}

Annotation: This guarantees the output slices stay in sync, as all fields travel together through the sort.

5. Add targeted benchmarks and tests

A new benchmark (BenchmarkMergeOp) was added to measure the performance of the operator in realistic scenarios. Two new tests ensure that the new sorting produces the same order as before (including edge cases), and that the tuple alignment is preserved.

From the review thread

No notable review discussion on this PR.

Takeaways

  • Avoid repeated map lookups or expensive operations inside tight loops or comparators. Instead, precompute or group the data to allow fast, direct access.
  • When refactoring for performance, preserve the exact ordering semantics by deriving new comparators directly from the old logic. This is especially important when subtle behaviors (like NaN handling or tie-breaking) are involved.
  • Comprehensive, property-based tests (e.g., comparing old and new output across edge cases and random data) are essential when changing core algorithms to ensure correctness is maintained.