Created
June 2, 2022 14:28
-
-
Save massahud/e899570a2e7089b43e3c6dbf5ad1d909 to your computer and use it in GitHub Desktop.
Chunk spliterator - transforms Stream<X> into Stream<List<X>>
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import java.util.ArrayList; | |
import java.util.List; | |
import java.util.Spliterator; | |
import java.util.function.Consumer; | |
import java.util.stream.Stream; | |
import java.util.stream.StreamSupport; | |
| |
public class ChunkSpliterator<T> implements Spliterator<List<T>> | |
{ | |
| |
private final Spliterator<T> wrapped; | |
private final int chunkSize; | |
| |
public ChunkSpliterator(Spliterator<T> spliterator, int chunkSize) | |
{ | |
if (chunkSize < 1) | |
{ | |
throw new IllegalArgumentException("non positive chunkSize"); | |
} | |
this.wrapped = spliterator; | |
this.chunkSize = chunkSize; | |
} | |
| |
public static <X> Stream<List<X>> chunk(Stream<X> stream, int chunkSize) | |
{ | |
return StreamSupport.stream( | |
new ChunkSpliterator<>(stream.spliterator(), chunkSize), | |
stream.isParallel() | |
); | |
} | |
| |
@Override | |
public boolean tryAdvance(Consumer<? super List<T>> action) | |
{ | |
List<T> result = new ArrayList<>(chunkSize); | |
while (result.size() < chunkSize) | |
{ | |
if (!wrapped.tryAdvance(result::add)) | |
{ | |
break; | |
} | |
} | |
if (result.isEmpty()) | |
{ | |
return false; | |
} | |
action.accept(result); | |
return true; | |
} | |
| |
@Override | |
public Spliterator<List<T>> trySplit() | |
{ | |
Spliterator<T> split = wrapped.trySplit(); | |
if (split == null) | |
{ | |
return null; | |
} | |
return new ChunkSpliterator<>(split, chunkSize); | |
} | |
| |
@Override | |
public long estimateSize() | |
{ | |
long wrappedSize = wrapped.estimateSize(); | |
if (wrappedSize == Long.MAX_VALUE) | |
{ | |
return Long.MAX_VALUE; | |
} | |
return wrappedSize / chunkSize; | |
} | |
| |
@Override | |
public int characteristics() | |
{ | |
return (wrapped.characteristics() & | |
(Spliterator.ORDERED | |
| Spliterator.DISTINCT | |
| Spliterator.SIZED | |
| Spliterator.IMMUTABLE | |
| Spliterator.CONCURRENT | |
) | |
) | Spliterator.NONNULL; | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment