Add SpanJoiner based on list of spanId

Author: AnandShivansh

hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/DefaultSpanJoinerBuilder.java

@@ -2,18 +2,15 @@
 
 import static com.google.common.collect.ImmutableList.copyOf;
 import static com.google.common.collect.Iterables.concat;
-import static java.util.Collections.unmodifiableMap;
 import static org.hypertrace.core.graphql.atttributes.scopes.HypertraceCoreAttributeScopeString.SPAN;
 import static org.hypertrace.core.graphql.span.joiner.SpanJoin.SPANS_KEY;
 
 import graphql.schema.DataFetchingFieldSelectionSet;
 import graphql.schema.SelectedField;
 import io.reactivex.rxjava3.core.Observable;
 import io.reactivex.rxjava3.core.Single;
-import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -109,18 +106,20 @@ private <T> Single<Map<T, Collection<Span>>> joinSpans(
 
     private <T> Map<T, Collection<Span>> buildSourceToSpansMap(
         Map<T, Collection<String>> sourceToSpanIdsMap, Map<String, Span> spanIdToSpanMap) {
-      Map<T, Collection<Span>> sourceToSpansMap = new HashMap<>();
-      for (Entry<T, Collection<String>> entry : sourceToSpanIdsMap.entrySet()) {
-        List<Span> spans = new ArrayList<>();
-        for (String spanId : entry.getValue()) {
-          if (spanIdToSpanMap.containsKey(spanId)) {
-            spans.add(spanIdToSpanMap.get(spanId));
-          }
-        }
-        sourceToSpansMap.put(
-            entry.getKey(), spans.stream().distinct().collect(Collectors.toUnmodifiableList()));
-      }
-      return unmodifiableMap(sourceToSpansMap);
+      return sourceToSpanIdsMap.entrySet().stream()
+          .collect(
+              Collectors.toUnmodifiableMap(
+                  Entry::getKey,
+                  entry -> buildSpansCollectionForEachSource(entry.getValue(), spanIdToSpanMap)));
+    }
+
+    private Collection<Span> buildSpansCollectionForEachSource(
+        Collection<String> spanIds, Map<String, Span> spanIdToSpanMap) {
+      return spanIds.stream()
+          .filter(spanIdToSpanMap::containsKey)
+          .distinct()
+          .map(spanIdToSpanMap::get)
+          .collect(Collectors.toUnmodifiableList());
     }
 
     private Map<String, Span> buildSpanIdToSpanMap(SpanResultSet resultSet) {