From 7f11037d55df9deaaf99960baa22cb774511983a Mon Sep 17 00:00:00 2001
From: Ben Yu <yujige@gmail.com>
Date: Thu, 19 Sep 2024 14:54:49 -0700
Subject: [PATCH] BiCollectors.partitioningBy()

---
 .../google/mu/util/stream/BiCollectors.java   | 77 +++++++++++++++++++
 .../google/mu/util/stream/MoreCollectors.java |  6 +-
 .../mu/util/stream/BiCollectorsTest.java      | 18 +++++
 3 files changed, 98 insertions(+), 3 deletions(-)

diff --git a/mug/src/main/java/com/google/mu/util/stream/BiCollectors.java b/mug/src/main/java/com/google/mu/util/stream/BiCollectors.java
index db6e1cf7db..0661e69aa6 100644
--- a/mug/src/main/java/com/google/mu/util/stream/BiCollectors.java
+++ b/mug/src/main/java/com/google/mu/util/stream/BiCollectors.java
@@ -29,6 +29,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
 import java.util.function.BinaryOperator;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -453,6 +454,82 @@ public static <K, V, G> BiCollector<K, V, BiStream<G, V>> groupingBy(
     };
   }
 
+/**
+ * Returns a BiCollector that partitions the incoming pairs into two groups: elements that match
+ * {@code predicate}, and those that don't. Both groups are stored in a BiStream.
+ *
+ * <p>For example:
+ *
+ * <pre>{@code
+ * timeSeries
+ *     .collect(partitioningBy((time, event) -> event.isImportant()))
+ *     .andThen((importantEvents, unimportantEvents) -> ...);
+ * }</pre>
+ *
+ * @since 8.1
+ */
+public static <K, V> BiCollector<K, V, Both<BiStream<K, V>, BiStream<K, V>>> partitioningBy(
+    BiPredicate<? super K, ? super V> predicate) {
+  return partitioningBy(predicate, BiStream::toBiStream);
+}
+
+/**
+ * Returns a BiCollector that partitions the incoming pairs into two groups: elements that match
+ * {@code predicate}, and those that don't, and use {@code downstream} collector to collect the
+ * pairs.
+ *
+ * <p>For example:
+ *
+ * <pre>{@code
+ * timeSeries
+ *     .collect(partitioningBy((time, event) -> event.isImportant(), toSortedImmutableMap()))
+ *     .andThen((importantEvents, unimportantEvents) -> ...);
+ * }</pre>
+ *
+ * @param <K> the input key type
+ * @param <V> the input value type
+ * @param <R> the result type of the downstream collector
+ * @since 8.1
+ */
+public static <K, V, R> BiCollector<K, V, Both<R, R>> partitioningBy(
+    BiPredicate<? super K, ? super V> predicate,
+    BiCollector<? super K, ? super V, ? extends R> downstream) {
+  return partitioningBy(predicate, downstream, downstream);
+}
+
+/**
+ * Returns a BiCollector that partitions the incoming pairs into two groups: elements that match
+ * {@code predicate}, and those that don't, and use {@code ifTrue} and {@code ifFalse} downstream
+ * collectors respectively to collect the pairs.
+ *
+ * <p>For example:
+ *
+ * <pre>{@code
+ * timeSeries
+ *     .collect(
+ *         partitioningBy((time, event) -> event.isImportant(), toImmutableMap(), counting()))
+ *     .andThen((importantEvents, unimportantCount) -> ...);
+ * }</pre>
+ *
+ * @param <K> the input key type
+ * @param <V> the input value type
+ * @param <T> the result type for the pairs that evaluate to true
+ * @param <F> the result type for the pairs that evaluate to false
+ * @since 8.1
+ */
+public static <K, V, T, F> BiCollector<K, V, Both<T, F>> partitioningBy(
+    BiPredicate<? super K, ? super V> predicate,
+    BiCollector<? super K, ? super V, ? extends T> ifTrue,
+    BiCollector<? super K, ? super V, ? extends F> ifFalse) {
+  requireNonNull(predicate);
+  return mapping(
+      AbstractMap.SimpleImmutableEntry<K, V>::new,
+      MoreCollectors.partitioningBy(
+          (Map.Entry<K, V> e) -> predicate.test(e.getKey(), e.getValue()),
+          ifTrue.collectorOf(Map.Entry::getKey, Map.Entry::getValue),
+          ifFalse.collectorOf(Map.Entry::getKey, Map.Entry::getValue)));
+}
+
   /**
    * Returns a {@link BiCollector} that maps the result of {@code upstream} collector using
    * {@code finisher}.
diff --git a/mug/src/main/java/com/google/mu/util/stream/MoreCollectors.java b/mug/src/main/java/com/google/mu/util/stream/MoreCollectors.java
index 989d40c77d..3c77344e07 100644
--- a/mug/src/main/java/com/google/mu/util/stream/MoreCollectors.java
+++ b/mug/src/main/java/com/google/mu/util/stream/MoreCollectors.java
@@ -534,7 +534,7 @@ int arity() {
    * @since 6.0
    */
   public static <E, R> Collector<E, ?, Both<R, R>> partitioningBy(
-      Predicate<? super E> predicate, Collector<E, ?, R> downstream) {
+      Predicate<? super E> predicate, Collector<E, ?, ? extends R> downstream) {
     return partitioningBy(predicate, downstream, downstream);
   }
 
