com.pivovarit.collectors.ParallelStreamCollector 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.Collection;
import java.util.EnumSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.function.BiConsumer;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collector;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

import static com.pivovarit.collectors.AsyncParallelCollector.requireValidParallelism;
import static com.pivovarit.collectors.BatchingStream.batching;
import static com.pivovarit.collectors.BatchingStream.partitioned;
import static com.pivovarit.collectors.Dispatcher.getDefaultParallelism;
import static com.pivovarit.collectors.Dispatcher.limiting;
import static com.pivovarit.collectors.Dispatcher.of;
import static java.util.Collections.emptySet;
import static java.util.Objects.requireNonNull;
import static java.util.stream.Collectors.collectingAndThen;
import static java.util.stream.Collectors.toList;

/**
 * @author Grzegorz Piwowarek
 */
class ParallelStreamCollector implements Collector>, Stream> {

    private static final EnumSet UNORDERED = EnumSet.of(Characteristics.UNORDERED);

    private final Dispatcher dispatcher;
    private final Function function;
    private final Function>, Stream> processor;
    private final Set characteristics;

    private ParallelStreamCollector(
      Function function,
      Function>, Stream> processor,
      Set characteristics,
      Dispatcher dispatcher) {
        this.processor = processor;
        this.characteristics = characteristics;
        this.dispatcher = dispatcher;
        this.function = function;
    }

    private void startConsuming() {
        if (!dispatcher.isRunning()) {
            dispatcher.start();
        }
    }

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

    @Override
    public BiConsumer>, T> accumulator() {
        return (acc, e) -> {
            startConsuming();
            acc.add(dispatcher.enqueue(() -> function.apply(e)));
        };
    }

    @Override
    public BinaryOperator>> combiner() {
        return (left, right) -> {
            throw new UnsupportedOperationException();
        };
    }

    @Override
    public Function>, Stream> finisher() {
        return acc -> {
            dispatcher.stop();
            return processor.apply(acc);
        };
    }

    @Override
    public Set characteristics() {
        return characteristics;
    }

    static  Collector> streaming(Function mapper, Executor executor) {
        return streaming(mapper, executor, getDefaultParallelism());
    }

    static  Collector> streaming(Function mapper, Executor executor, int parallelism) {
        requireNonNull(executor, "executor can't be null");
        requireNonNull(mapper, "mapper can't be null");
        requireValidParallelism(parallelism);

        return parallelism == 1
          ? BatchingCollectors.syncCollector(mapper)
          : new ParallelStreamCollector<>(mapper, streamInCompletionOrderStrategy(), UNORDERED, limiting(executor, parallelism));
    }

    static  Collector> streamingOrdered(Function mapper, Executor executor) {
        return streamingOrdered(mapper, executor, getDefaultParallelism());
    }

    static  Collector> streamingOrdered(Function mapper, Executor executor, int parallelism) {
        requireNonNull(executor, "executor can't be null");
        requireNonNull(mapper, "mapper can't be null");
        requireValidParallelism(parallelism);

        return parallelism == 1
          ? BatchingCollectors.syncCollector(mapper)
          : new ParallelStreamCollector<>(mapper, streamOrderedStrategy(), emptySet(), limiting(executor, parallelism));
    }

    private static  Function>, Stream> streamInCompletionOrderStrategy() {
        return futures -> StreamSupport.stream(new CompletionOrderSpliterator<>(futures), false);
    }

    private static  Function>, Stream> streamOrderedStrategy() {
        return futures -> futures.stream().map(CompletableFuture::join);
    }

    static final class BatchingCollectors {
        private BatchingCollectors() {
        }

        static  Collector> streaming(Function mapper, Executor executor, int parallelism) {
            requireNonNull(executor, "executor can't be null");
            requireNonNull(mapper, "mapper can't be null");
            requireValidParallelism(parallelism);

            return parallelism == 1
              ? syncCollector(mapper)
              : batched(new ParallelStreamCollector<>(batching(mapper), streamInCompletionOrderStrategy(), UNORDERED, of(executor)), parallelism);
        }

        static  Collector> streamingOrdered(Function mapper, Executor executor, int parallelism) {
            requireNonNull(executor, "executor can't be null");
            requireNonNull(mapper, "mapper can't be null");
            requireValidParallelism(parallelism);

            return parallelism == 1
              ? syncCollector(mapper)
              : batched(new ParallelStreamCollector<>(batching(mapper), streamOrderedStrategy(), emptySet(), of(executor)), parallelism);
        }

        private static  Collector> batched(ParallelStreamCollector, List> downstream, int parallelism) {
            return collectingAndThen(
              toList(),
              list -> partitioned(list, parallelism)
                .collect(collectingAndThen(downstream, s -> s.flatMap(Collection::stream))));
        }

        private static  Collector, Stream> syncCollector(Function mapper) {
            return Collector.of(ArrayList::new, (rs, t) -> rs.add(mapper.apply(t)), (rs, rs2) -> {
                throw new UnsupportedOperationException();
            }, Collection::stream);
        }
    }
}