Add HeapSet.peekn (#6947)

Author: gjoseph92

distributed/collections.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import heapq
+import itertools
 import weakref
 from collections import OrderedDict, UserDict
 from collections.abc import Callable, Hashable, Iterator
@@ -99,6 +100,20 @@ def peek(self) -> T:
                 return value
             heapq.heappop(self._heap)
 
+    def peekn(self, n: int) -> Iterator[T]:
+        "Iterator over the N smallest elements. This is O(1) for n == 1, O(n*logn) otherwise."
+        if n <= 0:
+            return  # empty iterator
+        if n == 1:
+            yield self.peek()
+        else:
+            # NOTE: we could pop N items off the queue, then push them back.
+            # But copying the list N times is probably slower than just sorting it
+            # with fast C code.
+            # If we had a `heappop` that sliced the list instead of popping from it,
+            # we could implement an optimized version for small `n`s.
+            yield from itertools.islice(self.sorted(), n)
+
     def pop(self) -> T:
         if not self._data:
             raise KeyError("pop from an empty set")

distributed/tests/test_collections.py

@@ -38,6 +38,9 @@ def __hash__(self):
     def __eq__(self, other):
         return isinstance(other, C) and other.k == self.k
 
+    def __repr__(self):
+        return f"C({self.k}, {self.i})"
+
 
 def test_heapset():
     heap = HeapSet(key=operator.attrgetter("i"))
@@ -131,6 +134,21 @@ def test_heapset():
     heap.add(cx)
     assert cx in heap
 
+    # Test peekn()
+    heap.add(cy)
+    heap.add(cw)
+    heap.add(cz)
+    heap.add(cx)
+    assert list(heap.peekn(3)) == [cy, cx, cz]
+    heap.remove(cz)
+    assert list(heap.peekn(10)) == [cy, cx, cw]
+    assert list(heap.peekn(0)) == []
+    assert list(heap.peekn(-1)) == []
+    heap.remove(cy)
+    assert list(heap.peekn(1)) == [cx]
+    heap.remove(cw)
+    assert list(heap.peekn(1)) == [cx]
+
     # Test resilience to failure in key()
     bad_key = C("bad_key", 0)
     del bad_key.i