com.pivovarit.collectors.BatchingSpliterator Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of parallel-collectors Show documentation
Parallel collection processing with customizable thread pools
The newest version!
package com.pivovarit.collectors;

import java.util.ArrayList;
import java.util.List;
import java.util.Spliterator;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.stream.Stream;

import static java.util.stream.Stream.empty;
import static java.util.stream.Stream.of;
import static java.util.stream.StreamSupport.stream;

/**
 * @author Grzegorz Piwowarek
 */
final class BatchingSpliterator implements Spliterator> {

    private final List source;
    private final int maxChunks;

    private int chunks;
    private int chunkSize;
    private int consumed;

    private BatchingSpliterator(List list, int batches) {
        if (batches < 1) {
            throw new IllegalArgumentException("batches can't be lower than one");
        }
        source = list;
        chunks = batches;
        maxChunks = Math.min(list.size(), batches);
        chunkSize = (int) Math.ceil(((double) source.size()) / batches);
    }

    static  Stream> partitioned(List list, int numberOfParts) {
        int size = list.size();

        if (size <= numberOfParts) {
            return asSingletonListStream(list);
        } else if (size == 0) {
            return empty();
        } else if (numberOfParts == 1) {
            return of(list);
        } else {
            return stream(new BatchingSpliterator<>(list, numberOfParts), false);
        }
    }

    private static  Stream> asSingletonListStream(List list) {
        Stream.Builder> acc = Stream.builder();
        for (T t : list) {
            acc.add(List.of(t));
        }
        return acc.build();
    }

    static  Function, List> batching(Function mapper) {
        return batch -> {
            List list = new ArrayList<>(batch.size());
            for (T t : batch) {
                list.add(mapper.apply(t));
            }
            return list;
        };
    }

    @Override
    public boolean tryAdvance(Consumer> action) {
        if (consumed < source.size() && chunks != 0) {
            List batch = source.subList(consumed, consumed + chunkSize);
            consumed += chunkSize;
            chunkSize = (int) Math.ceil(((double) (source.size() - consumed)) / --chunks);
            action.accept(batch);
            return true;
        } else {
            return false;
        }
    }

    @Override
    public Spliterator> trySplit() {
        return null;
    }

    @Override
    public long estimateSize() {
        return maxChunks;
    }

    @Override
    public int characteristics() {
        return ORDERED | SIZED;
    }
}