chore | Add SpanJoiner based on list of spanId

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

@@ -3,8 +3,12 @@
 import static com.google.common.collect.ImmutableList.copyOf;
 import static com.google.common.collect.Iterables.concat;
 import static org.hypertrace.core.graphql.atttributes.scopes.HypertraceCoreAttributeScopeString.SPAN;
+import static org.hypertrace.core.graphql.span.joiner.SpanJoin.SPANS_KEY;
 import static org.hypertrace.core.graphql.span.joiner.SpanJoin.SPAN_KEY;
 
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.ListMultimap;
+import com.google.common.collect.Multimaps;
 import graphql.schema.DataFetchingFieldSelectionSet;
 import graphql.schema.SelectedField;
 import io.reactivex.rxjava3.core.Observable;
@@ -46,9 +50,9 @@
 public class DefaultSpanJoinerBuilder implements SpanJoinerBuilder {
 
   private static final int ZERO_OFFSET = 0;
-
   private final SpanDao spanDao;
   private final GraphQlSelectionFinder selectionFinder;
+
   private final ResultSetRequestBuilder resultSetRequestBuilder;
   private final FilterRequestBuilder filterRequestBuilder;
 
@@ -70,61 +74,122 @@ public Single<SpanJoiner> build(
       TimeRangeArgument timeRange,
       DataFetchingFieldSelectionSet selectionSet,
       List<String> pathToSpanJoin) {
-    return Single.just(
-        new DefaultSpanJoiner(
-            context, timeRange, this.getSelections(selectionSet, pathToSpanJoin)));
-  }
-
-  private List<SelectedField> getSelections(
-      DataFetchingFieldSelectionSet selectionSet, List<String> pathToSpanJoin) {
-    List<String> fullPath = copyOf(concat(pathToSpanJoin, List.of(SPAN_KEY)));
-    return selectionFinder
-        .findSelections(selectionSet, SelectionQuery.builder().selectionPath(fullPath).build())
-        .collect(Collectors.toUnmodifiableList());
+    return Single.just(new DefaultSpanJoiner(context, timeRange, selectionSet, pathToSpanJoin));
   }
 
   @AllArgsConstructor
   private class DefaultSpanJoiner implements SpanJoiner {
 
     private final GraphQlRequestContext context;
     private final TimeRangeArgument timeRange;
-    private final List<SelectedField> selectedFields;
+    private final DataFetchingFieldSelectionSet selectionSet;
+    private final List<String> pathToJoin;
 
     @Override
-    public <T> Single<Map<T, Span>> joinSpans(
+    public <T> Single<Map<T, Span>> joinSpan(
         Collection<T> joinSources, SpanIdGetter<T> spanIdGetter) {
-      return this.buildSourceToIdMap(joinSources, spanIdGetter).flatMap(this::joinSpans);
+      return this.buildSourceToIdMap(joinSources, spanIdGetter)
+          .flatMap(
+              sourceToSpanIdsMap ->
+                  this.buildSpanRequest(sourceToSpanIdsMap, SPAN_KEY)
+                      .flatMap(spanDao::getSpans)
+                      .map(this::buildSpanIdToSpanMap)
+                      .map(
+                          spanIdToSpanMap ->
+                              this.buildSourceToSpanListMultiMap(
+                                  sourceToSpanIdsMap, spanIdToSpanMap)))
+          .map(Multimaps::asMap)
+          .map(this::reduceMap);
     }
 
-    private <T> Single<Map<T, Span>> joinSpans(Map<T, String> sourceToSpanIdMap) {
-      return this.buildSpanRequest(sourceToSpanIdMap)
-          .flatMap(spanDao::getSpans)
-          .map(this::buildSpanIdToSpanMap)
-          .map(spanIdToSpanMap -> buildSourceToSpanMap(sourceToSpanIdMap, spanIdToSpanMap));
+    @Override
+    public <T> Single<ListMultimap<T, Span>> joinSpans(
+        Collection<T> joinSources, MultipleSpanIdGetter<T> multipleSpanIdGetter) {
+      return this.buildSourceToIdsMap(joinSources, multipleSpanIdGetter)
+          .flatMap(
+              sourceToSpanIdsMap ->
+                  this.buildSpanRequest(sourceToSpanIdsMap, SPANS_KEY)
+                      .flatMap(spanDao::getSpans)
+                      .map(this::buildSpanIdToSpanMap)
+                      .map(
+                          spanIdToSpanMap ->
+                              this.buildSourceToSpanListMultiMap(
+                                  sourceToSpanIdsMap, spanIdToSpanMap)));
     }
 
-    private <T> Map<T, Span> buildSourceToSpanMap(
-        Map<T, String> sourceToSpanIdMap, Map<String, Span> spanIdToSpanMap) {
-      return sourceToSpanIdMap.entrySet().stream()
-          .filter(entry -> spanIdToSpanMap.containsKey(entry.getValue()))
-          .collect(
-              Collectors.toUnmodifiableMap(
-                  Entry::getKey, entry -> spanIdToSpanMap.get(entry.getValue())));
+    private <T> Map<T, Span> reduceMap(Map<T, List<Span>> multiMap) {
+      return multiMap.entrySet().stream()
+          .filter(entry -> !entry.getValue().isEmpty())
+          .collect(Collectors.toUnmodifiableMap(Entry::getKey, entry -> entry.getValue().get(0)));
+    }
+
+    private <T> Single<Map<T, List<String>>> buildSourceToIdMap(
+        Collection<T> joinSources, SpanIdGetter<T> spanIdGetter) {
+      return Observable.fromIterable(joinSources)
+          .flatMapSingle(source -> this.maybeBuildMapEntry(source, spanIdGetter))
+          .collect(Collectors.toMap(Entry::getKey, Entry::getValue));
+    }
+
+    private <T> Single<Map<T, List<String>>> buildSourceToIdsMap(
+        Collection<T> joinSources, MultipleSpanIdGetter<T> multipleSpanIdGetter) {
+      return Observable.fromIterable(joinSources)
+          .flatMapSingle(source -> this.maybeBuildMapEntry(source, multipleSpanIdGetter))
+          .collect(Collectors.toMap(Entry::getKey, Entry::getValue));
+    }
+
+    private <T> Single<Entry<T, List<String>>> maybeBuildMapEntry(
+        T source, SpanIdGetter<T> spanIdGetter) {
+      return spanIdGetter.getSpanId(source).map(List::of).map(ids -> Map.entry(source, ids));
+    }
+
+    private <T> Single<Entry<T, List<String>>> maybeBuildMapEntry(
+        T source, MultipleSpanIdGetter<T> multipleSpanIdGetter) {
+      return multipleSpanIdGetter.getSpanIds(source).map(ids -> Map.entry(source, ids));
+    }
+
+    private <T> ListMultimap<T, Span> buildSourceToSpanListMultiMap(
+        Map<T, List<String>> sourceToSpanIdsMap, Map<String, Span> spanIdToSpanMap) {
+      ListMultimap<T, Span> listMultimap = ArrayListMultimap.create();
+      for (Entry<T, List<String>> entry : sourceToSpanIdsMap.entrySet()) {
+        T source = entry.getKey();
+        for (String spanId : entry.getValue()) {
+          if (spanIdToSpanMap.containsKey(spanId)) {
+            listMultimap.put(source, spanIdToSpanMap.get(spanId));
+          }
+        }
+      }
+      return Multimaps.unmodifiableListMultimap(listMultimap);
+    }
+
+    private List<SelectedField> getSelections(String joinSpanKey) {
+      List<String> fullPath = copyOf(concat(pathToJoin, List.of(joinSpanKey)));
+      return selectionFinder
+          .findSelections(selectionSet, SelectionQuery.builder().selectionPath(fullPath).build())
+          .collect(Collectors.toUnmodifiableList());
     }
 
     private Map<String, Span> buildSpanIdToSpanMap(SpanResultSet resultSet) {
       return resultSet.results().stream()
           .collect(Collectors.toUnmodifiableMap(Identifiable::id, Function.identity()));
     }
 
-    private <T> Single<SpanRequest> buildSpanRequest(Map<T, String> sourceToSpanIdMap) {
-      Collection<String> spanIds = sourceToSpanIdMap.values();
-      return buildSpanIdsFilter(spanIds)
-          .flatMap(filterArguments -> buildSpanRequest(spanIds.size(), filterArguments));
+    private <T> Single<SpanRequest> buildSpanRequest(
+        Map<T, List<String>> sourceToSpanIdsMap, String joinSpanKey) {
+      Collection<String> spanIds =
+          sourceToSpanIdsMap.values().stream()
+              .flatMap(List::stream)
+              .distinct()
+              .collect(Collectors.toUnmodifiableList());
+      List<SelectedField> selectedFields = getSelections(joinSpanKey);
+      return buildSpanIdsFilter(context, spanIds)
+          .flatMap(
+              filterArguments -> buildSpanRequest(spanIds.size(), filterArguments, selectedFields));
     }
 
     private Single<SpanRequest> buildSpanRequest(
-        int size, List<AttributeAssociation<FilterArgument>> filterArguments) {
+        int size,
+        List<AttributeAssociation<FilterArgument>> filterArguments,
+        List<SelectedField> selectedFields) {
       return resultSetRequestBuilder
           .build(
               context,
@@ -140,21 +205,9 @@ private Single<SpanRequest> buildSpanRequest(
     }
 
     private Single<List<AttributeAssociation<FilterArgument>>> buildSpanIdsFilter(
-        Collection<String> spanIds) {
+        GraphQlRequestContext context, Collection<String> spanIds) {
       return filterRequestBuilder.build(context, SPAN, Set.of(new SpanIdFilter(spanIds)));
     }
-
-    private <T> Single<Map<T, String>> buildSourceToIdMap(
-        Collection<T> joinSources, SpanIdGetter<T> spanIdGetter) {
-      return Observable.fromIterable(joinSources)
-          .flatMapSingle(source -> this.maybeBuildMapEntry(source, spanIdGetter))
-          .collect(Collectors.toMap(Entry::getKey, Entry::getValue));
-    }
-
-    private <T> Single<Entry<T, String>> maybeBuildMapEntry(
-        T source, SpanIdGetter<T> spanIdGetter) {
-      return spanIdGetter.getSpanId(source).map(id -> Map.entry(source, id));
-    }
   }
 
   @Value

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

@@ -2,12 +2,18 @@
 
 import graphql.annotations.annotationTypes.GraphQLField;
 import graphql.annotations.annotationTypes.GraphQLName;
+import java.util.List;
 import org.hypertrace.core.graphql.span.schema.Span;
 
 public interface SpanJoin {
   String SPAN_KEY = "span";
+  String SPANS_KEY = "spans";
 
   @GraphQLField
   @GraphQLName(SPAN_KEY)
   Span span();
+
+  @GraphQLField
+  @GraphQLName(SPANS_KEY)
+  List<Span> spans();
 }

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

@@ -1,8 +1,11 @@
 package org.hypertrace.core.graphql.span.joiner;
 
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.ListMultimap;
 import io.reactivex.rxjava3.core.Single;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 import org.hypertrace.core.graphql.span.schema.Span;
 
@@ -12,16 +15,30 @@ public interface SpanJoiner {
   SpanJoiner NO_OP_JOINER =
       new SpanJoiner() {
         @Override
-        public <T> Single<Map<T, Span>> joinSpans(
+        public <T> Single<Map<T, Span>> joinSpan(
             Collection<T> joinSources, SpanIdGetter<T> spanIdGetter) {
           return Single.just(Collections.emptyMap());
         }
+
+        @Override
+        public <T> Single<ListMultimap<T, Span>> joinSpans(
+            Collection<T> joinSources, MultipleSpanIdGetter<T> multipleSpanIdGetter) {
+          return Single.just(ArrayListMultimap.create());
+        }
       };
 
-  <T> Single<Map<T, Span>> joinSpans(Collection<T> joinSources, SpanIdGetter<T> spanIdGetter);
+  <T> Single<Map<T, Span>> joinSpan(Collection<T> joinSources, SpanIdGetter<T> spanIdGetter);
+
+  <T> Single<ListMultimap<T, Span>> joinSpans(
+      Collection<T> joinSources, MultipleSpanIdGetter<T> multipleSpanIdGetter);
 
   @FunctionalInterface
   interface SpanIdGetter<T> {
     Single<String> getSpanId(T source);
   }
+
+  @FunctionalInterface
+  interface MultipleSpanIdGetter<T> {
+    Single<List<String>> getSpanIds(T source);
+  }
 }