org.jobrunr.utils.streams.BatchCollector Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of jobrunr Show documentation

An easy way to perform background processing on the JVM. Backed by persistent storage. Open and free for commercial use.

There is a newer version: 7.2.3

Show newest version

package org.jobrunr.utils.streams;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Set;
import java.util.function.BiConsumer;
import java.util.function.BinaryOperator;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collector;

import static java.util.Objects.requireNonNull;

/**
 * Collects elements in the stream and calls the supplied batch processor
 * after the configured batch size is reached.
 * 
 * In case of a parallel stream, the batch processor may be called with
 * elements less than the batch size.
 * 
 * The elements are not kept in memory, and the final result will be an
 * empty list.
 *
 * @param  Type of the elements being collected
 */
class BatchCollector implements Collector, List> {

    private final int batchSize;
    private final Consumer> batchProcessor;

    /**
     * Constructs the batch collector
     *
     * @param batchSize      the batch size after which the batchProcessor should be called
     * @param batchProcessor the batch processor which accepts batches of records to process
     */
    BatchCollector(int batchSize, Consumer> batchProcessor) {
        this.batchSize = batchSize;
        this.batchProcessor = requireNonNull(batchProcessor);
    }

    @Override
    public Supplier> supplier() {
        return ArrayList::new;
    }

    @Override
    public BiConsumer, T> accumulator() {
        return (ts, t) -> {
            ts.add(t);
            if (ts.size() >= batchSize) {
                batchProcessor.accept(ts);
                ts.clear();
            }
        };
    }

    @Override
    public BinaryOperator> combiner() {
        return (ts, ots) -> {
            // process each parallel list without checking for batch size
            // avoids adding all elements of one to another
            // can be modified if a strict batching mode is required
            batchProcessor.accept(ts);
            batchProcessor.accept(ots);
            return Collections.emptyList();
        };
    }

    @Override
    public Function, List> finisher() {
        return ts -> {
            batchProcessor.accept(ts);
            return Collections.emptyList();
        };
    }

    @Override
    public Set characteristics() {
        return Collections.emptySet();
    }
}