@@ -564,8 +564,8 @@ int arity() {
    */
   public static <E, A1, A2, T, F> Collector<E, ?, Both<T, F>> partitioningBy(
       Predicate<? super E> predicate,
-      Collector<E, A1, T> downstreamIfTrue,
-      Collector<E, A2, F> downstreamIfFalse) {
+      Collector<E, A1, ? extends T> downstreamIfTrue,
+      Collector<E, A2, ? extends F> downstreamIfFalse) {
     requireNonNull(predicate);
     Supplier<A1> factory1 = downstreamIfTrue.supplier();
     Supplier<A2> factory2 = downstreamIfFalse.supplier();
diff --git a/mug/src/test/java/com/google/mu/util/stream/BiCollectorsTest.java b/mug/src/test/java/com/google/mu/util/stream/BiCollectorsTest.java
index 5183d1085d..dc758b99dc 100644
--- a/mug/src/test/java/com/google/mu/util/stream/BiCollectorsTest.java
+++ b/mug/src/test/java/com/google/mu/util/stream/BiCollectorsTest.java
@@ -17,11 +17,13 @@
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth8.assertThat;
 import static com.google.mu.util.stream.BiCollectors.collectingAndThen;
+import static com.google.mu.util.stream.BiCollectors.counting;
 import static com.google.mu.util.stream.BiCollectors.groupingBy;
 import static com.google.mu.util.stream.BiCollectors.maxByKey;
 import static com.google.mu.util.stream.BiCollectors.maxByValue;
 import static com.google.mu.util.stream.BiCollectors.minByKey;
 import static com.google.mu.util.stream.BiCollectors.minByValue;
+import static com.google.mu.util.stream.BiCollectors.partitioningBy;
 import static com.google.mu.util.stream.BiCollectors.toMap;
 import static com.google.mu.util.stream.BiStream.biStream;
 import static com.google.mu.util.stream.BiStreamTest.assertKeyValues;
@@ -306,6 +308,22 @@ public class BiCollectorsTest {
         .inOrder();
   }
 
+  @Test public void testPartitioningBy_sameDownstreamCollector() {
+    String result =
+        BiStream.of(1, "one", 2, "two", 3, "three", 4, "four", 5, "five")
+            .collect(partitioningBy((i, n) -> i % 2 == 1))
+            .andThen((odds, evens) -> "odd:" + odds.toMap() + "; even:" + evens.toMap());
+    assertThat(result).isEqualTo("odd:{1=one, 3=three, 5=five}; even:{2=two, 4=four}");
+  }
+
+  @Test public void testPartitioningBy_differentDownstreamCollectors() {
+    String result =
+        BiStream.of(1, "one", 2, "two", 3, "three", 4, "four", 5, "five")
+            .collect(partitioningBy((i, n) -> i % 2 == 1, toMap(), counting()))
+            .andThen((odds, evens) -> "odd:" + odds + "; count of even:" + evens);
+    assertThat(result).isEqualTo("odd:{1=one, 3=three, 5=five}; count of even:2");
+  }
+
   @Test public void testMapping_downstreamCollector() {
     BiStream<String, Integer> salaries = BiStream.of("Joe", 100, "Tom", 200);
     assertThat(salaries.collect(BiCollectors.mapping(Joiner.on(':')::join, toList())))