Maintaining PriorityQueue Order With Java Streams

The tricky thing about working with PriorityQueues is that, ironically, they do not always behave according to the PriorityQueue semantics.

PriorityQueue Traversal

If we have a look at PriorityQueue.iterator() documentation, we'll see that, unintuitively, iterator() is not traversing the queue according to its priority order:

Returns an iterator over the elements in this queue. The iterator does not return the elements in any particular order.

The same behavior can be noticed when trying to use Java Stream API for processing PriorityQueue's elements using the instance obtained by the stream() method — the Stream instance depends on the Spliterator instance which does not guarantee the traversal order according to the priority.

PriorityQueue<String> queue = new PriorityQueue<>(comparing(String::length));
List<String> content = Arrays.asList("1", "333", "22", "55555", "4444");
queue.addAll(content);

assertThat(queue.stream())
  .containsExactlyElementsOf(content);

We can see that the insertion order gets preserved regardless of the fact is that our queue should be providing Strings according to their length.

There are two solutions for dealing with this.

Solution #1

To take elements from the queue according to their priority order, we can use the poll() method.

In this case, we can generate a Stream from consecutive elements returned by poll() using Stream.generate() method:

List<String> result = Stream.generate(queue::poll)
  .limit(queue.size())
  .collect(Collectors.toList());

assertThat(result)
  .containsExactly("1", "22", "333", "4444", "55555");
assertThat(queue)
  .isEmpty();

The problem with this implementation is that it's not concurrent-modification-friendly. Until Java 9 gets released, we can't terminate the generated Stream dynamically so we need to rely on limiting the Stream size to the queue size.

After consuming a Stream instance, we end up with a modified queue — the poll() method removes queue elements until it's empty.

Eventually, this functionality can be extracted to the separate utility method:

static <T> Stream<T> drainToStream(PriorityQueue<T> queue) {
    Objects.requireNonNull(queue);
    return Stream.generate(queue::poll)
      .limit(queue.size());
}

Java 9

Since Java 9, it'll be possible to rewrite it in a concurrent-friendly way:

static <T> Stream<T> drainToStream(PriorityQueue<T> queue) {
    Objects.requireNonNull(queue);
    return Stream.generate(queue::poll)
      .takeWhile(Objects::nonNull)
}

Solution #2

The other approach involves just sorting the Stream instance, obtained by calling the stream() method, using the same comparator that the queue uses. We need to remember that this will work as long as the queue was initialized using a custom comparator:

List<String> result = queue.stream()
  .sorted(queue.comparator())
  .collect(Collectors.toList());

assertThat(result)
  .containsExactly("1", "22", "333", "4444", "55555");
assertThat(queue)
  .isNotEmpty();

If we store Comparable objects in the queue and depend on their natural order, this becomes even simpler because we do not need to pass the Comparator instance around:

PriorityQueue<String> queue = new PriorityQueue<>();
queue.addAll(Arrays.asList("1", "333", "22", "55555", "4444"));

List<String> result = queue.stream()
  .sorted()
  .collect(Collectors.toList());

assertThat(result)
  .containsExactly("1", "22", "333", "4444", "55555");
assertThat(queue)
  .isNotEmpty();

The important part of this implementation is that after creating a Stream instance, our original queue remains intact.

Eventually, this functionality can be extracted to the separate utility method:

static <T> Stream<T> asStream(PriorityQueue<T> queue) {
    Objects.requireNonNull(queue);
    Comparator<? super T> comparator = queue.comparator();
    return comparator != null
      ? queue.stream().sorted(comparator)
      : queue.stream().sorted();
}

Conclusion

The priority order of the PriorityQueue is not preserved when iterating/traversing so, essentially, we need to create our Stream instances ourselves.

The working code snippets can be found on GitHub.