Use top-k sorted list builder with heap for search. (#8826)

Author: isoos

app/lib/search/heap.dart

@@ -0,0 +1,112 @@
+// 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>[];
+  bool _isValid = true;
+
+  Heap(this._compare);
+
+  int get length => _items.length;
+
+  /// Collects [item] and adds it to the end of the internal list, marks the [Heap]
+  /// as non-valid.
+  ///
+  /// A separate operation may trigger the restoration of the heap proprerty.
+  void collect(T item) {
+    _items.add(item);
+    _isValid = false;
+  }
+
+  /// Collects [items] and adds them the end of the internal list, marks the [Heap]
+  /// as non-valid.
+  ///
+  /// A separate operation may trigger the restoration of the heap proprerty.
+  void collectAll(Iterable<T> items) {
+    _items.addAll(items);
+    _isValid = false;
+  }
+
+  /// Ensures that the tree structure below the [index] is a valid heap.
+  void _heapify(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;
+    }
+  }
+
+  /// (Re-)builds the heap property if needed.
+  void _buildHeapIfNeeded() {
+    if (_isValid) {
+      assert(_isValidHeap());
+      return;
+    }
+
+    if (_items.isEmpty) {
+      _isValid = true;
+      return;
+    }
+    for (var i = (_items.length >> 1); i >= 0; i--) {
+      _heapify(i);
+    }
+
+    assert(_isValidHeap());
+    _isValid = true;
+  }
+
+  /// Verifies the heap property is true for all items.
+  bool _isValidHeap() {
+    for (var i = 1; i < _items.length; i++) {
+      final parentIndex = (i - 1) >> 1;
+      if (_compare(_items[parentIndex], _items[i]) > 0) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  /// Creates a sorted list of the top-k items and removes them from the [Heap].
+  ///
+  /// The algorithm builds a max-heap in `O(N)` steps on the already collected items,
+  /// and then selects the top-k items by removing the largest item from the [Heap]
+  /// and restoring the heap property again in `O(k * log(N))` steps.
+  Iterable<T> getAndRemoveTopK(int k) sync* {
+    _buildHeapIfNeeded();
+    var remaining = k;
+    while (remaining > 0 && _items.isNotEmpty) {
+      yield _items[0];
+      remaining--;
+      final last = _items.removeLast();
+      if (_items.isEmpty) {
+        break;
+      }
+      _items[0] = last;
+      _heapify(0);
+    }
+    assert(_isValidHeap());
+  }
+}

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 heap = Heap<int>((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 [];
+      heap.collect(i);
     }
-    list.sort(compare);
-    return list.map((i) => IndexedPackageHit(
+    return heap.getAndRemoveTopK(topK).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 heap = Heap<int>((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;
+      heap.collect(i);
     }
-    return Map.fromEntries(list.map((i) => MapEntry(_keys[i], _values[i])));
+    return Map.fromEntries(heap
+        .getAndRemoveTopK(count)
+        .map((i) => MapEntry(_keys[i], _values[i])));
   }
 
   Map<K, double> toMap() {

app/test/search/heap_test.dart

@@ -0,0 +1,60 @@
+// 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 heap = Heap(compare);
+      expect(heap.getAndRemoveTopK(5).toList(), []);
+    });
+
+    test('single item', () {
+      final heap = Heap(compare);
+      heap.collect(1);
+      expect(heap.getAndRemoveTopK(5).toList(), [1]);
+    });
+
+    test('three items ascending', () {
+      final builder = Heap(compare);
+      builder.collectAll([1, 2, 3]);
+      expect(builder.getAndRemoveTopK(5).toList(), [3, 2, 1]);
+    });
+
+    test('three items descending', () {
+      final heap = Heap(compare);
+      heap.collectAll([3, 2, 1]);
+      expect(heap.getAndRemoveTopK(5).toList(), [3, 2, 1]);
+    });
+
+    test('10 items + repeated', () {
+      final heap = Heap(compare);
+      heap.collectAll([1, 10, 2, 9, 3, 8, 4, 7, 6, 5, 9]);
+      expect(heap.getAndRemoveTopK(5).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 b1 = Heap(compare)..collectAll(items);
+        final r1 = b1.getAndRemoveTopK(k).toList();
+        expect(r1, List.generate(k, (i) => length - 1 - i));
+
+        items.shuffle(r);
+        final b2 = Heap(compare)..collectAll(items);
+        final r2 = b2.getAndRemoveTopK(k).toList();
+        expect(r2, r1);
+      }
+    });
+  });
+}