Limit a stream by a predicate

java

stream-engineering

java-8

performance

byAnton Shumikhin·Sep 7, 2024

Supercharge your streams in Java with the takeWhile method. When it encounters a false from the predicate, it wraps up its operation:

List<String> filtered = Stream.of("first", "second", "halt", "third")
                              .takeWhile(s -> !"halt".equals(s))
                              .collect(Collectors.toList());
// Output: ["first", "second"]

Consider this the direct approach to putting a leash on your stream.

Java 8 workaround

Pre-JDK 9 DIY `takeWhile`

In Java 8, we can roll up our sleeves and craft a similar functionality with Spliterator and StreamSupport:

public static <T> Stream<T> takeWhile(Stream<T> stream, Predicate<T> predicate) {
    Spliterator<T> splitr = stream.spliterator();
    return StreamSupport.stream(
        new Spliterators.AbstractSpliterator<T>(splitr.estimateSize(), 0) {
            boolean stillWorking = true; // coffee breaks are not supported yet

            @Override
            public boolean tryAdvance(Consumer<? super T> action) {
                if (stillWorking) {
                    boolean hasNext = splitr.tryAdvance(elem -> {
                        if (predicate.test(elem)) {
                            action.accept(elem);
                        } else {
                            stillWorking = false; // time for that coffee break
                        }
                    });
                    return hasNext && stillWorking;
                }
                return false;
            }
        }, false); // parallel streams don't support tea breaks too
}

StreamEx to the rescue

StreamEx library does a good "Hold my beer" impression of JDK 9 and offers a takeWhile method:

List<String> list = StreamEx.of("one", "two", "stop", "three")
                            .takeWhile(s -> !"stop".equals(s))
                            .toList();
// Output: ["one", "two"]

StreamEx pulls the rabbit out of the hat with a stellar performance using MethodHandle.invokeExact.

Protonpack for the win

Protonpack library also has some tricks up its sleeve with a takeWhile function for Java 8:

List<String> result = StreamUtils.takeWhile(Stream.of("1", "2", "stop", "3"), s -> !s.equals("stop"))
                                 .collect(Collectors.toList());
// Output: ["1", "2"]

Both libraries would make any Java 8 setup punch above its weight by providing ready to use takeWhile creations.

IntStream and predicate combo

The IntStream iterate method, sprinkled with some peek and allMatch, can also wrangle streams:

int[] nums = IntStream.iterate(1, i -> i + 1)
                      .peek(i -> { if (i > 5) throw new BreakerException(); }) // angry bird moment
                      .limit(10)
                      .toArray();