Use top-k sorted list builder with min/max-heaps for search.

Author: isoos

app/lib/search/heap.dart

@@ -0,0 +1,111 @@
+// Copyright (c) 2025, the Dart project authors.  Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+/// A data structure that keep satisfies the "heap property". This property
+/// dictates that in a max-heap, each node's value is greater than or equal
+/// to its children's values, and in a min-heap, each node's value is less
+/// than or equal to its children's values.
+///
+/// The provided comparator decides which kind of heap is being built.
+class Heap<T> {
+  final Comparator<T> _compare;
+  final _items = <T>[];
+
+  Heap(this._compare);
+
+  int get length => _items.length;
+
+  /// Insert a new [item] at the bottom of the heap, and move it up if needed.
+  void insert(T item) {
+    _items.add(item);
+    _up(_items.length - 1);
+  }
+
+  /// Replace the first entry of the with [item] and move it down if needed.
+  ///
+  /// WARNING: when maxLength is specified, the heap property will only applied
+  ///          for items before [maxLength] index.
+  void replaceFirst(T item, [int? maxLength]) {
+    assert(_items.isNotEmpty);
+    _items[0] = item;
+    _down(0, maxLength);
+  }
+
+  void _up(int index) {
+    final item = _items[index];
+    while (index > 0) {
+      final parentIndex = (index - 1) >> 1;
+      final parent = _items[parentIndex];
+      if (_compare(parent, item) <= 0) {
+        return;
+      }
+      _items[parentIndex] = item;
+      _items[index] = parent;
+      index = parentIndex;
+    }
+  }
+
+  void _down(int index, [int? maxLength]) {
+    maxLength ??= _items.length;
+    final item = _items[index];
+    while (index < maxLength) {
+      final leftIndex = (index << 1) + 1;
+      if (leftIndex >= maxLength) {
+        return;
+      }
+      var childIndex = leftIndex;
+      final rightIndex = leftIndex + 1;
+      if (rightIndex < maxLength &&
+          _compare(_items[leftIndex], _items[rightIndex]) > 0) {
+        childIndex = rightIndex;
+      }
+      if (_compare(item, _items[childIndex]) <= 0) {
+        return;
+      }
+      _items[index] = _items[childIndex];
+      _items[childIndex] = item;
+      index = childIndex;
+    }
+  }
+}
+
+/// Builds a sorted list of the top-k items using the provided comparator.
+///
+/// The algorithm uses min-heap to select the top-k items, and then builds
+/// a max-heap and uses heap sort to return the items in descending order.
+class TopKSortedListBuilder<T> {
+  final int _k;
+  final Comparator<T> _compare;
+  final Heap<T> _heap;
+
+  TopKSortedListBuilder(this._k, this._compare)
+      : _heap = Heap<T>((a, b) => -_compare(a, b));
+
+  void addAll(Iterable<T> items) {
+    for (final item in items) {
+      add(item);
+    }
+  }
+
+  void add(T item) {
+    final isAtCapacity = _heap.length >= _k;
+    if (isAtCapacity && _compare(_heap._items.first, item) <= 0) {
+      return;
+    }
+    if (isAtCapacity) {
+      _heap.replaceFirst(item);
+    } else {
+      _heap.insert(item);
+    }
+  }
+
+  Iterable<T> getTopK() {
+    for (var i = _heap.length - 1; i > 0; i--) {
+      final item = _heap._items[i];
+      _heap._items[i] = _heap._items[0];
+      _heap.replaceFirst(item, i);
+    }
+    return _heap._items;
+  }
+}

app/lib/search/mem_index.dart

@@ -9,6 +9,7 @@ import 'package:clock/clock.dart';
 import 'package:collection/collection.dart';
 import 'package:logging/logging.dart';
 import 'package:meta/meta.dart';
+import 'package:pub_dev/search/heap.dart';
 import 'package:pub_dev/service/topics/models.dart';
 import 'package:pub_dev/third_party/bit_array/bit_array.dart';
 
@@ -292,8 +293,8 @@ class InMemoryPackageIndex {
         }
         indexedHits = _rankWithValues(
           packageScores,
-          requiredLengthThreshold: query.offset,
           bestNameIndex: bestNameIndex ?? -1,
+          topK: query.offset + query.limit,
         );
         break;
       case SearchOrder.created:
