Add select key and flatMapToPair methods (#9)

bakdata · Nov 22, 2023 · 7246dd4 · 7246dd4
1 parent 0891165
commit 7246dd4
Show file tree

Hide file tree

Showing 4 changed files with 160 additions and 2 deletions.
diff --git a/src/main/java/com/bakdata/util/seq2/PairSeq.java b/src/main/java/com/bakdata/util/seq2/PairSeq.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c), 2019 bakdata GmbH
+ * Copyright (c), 2023 bakdata GmbH
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
@@ -911,6 +911,16 @@ default <K2, V2> PairSeq<K2, V2> flatMapKeysToPair(
         return seq(this.flatMapKeys(function).keys());
     }
 
+    default <K2, V2> PairSeq<K2, V2> flatMapKeysToOptionalPair(
+            final Function<? super K, ? extends Optional<? extends Tuple2<K2, V2>>> function) {
+        return this.flatMapKeysToPair(k -> function.apply(k).stream());
+    }
+
+    default <K2, V2> PairSeq<K2, V2> flatMapKeysToIterablePair(
+            final Function<? super K, ? extends Iterable<? extends Tuple2<K2, V2>>> function) {
+        return this.flatMapKeysToPair(k -> seq(function.apply(k)));
+    }
+
     default <K2> PairSeq<K2, V> flatMapKeysToOptional(final Function<? super K, Optional<? extends K2>> function) {
         return this.flatMapKeys(k -> function.apply(k).stream());
     }
@@ -973,6 +983,16 @@ default <K2, V2> PairSeq<K2, V2> flatMapToPair(
         return seq(seq2);
     }
 
+    default <K2, V2> PairSeq<K2, V2> flatMapToOptionalPair(
+            final BiFunction<? super K, ? super V, ? extends Optional<? extends Tuple2<K2, V2>>> mapper) {
+        return this.flatMapToPair((k, v) -> mapper.apply(k, v).stream());
+    }
+
+    default <K2, V2> PairSeq<K2, V2> flatMapToIterablePair(
+            final BiFunction<? super K, ? super V, ? extends Iterable<? extends Tuple2<K2, V2>>> mapper) {
+        return this.flatMapToPair((k, v) -> seq(mapper.apply(k, v)));
+    }
+
     default <V2> PairSeq<K, V2> flatMapValues(final Function<? super V, ? extends Stream<? extends V2>> function) {
         return this.flatMapToPair((k, v) -> function.apply(v).map(v2 -> new Tuple2<>(k, v2)));
     }
@@ -982,6 +1002,16 @@ default <K2, V2> PairSeq<K2, V2> flatMapValuesToPair(
         return seq(this.flatMapValues(function).values());
     }
 
+    default <K2, V2> PairSeq<K2, V2> flatMapValuesToOptionalPair(
+            final Function<? super V, ? extends Optional<? extends Tuple2<K2, V2>>> function) {
+        return this.flatMapValuesToPair(v -> function.apply(v).stream());
+    }
+
+    default <K2, V2> PairSeq<K2, V2> flatMapValuesToIterablePair(
+            final Function<? super V, ? extends Iterable<? extends Tuple2<K2, V2>>> function) {
+        return this.flatMapValuesToPair(v -> seq(function.apply(v)));
+    }
+
     default <V2> PairSeq<K, V2> flatMapValuesToOptional(final Function<? super V, Optional<? extends V2>> function) {
         return this.flatMapValues(v -> function.apply(v).stream());
     }
@@ -2298,6 +2328,10 @@ default <U> Seq2<U> scanRight(final U seed,
         return this.toSeq2().scanRight(seed, function);
     }
 
+    default <K2> PairSeq<K2, Tuple2<K, V>> selectKey(final BiFunction<? super K, ? super V, ? extends K2> keyMapper) {
+        return this.mapToPair(keyMapper, Tuple2::new);
+    }
+
     @Override
     default PairSeq<K, V> sequential() {
         return seq(this.toSeq2().sequential());

diff --git a/src/main/java/com/bakdata/util/seq2/Seq2.java b/src/main/java/com/bakdata/util/seq2/Seq2.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c), 2019 bakdata GmbH
+ * Copyright (c), 2023 bakdata GmbH
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
@@ -471,6 +471,22 @@ default <K, V> PairSeq<K, V> flatMapToPair(
         return PairSeq.seq(seq);
     }
 
