Skip to content

Instantly share code, notes, and snippets.

@massahud
Created June 2, 2022 14:28
Show Gist options
  • Save massahud/e899570a2e7089b43e3c6dbf5ad1d909 to your computer and use it in GitHub Desktop.
Save massahud/e899570a2e7089b43e3c6dbf5ad1d909 to your computer and use it in GitHub Desktop.
Chunk spliterator - transforms Stream<X> into Stream<List<X>>
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