@@ -521,12 +522,13 @@ class InMemoryPackageIndex {
 
   Iterable<IndexedPackageHit> _rankWithValues(
     IndexedScore<String> score, {
-    // if the item count is fewer than this threshold, an empty list will be returned
-    required int requiredLengthThreshold,
-    // When no best name match is applied, this parameter will be `-1`
+    /// When no best name match is applied, this parameter will be `-1`
     required int bestNameIndex,
+
+    /// Return (and sort) only the top-k results.
+    required int topK,
   }) {
-    int compare(int aIndex, int bIndex) {
+    final builder = TopKSortedListBuilder<int>(topK, (aIndex, bIndex) {
       if (aIndex == bestNameIndex) return -1;
       if (bIndex == bestNameIndex) return 1;
       final aScore = score.getValue(aIndex);
@@ -535,20 +537,13 @@ class InMemoryPackageIndex {
       if (scoreCompare != 0) return scoreCompare;
       // if two packages got the same score, order by last updated
       return _compareUpdated(_documents[aIndex], _documents[bIndex]);
-    }
-
-    final list = <int>[];
+    });
     for (var i = 0; i < score.length; i++) {
       final value = score.getValue(i);
       if (value <= 0.0 && i != bestNameIndex) continue;
-      list.add(i);
-    }
-    if (requiredLengthThreshold > list.length) {
-      // There is no point to sort or even keep the results, as the search query offset ignores these anyway.
-      return [];
+      builder.add(i);
     }
-    list.sort(compare);
-    return list.map((i) => IndexedPackageHit(
+    return builder.getTopK().map((i) => IndexedPackageHit(
         i, PackageHit(package: score.keys[i], score: score.getValue(i))));
   }
 

app/lib/search/token_index.dart

@@ -5,6 +5,7 @@
 import 'dart:math' as math;
 
 import 'package:meta/meta.dart';
+import 'package:pub_dev/search/heap.dart';
 import 'package:pub_dev/third_party/bit_array/bit_array.dart';
 
 import 'text_utils.dart';
@@ -313,21 +314,16 @@ class IndexedScore<K> {
   }
 
   Map<K, double> top(int count, {double? minValue}) {
-    final list = <int>[];
-    double? lastValue;
+    minValue ??= 0.0;
+    final builder = TopKSortedListBuilder<int>(
+        count, (a, b) => -_values[a].compareTo(_values[b]));
     for (var i = 0; i < length; i++) {
       final v = _values[i];
-      if (minValue != null && v < minValue) continue;
-      if (list.length == count) {
-        if (lastValue != null && lastValue >= v) continue;
-        list[count - 1] = i;
-      } else {
-        list.add(i);
-      }
-      list.sort((a, b) => -_values[a].compareTo(_values[b]));
-      lastValue = _values[list.last];
+      if (v < minValue) continue;
+      builder.add(i);
     }
-    return Map.fromEntries(list.map((i) => MapEntry(_keys[i], _values[i])));
+    return Map.fromEntries(
+        builder.getTopK().map((i) => MapEntry(_keys[i], _values[i])));
   }
 
   Map<K, double> toMap() {

app/test/search/heap_test.dart

@@ -0,0 +1,41 @@
+// Copyright (c) 2025, the Dart project authors.  Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+import 'package:pub_dev/search/heap.dart';
+import 'package:test/test.dart';
+
+void main() {
+  group('top-k sorted list', () {
+    int compare(int a, int b) => -a.compareTo(b);
+
+    test('no items', () {
+      final builder = TopKSortedListBuilder(5, compare);
+      expect(builder.getTopK().toList(), []);
+    });
+
+    test('single item', () {
+      final builder = TopKSortedListBuilder(5, compare);
+      builder.add(1);
+      expect(builder.getTopK().toList(), [1]);
+    });
+
+    test('three items ascending', () {
+      final builder = TopKSortedListBuilder(5, compare);
+      builder.addAll([1, 2, 3]);
+      expect(builder.getTopK().toList(), [3, 2, 1]);
+    });
+
+    test('three items descending', () {
+      final builder = TopKSortedListBuilder(5, compare);
+      builder.addAll([3, 2, 1]);
+      expect(builder.getTopK().toList(), [3, 2, 1]);
+    });
+
+    test('10 items + repeated', () {
+      final builder = TopKSortedListBuilder(5, compare);
+      builder.addAll([1, 10, 2, 9, 3, 8, 4, 7, 6, 5, 9]);
+      expect(builder.getTopK().toList(), [10, 9, 9, 8, 7]);
+    });
+  });
+}