com.yahoo.stream.CustomCollectors Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of vespajlib Show documentation
Library for use in Java components of Vespa. Shared code which do not fit anywhere else.
There is a newer version: 8.441.21
// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.stream;

import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Optional;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collector;
import java.util.stream.Collectors;

/**
 * The purpose of this class is to fill gaps in the Java {@link Collectors} api
 * by offering convenient ways to retrieve implementations of {@link Collector}.
 *
 * For example, to get a collector that accumulates elements into a map with
 * predictable iteration order:
 * 
{@code
 *
 *     Map idToPerson =
 *         persons.stream().collect(toLinkedMap(Person::id, Functions.identity());
 * }
 *
 * @author gjoranv
 */
public class CustomCollectors {

    private CustomCollectors() { }

    /**
     * Returns a {@code Collector} that accumulates elements into a {@code Map}
     * that provides insertion order iteration. This a convenience that can be used
     * instead of calling {@link java.util.stream.Collectors#toMap(Function, Function, BinaryOperator, Supplier)}.
     * with a merger that throws upon duplicate keys.
     *
     * @param keyMapper Mapping function to produce keys.
     * @param valueMapper Mapping function to produce values.
     * @param  Type of the input elements.
     * @param  Output type of the key mapping function.
     * @param  Output type of the value mapping function.
     * @return A collector which collects elements into a map with insertion order iteration.
     * @throws DuplicateKeyException If two elements map to the same key.
     */
    public static 
    Collector> toLinkedMap(Function keyMapper,
                                          Function valueMapper) {
        return Collectors.toMap(keyMapper, valueMapper, throwingMerger(), LinkedHashMap::new);
    }

    /**
     * Returns a {@code Collector} that accumulates elements into a {@code Map}
     * created by the given supplier. This a convenience that can be used
     * instead of calling {@link java.util.stream.Collectors#toMap(Function, Function, BinaryOperator, Supplier)}.
     * with a merger that throws upon duplicate keys.
     *
     * @param keyMapper Mapping function to produce keys.
     * @param valueMapper Mapping function to produce values.
     * @param mapSupplier Supplier of a new map.
     * @param  Type of the input elements.
     * @param  Output type of the key mapping function.
     * @param  Output type of the value mapping function.
     * @param  Type of the resulting map.
     * @return A collector which collects elements into a map created by the given supplier.
     * @throws DuplicateKeyException If two elements map to the same key.
     */
    public static >
    Collector toCustomMap(Function keyMapper,
                                   Function valueMapper,
                                   Supplier mapSupplier) {
        return Collectors.toMap(keyMapper, valueMapper, throwingMerger(), mapSupplier);
    }

    /**
     * Returns a {@code Collector} that returns a singleton, or throws an {@code IllegalArgumentException} if there are more than one item.
     *
     * @return A collector returning an optional element
     * @param  Type of the input elements.
     * @throws IllegalArgumentException if there are more than one element
     */
    public static  Collector> singleton() {
        return Collectors.collectingAndThen(
                Collectors.toList(),
                list -> {
                    if (list.size() > 1) throw new IllegalArgumentException("More than one element");
                    return list.stream().findAny();
                }
        );
    }

    private static  BinaryOperator throwingMerger() {
        return (u,v) -> { throw new DuplicateKeyException(u); };
    }

    public static class DuplicateKeyException extends IllegalStateException {
        private static final long serialVersionUID = 1L;

        DuplicateKeyException(Object key) {
            super(String.format("Duplicate keys: %s", key));
        }
    }
}