com.pivovarit.collectors.BatchingStream 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

There is a newer version: 3.2.0

package com.pivovarit.collectors;

import java.util.ArrayList;
import java.util.Collections;
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 BatchingStream implements Spliterator> {

    private final List source;
    private final int size;
    private final int maxChunks;

    private int chunks;
    private int chunkSize;
    private int leftElements;
    private int i;

    private BatchingStream(List list, int numberOfParts) {
        source = list;
        size = list.size();
        chunks = numberOfParts;
        maxChunks = numberOfParts;
        chunkSize = (int) Math.ceil(((double) size) / numberOfParts);
        leftElements = size;
    }

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

        if (size == numberOfParts) {
            return list.stream().map(Collections::singletonList);
        } else if (size == 0) {
            return empty();
        } else if (numberOfParts == 1) {
            return of(list);
        } else {
            return stream(new BatchingStream<>(list, numberOfParts), false);
        }
    }

    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 (i < size && chunks != 0) {
            List batch = source.subList(i, i + chunkSize);
            i = i + chunkSize;
            leftElements = leftElements - chunkSize;
            chunkSize = (int) Math.ceil(((double) leftElements) / --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;
    }
}