Add com.yahoo.stream.CustomCollectors.

2 files changed, 148 insertions, 0 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/stream/CustomCollectors.java b/vespajlib/src/main/java/com/yahoo/stream/CustomCollectors.java
new file mode 100644
index 00000000000..8198c519dfd
--- /dev/null
+++ b/vespajlib/src/main/java/com/yahoo/stream/CustomCollectors.java
@@ -0,0 +1,85 @@
+// Copyright 2017 Yahoo Inc. 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.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}.
+ *
+ * <p>For example, to get a collector that accumulates elements into an insertion
+ * order map:
+ * <pre>{@code
+ *
+ *     Map<String, Person> idToPerson =
+ *         persons.stream().collect(toLinkedMap(Person::id, Functions.identity());
+ * }</pre>
+ *
+ * @author gjoranv
+ */
+public class CustomCollectors {
+
+    private CustomCollectors() { }
+
+    /**
+     * Returns a {@code Collector} that accumulates elements into a {@code Map}
+     * with predictable iteration order. 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 <T> Type of the input elements.
+     * @param <K> Output type of the key mapping function.
+     * @param <U> Output type of the value mapping function.
+     * @return A collector which collects elements into a map with predictable iteration order.
+     * @throws DuplicateKeyException If two elements map to the same key.
+     */
+    public static <T, K, U>
+    Collector<T, ?, Map<K,U>> toLinkedMap(Function<? super T, ? extends K> keyMapper,
+                                          Function<? super T, ? extends U> 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 <T> Type of the input elements.
+     * @param <K> Output type of the key mapping function.
+     * @param <U> Output type of the value mapping function.
+     * @param <M> 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 <T, K, U, M extends Map<K,U>>
+    Collector<T, ?, M> toCustomMap(Function<? super T, ? extends K> keyMapper,
+                                   Function<? super T, ? extends U> valueMapper,
+                                   Supplier<M> mapSupplier) {
+        return Collectors.toMap(keyMapper, valueMapper, throwingMerger(), mapSupplier);
+    }
+
+
+    private static <T> BinaryOperator<T> throwingMerger() {
+        return (u,v) -> { throw new DuplicateKeyException(u); };
+    }
+
+    public static class DuplicateKeyException extends RuntimeException {
+        private static final long serialVersionUID = 1L;
+
+        DuplicateKeyException(Object key) {
+            super(String.format("Duplicate keys: %s", key));
+        }
+    }
+}
diff --git a/vespajlib/src/test/java/com/yahoo/stream/CustomCollectorsTest.java b/vespajlib/src/test/java/com/yahoo/stream/CustomCollectorsTest.java
new file mode 100644
index 00000000000..8cf225c1e34
--- /dev/null
+++ b/vespajlib/src/test/java/com/yahoo/stream/CustomCollectorsTest.java
@@ -0,0 +1,63 @@
+// Copyright 2017 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.stream;
+
+import com.google.common.collect.Lists;
+import com.yahoo.stream.CustomCollectors.DuplicateKeyException;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
+
+import static com.yahoo.stream.CustomCollectors.toCustomMap;
+import static com.yahoo.stream.CustomCollectors.toLinkedMap;
+import static java.util.function.Function.identity;
+import static org.junit.Assert.assertEquals;
+
+/**
+ * @author gjoranv
+ */
+public class CustomCollectorsTest {
+
+    @Rule
+    public ExpectedException thrown= ExpectedException.none();
+
+    @Test
+    public void linked_map_collector_returns_map_with_insertion_order() {
+        List<String> stringList = numberList();
+        Map<String, String> orderedMap = stringList.stream().collect(toLinkedMap(identity(), identity()));
+        int i = 0;
+        for (String val : orderedMap.keySet()) {
+            assertEquals(stringList.get(i), val);
+            i++;
+        }
+    }
+
+    @Test
+    public void custom_map_collector_returns_map_from_given_supplier() {
+        List<String> stringList = numberList();
+        Map<String, String> customMap = stringList.stream().collect(toCustomMap(identity(), identity(), CustomHashMap::new));
+
+        assertEquals(CustomHashMap.class, customMap.getClass());
+    }
+
+    @Test
+    public void custom_map_collector_throws_exception_upon_duplicate_keys() {
+        List<String> duplicates = Lists.newArrayList("same", "same");
+
+        thrown.expect(DuplicateKeyException.class);
+        duplicates.stream().collect(toCustomMap(Function.identity(), Function.identity(), HashMap::new));
+    }
+
+    private static List<String> numberList() {
+        return Lists.newArrayList("zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine", "ten");
+    }
+
+    private static class CustomHashMap<K,V> extends HashMap<K,V> {
+        private static final long serialVersionUID = 1L;  // To avoid compiler warning
+    }
+
+}