fix(entityservice): fix merging sideeffects (datahub-project#10937)

ajoymajumdar · Jul 17, 2024 · 452b94f · 452b94f
1 parent 0b64de8
commit 452b94f
Showing 1 changed file with 7 additions and 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
@@ -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) {