+    /**
+     * Flat map a {@code Seq2} to a {@code PairSeq}
+     */
+    default <K, V> PairSeq<K, V> flatMapToOptionalPair(
+            final Function<? super T, ? extends Optional<? extends Tuple2<K, V>>> mapper) {
+        return this.flatMapToPair(t -> mapper.apply(t).stream());
+    }
+
+    /**
+     * Flat map a {@code Seq2} to a {@code PairSeq}
+     */
+    default <K, V> PairSeq<K, V> flatMapToIterablePair(
+            final Function<? super T, ? extends Iterable<? extends Tuple2<K, V>>> mapper) {
+        return this.flatMapToPair(t -> seq(mapper.apply(t)));
+    }
+
     @Override
     default void forEach(final Consumer<? super T> action) {
         this.toSeq().forEach(action);
@@ -1044,6 +1060,15 @@ default <U> Seq2<U> scanRight(final U seed, final BiFunction<? super T, ? super
         return seq(this.toSeq().scanRight(seed, function));
     }
 
+    /**
+     * Map a {@code Seq2} to a {@code PairSeq}
+     *
+     * @see Seq#map(Function)
+     */
+    default <K> PairSeq<K, T> selectKey(final Function<? super T, ? extends K> keyMapper) {
+        return this.mapToPair(keyMapper, Function.identity());
+    }
+
     @Override
     default Seq2<T> sequential() {
         return seq(this.toSeq().sequential());

diff --git a/src/test/java/com/bakdata/util/seq2/PairSeqTest.java b/src/test/java/com/bakdata/util/seq2/PairSeqTest.java
@@ -19,6 +19,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Function;
 import java.util.stream.DoubleStream;
 import java.util.stream.IntStream;
 import java.util.stream.LongStream;
@@ -112,6 +113,32 @@ void shouldFlatMapToDouble() {
                 .containsExactlyInAnyOrder(1.0, 2.0);
     }
 
+    @Test
+    void shouldFlatMapToPair() {
+        assertThat((Stream<Tuple2<Integer, String>>) PairSeq.seq(Map.of(1, "a", 2, "b"))
+                .flatMapToPair((k, v) -> PairSeq.seq(Map.of(k, v, k + 1, v + v))))
+                .hasSize(4)
+                .containsExactlyInAnyOrder(new Tuple2<>(1, "a"), new Tuple2<>(2, "aa"),
+                        new Tuple2<>(2, "b"), new Tuple2<>(3, "bb"));
+    }
+
+    @Test
+    void shouldFlatMapToOptionalPair() {
+        assertThat((Stream<Tuple2<Integer, String>>) PairSeq.seq(Map.of(1, "a", 2, "b"))
+                .flatMapToOptionalPair((k, v) -> k == 1 ? Optional.empty() : Optional.of(new Tuple2<>(k + 1, v + v))))
+                .hasSize(1)
+                .containsExactlyInAnyOrder(new Tuple2<>(3, "bb"));
+    }
+
+    @Test
+    void shouldFlatMapToIterablePair() {
+        assertThat((Stream<Tuple2<Integer, String>>) PairSeq.seq(Map.of(1, "a", 2, "b"))
+                .flatMapToIterablePair((k, v) -> List.of(new Tuple2<>(k, v), new Tuple2<>(k + 1, v + v))))
+                .hasSize(4)
+                .containsExactlyInAnyOrder(new Tuple2<>(1, "a"), new Tuple2<>(2, "aa"),
+                        new Tuple2<>(2, "b"), new Tuple2<>(3, "bb"));
+    }
+
     @Test
     void shouldFlatMapValues() {
         assertThat((Stream<Tuple2<Integer, Character>>) PairSeq.seq(Map.of(1, "ab", 2, "bc"))
@@ -130,6 +157,23 @@ void shouldFlatMapKeysToPair() {
                         new Tuple2<>(2, 2), new Tuple2<>(3, 3));
     }
 
+    @Test
+    void shouldFlatMapKeysToOptionalPair() {
+        assertThat((Stream<Tuple2<Integer, Integer>>) PairSeq.seq(Map.of(1, "a", 2, "b"))
+                .flatMapKeysToOptionalPair(k -> k == 1 ? Optional.empty() : Optional.of(new Tuple2<>(k + 1, k))))
+                .hasSize(1)
+                .containsExactlyInAnyOrder(new Tuple2<>(3, 2));
+    }
+
+    @Test
+    void shouldFlatMapKeysToIterablePair() {
+        assertThat((Stream<Tuple2<Integer, Integer>>) PairSeq.seq(Map.of(1, "a", 2, "b"))
+                .flatMapKeysToIterablePair(k -> List.of(new Tuple2<>(k, k), new Tuple2<>(k + 1, k + 1))))
+                .hasSize(4)
+                .containsExactlyInAnyOrder(new Tuple2<>(1, 1), new Tuple2<>(2, 2),
+                        new Tuple2<>(2, 2), new Tuple2<>(3, 3));
+    }
+
     @Test
     void shouldFlatMapValuesToPair() {
         assertThat((Stream<Tuple2<Integer, String>>) PairSeq.seq(Map.of(1, Map.of(3, "a"), 2, Map.of(4, "b", 5, "c")))
@@ -138,6 +182,26 @@ void shouldFlatMapValuesToPair() {
                 .containsExactlyInAnyOrder(new Tuple2<>(3, "a"), new Tuple2<>(4, "b"), new Tuple2<>(5, "c"));
     }
 
+    @Test
+    void shouldFlatMapValuesToOptionalPair() {
+        assertThat((Stream<Tuple2<Integer, String>>) PairSeq.seq(
+                        Map.<Integer, Optional<Tuple2<Integer, String>>>of(1, Optional.of(new Tuple2<>(3, "a")), 2,
+                                Optional.empty()))
+                .flatMapValuesToOptionalPair(Function.identity()))
+                .hasSize(1)
+                .containsExactlyInAnyOrder(new Tuple2<>(3, "a"));
+    }
+
+    @Test
+    void shouldFlatMapValuesToIterablePair() {
+        assertThat((Stream<Tuple2<Integer, String>>) PairSeq.seq(
+                        Map.of(1, List.of(new Tuple2<>(3, "a")), 2, List.of(new Tuple2<>(4, "b"), new Tuple2<>(5,
+                                "c"))))
+                .flatMapValuesToIterablePair(Function.identity()))
+                .hasSize(3)
+                .containsExactlyInAnyOrder(new Tuple2<>(3, "a"), new Tuple2<>(4, "b"), new Tuple2<>(5, "c"));
+    }
+
     @Test
     void shouldFlatMapKeys() {
         assertThat((Stream<Tuple2<Integer, String>>) PairSeq.seq(Map.of(1, "a", 2, "b"))
@@ -394,4 +458,13 @@ void shouldJoinToStringOrEmptyWithPrefixAndSuffix() {
         assertThat(PairSeq.empty().toStringOrEmpty(" ", "^", "$")).isNotPresent();
     }
 
+    @Test
+    void shouldSelectKey() {
+        assertThat((Stream<Tuple2<String, Tuple2<Integer, String>>>) PairSeq.seq(Map.of(1, "a", 2, "b"))
+                .selectKey((k, v) -> k + v))
+                .hasSize(2)
+                .containsExactlyInAnyOrder(new Tuple2<>("1a", new Tuple2<>(1, "a")),
+                        new Tuple2<>("2b", new Tuple2<>(2, "b")));
+    }
+
 }
diff --git a/src/test/java/com/bakdata/util/seq2/Seq2Test.java b/src/test/java/com/bakdata/util/seq2/Seq2Test.java
@@ -18,6 +18,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.stream.Stream;
+import org.jooq.lambda.tuple.Tuple2;
 import org.junit.jupiter.api.Test;
 
 class Seq2Test {
@@ -58,4 +59,29 @@ void shouldJoinToStringWithPrefixAndSuffix() {
         assertThat(Seq2.empty().toStringOrEmpty(" ", "^", "$")).isNotPresent();
     }
 
+    @Test
+    void shouldSelectKey() {
+        assertThat((Stream<Tuple2<Integer, Integer>>) Seq2.seq(List.of(1, 2))
+                .selectKey(i -> i + 1))
+                .hasSize(2)
+                .containsExactlyInAnyOrder(new Tuple2<>(2, 1), new Tuple2<>(3, 2));
+    }
+
+    @Test
+    void shouldFlatMapToOptionalPair() {
+        assertThat((Stream<Tuple2<Integer, Integer>>) Seq2.seq(List.of(1, 2))
+                .flatMapToOptionalPair(i -> i == 1 ? Optional.empty() : Optional.of(new Tuple2<>(i + 1, i))))
+                .hasSize(1)
+                .containsExactlyInAnyOrder(new Tuple2<>(3, 2));
+    }
+
+    @Test
+    void shouldFlatMapToIterablePair() {
+        assertThat((Stream<Tuple2<Integer, Integer>>) Seq2.seq(List.of(1, 2))
+                .flatMapToIterablePair(i -> List.of(new Tuple2<>(i, i), new Tuple2<>(i + 1, i + 1))))
+                .hasSize(4)
+                .containsExactlyInAnyOrder(new Tuple2<>(1, 1), new Tuple2<>(2, 2),
+                        new Tuple2<>(2, 2), new Tuple2<>(3, 3));
+    }
+
 }