From 452b94fb023356b7a3392c878a67aed9c6b722cf Mon Sep 17 00:00:00 2001
From: david-leifker <114954101+david-leifker@users.noreply.github.com>
Date: Wed, 17 Jul 2024 17:09:36 -0500
Subject: [PATCH] fix(entityservice): fix merging sideeffects (#10937)

---
 .../metadata/aspect/batch/AspectsBatch.java     | 17 +++++++----------
 1 file changed, 7 insertions(+), 10 deletions(-)

diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/AspectsBatch.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/AspectsBatch.java
index 77820948b00cbc..fc4ac90dfabad8 100644
--- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/AspectsBatch.java
+++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/AspectsBatch.java
@@ -9,6 +9,7 @@
 import com.linkedin.util.Pair;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -198,16 +199,12 @@ default Map<String, Set<String>> getNewUrnAspectsMap(
 
   static <T> Map<String, Map<String, T>> merge(
       @Nonnull Map<String, Map<String, T>> a, @Nonnull Map<String, Map<String, T>> b) {
-    return Stream.concat(a.entrySet().stream(), b.entrySet().stream())
-        .flatMap(
-            entry ->
-                entry.getValue().entrySet().stream()
-                    .map(innerEntry -> Pair.of(entry.getKey(), innerEntry)))
-        .collect(
-            Collectors.groupingBy(
-                Pair::getKey,
-                Collectors.mapping(
-                    Pair::getValue, Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))));
+    Map<String, Map<String, T>> mergedMap = new HashMap<>();
+    for (Map.Entry<String, Map<String, T>> entry :
+        Stream.concat(a.entrySet().stream(), b.entrySet().stream()).collect(Collectors.toList())) {
+      mergedMap.computeIfAbsent(entry.getKey(), k -> new HashMap<>()).putAll(entry.getValue());
+    }
+    return mergedMap;
   }
 
   default String toAbbreviatedString(int maxWidth) {