diff --git a/src/main/java/com/annimon/stream/Stream.java b/src/main/java/com/annimon/stream/Stream.java index 3c4f96d4..b0ed6bd4 100644 --- a/src/main/java/com/annimon/stream/Stream.java +++ b/src/main/java/com/annimon/stream/Stream.java @@ -659,7 +659,59 @@ public int compare(T o1, T o2) { public Stream>> groupBy(final Function classifier) { return Stream.of( collect(Collectors.groupingBy(classifier)) ); } + + /** + * Partitions {@code Stream} into {@code List}s according to the given classifier function. In contrast + * to {@link #groupBy(Function)}, this method assumes that the elements of the stream are sorted. + * Because of this assumption, it does not need to first collect all elements and then partition them. + * Instead, it can emit a {@code List} of elements when it reaches the first element that does not + * belong to the same chunk as the previous elements. + *

+ *

This is an intermediate operation. + * + * @param the type of the keys, which are the result of the classifier function + * @param classifier the classifier function + * @return the new stream + */ + public Stream> chunkBy(final Function classifier) { + return new Stream>(new LsaIterator>() { + private T next; + private boolean peekedNext; + + @Override + public boolean hasNext() { + return peekedNext || iterator.hasNext(); + } + + @Override + public List next() { + K key = classifier.apply(peek()); + List list = new ArrayList(); + do { + list.add(takeNext()); + } + while (iterator.hasNext() && key.equals(classifier.apply(peek()))); + + return list; + } + + private T takeNext() { + T element = peek(); + peekedNext = false; + return element; + } + + private T peek() { + if (!peekedNext) { + next = iterator.next(); + peekedNext = true; + } + return next; + } + }); + } + /** * Perform provided action to each elements. * diff --git a/src/test/java/com/annimon/stream/StreamTest.java b/src/test/java/com/annimon/stream/StreamTest.java index 9d48f884..ddcdc9f3 100644 --- a/src/test/java/com/annimon/stream/StreamTest.java +++ b/src/test/java/com/annimon/stream/StreamTest.java @@ -489,6 +489,20 @@ public void accept(Map.Entry> entry) { assertEquals("[2, 3, 2, 3, 2, 3]", pc2.toString()); } + @Test + public void testChunkBy() { + final PrintConsumer> consumer = new PrintConsumer>(); + + Stream.of(1, 1, 2, 2, 2, 3, 1).chunkBy(new Function() { + @Override + public Integer apply(Integer value) { + return value; + } + }).forEach(consumer); + + assertEquals("[1, 1][2, 2, 2][3][1]", consumer.toString()); + } + @Test public void testPeek() { final PrintConsumer consumer = new PrintConsumer();