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

Author: isoos

app/lib/search/heap.dart

@@ -0,0 +1,97 @@
+// 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;
+
+  void _pushDown(int index) {
+    final 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;
+    }
+  }
+
+  bool _isValidHeap() {
+    for (var i = 1; i < _items.length; i++) {
+      final parentIndex = (i - 1) >> 1;
+      if (_compare(_items[parentIndex], _items[i]) > 0) {
+        print(parentIndex);
+        print(_items);
+        return false;
+      }
+    }
+    return true;
+  }
+}
+
+/// 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 _Heap<T> _heap;
+
+  TopKSortedListBuilder(this._k, Comparator<T> compare)
+      : _heap = _Heap<T>(compare);
+
+  void addAll(Iterable<T> items) {
+    for (final item in items) {
+      add(item);
+    }
+  }
+
+  void add(T item) {
+    _heap._items.add(item);
+  }
+
+  /// Gets and removes the top-k items from the current list.
+  Iterable<T> getTopK() sync* {
+    if (_heap._items.isEmpty) {
+      return;
+    }
+    for (var i = (_heap._items.length >> 1); i >= 0; i--) {
+      _heap._pushDown(i);
+    }
+    assert(_heap._isValidHeap());
+    var count = _k;
+    while (count > 0 && _heap._items.isNotEmpty) {
+      yield _heap._items[0];
+      count--;
+      final last = _heap._items.removeLast();
+      if (_heap._items.isEmpty) {
+        break;
+      }
+      _heap._items[0] = last;
+      _heap._pushDown(0);
+    }
+  }
+}

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,56 @@
+// 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 'dart:math';
+
+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]);
+    });
+
+    test('randomized verification', () {
+      for (var i = 0; i < 1000; i++) {
+        final r = Random(i);
+        final length = 1000 + r.nextInt(1000);
+        final k = 10 + r.nextInt(200);
+        final items = List.generate(length, (i) => i);
+        final builder = TopKSortedListBuilder(k, compare);
+        builder.addAll(items);
+        final result = builder.getTopK().toList();
+        expect(result, List.generate(k, (i) => length - 1 - i));
+      }
+    });
+  });
+}