org.vertexium.util.StreamUtils Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of vertexium-core Show documentation
There is a newer version: 4.10.0
package org.vertexium.util;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import org.vertexium.Element;
import org.vertexium.VertexiumException;
import org.vertexium.query.Query;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.function.Supplier;
import java.util.stream.Collector;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

public class StreamUtils {
    private StreamUtils() {

    }

    /**
     * Create a {@link java.util.stream.Stream} containing the results of executing the queries, in order. The results
     * are not loaded into memory first.
     */
    public static Stream stream(Query... queries) {
        return Arrays.stream(queries)
            .map(query -> StreamSupport.stream(query.elements().spliterator(), false))
            .reduce(Stream::concat)
            .orElseGet(Stream::empty);
    }

    /**
     * Create a {@link java.util.stream.Stream} over the elements of the iterables, in order.  A list of iterators
     * is first created from the iterables, and passed to {@link #stream(Iterator[])}. The iterable elements are not
     * loaded into memory first.
     */
    @SafeVarargs
    @SuppressWarnings("unchecked")
    public static  Stream stream(Iterable... iterables) {
        List> iterators = Arrays.stream(iterables)
            .map(Iterable::iterator)
            .collect(Collectors.toList());

        return stream(iterators.toArray(new Iterator[iterables.length]));
    }

    /**
     * Create a {@link java.util.stream.Stream} over the elements of the iterators, in order.  The iterator elements
     * are not loaded into memory first.
     */
    @SafeVarargs
    public static  Stream stream(Iterator... iterators) {
        return withCloseHandler(
            Arrays.stream(iterators)
                .map(StreamUtils::streamForIterator)
                .reduce(Stream::concat)
                .orElseGet(Stream::empty),
            iterators
        );
    }

    @SafeVarargs
    private static  Stream withCloseHandler(Stream stream, Iterator... iterators) {
        return stream.onClose(() -> {
            for (Iterator iterator : iterators) {
                if (iterator instanceof AutoCloseable) {
                    try {
                        ((AutoCloseable) iterator).close();
                    } catch (Exception ex) {
                        throw new VertexiumException(
                            String.format("exception occurred when closing %s", iterator.getClass().getName()),
                            ex
                        );
                    }
                }
            }
        });
    }

    private static  Stream streamForIterator(Iterator iterator) {
        return StreamSupport.stream(Spliterators.spliteratorUnknownSize(iterator, 0), false);
    }

    public static  Collector, ImmutableList> toImmutableList() {
        return Collector.of(
            ImmutableList.Builder::new,
            ImmutableList.Builder::add,
            (l, r) -> l.addAll(r.build()),
            ImmutableList.Builder::build
        );
    }

    public static  Collector, ImmutableSet> toImmutableSet() {
        return Collector.of(
            ImmutableSet.Builder::new,
            ImmutableSet.Builder::add,
            (l, r) -> l.addAll(r.build()),
            ImmutableSet.Builder::build,
            Collector.Characteristics.UNORDERED
        );
    }

    public static  Collector, ImmutableMap> toImmutableMap(
        Function keyMapper,
        Function valueMapper
    ) {
        return Collector.of(
            ImmutableMap.Builder::new,
            (r, t) -> r.put(keyMapper.apply(t), valueMapper.apply(t)),
            (l, r) -> l.putAll(r.build()),
            ImmutableMap.Builder::build,
            Collector.Characteristics.UNORDERED
        );
    }

    public static  Collector, LinkedHashSet> toLinkedHashSet() {
        return Collector.of(
            LinkedHashSet::new,
            HashSet::add,
            (a1, a2) -> {
                LinkedHashSet results = new LinkedHashSet();
                results.addAll(a1);
                results.addAll(a2);
                return results;
            },
            ts -> ts
        );
    }

    public static  TReturn ifEmpty(
        Stream stream,
        Supplier trueFunc,
        Function, TReturn> falseFunc
    ) {
        Spliterator split = stream.spliterator();
        AtomicReference firstItem = new AtomicReference<>();
        if (split.tryAdvance(firstItem::set)) {
            Stream newStream = Stream.concat(Stream.of(firstItem.get()), StreamSupport.stream(split, stream.isParallel()));
            return falseFunc.apply(newStream);
        } else {
            return trueFunc.get();
        }
    }

    public static  Predicate distinctBy(Function fn) {
        Set