optimize metric aggregations order by

Author: skjindal93

gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/ExecutionContext.java

@@ -206,14 +206,18 @@ public void removePendingSelectionSource(String source) {
     pendingSelectionSources.remove(source);
   }
 
-  public void removePendingMetricAggregationSources(String source) {
+  public void removePendingMetricAggregationSource(String source) {
     pendingMetricAggregationSources.remove(source);
   }
 
   public void removePendingSelectionSourceForOrderBy(String source) {
     pendingSelectionSourcesForOrderBy.remove(source);
   }
 
+  public void removePendingMetricAggregationSourceForOrderBy(String source) {
+    pendingMetricAggregationSourcesForOrderBy.remove(source);
+  }
+
   public Map<String, List<Expression>> getSourceToFilterExpressionMap() {
     return sourceToFilterExpressionMap;
   }

gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/ExecutionTreeBuilder.java

@@ -197,7 +197,7 @@ QueryNode buildExecutionTree(ExecutionContext executionContext, QueryNode filter
           new SelectionNode.Builder(rootNode)
               .setAggMetricSelectionSources(metricSourcesForOrderBy)
               .build();
-      metricSourcesForOrderBy.forEach(executionContext::removePendingMetricAggregationSources);
+      metricSourcesForOrderBy.forEach(executionContext::removePendingMetricAggregationSource);
     }
 
     // Try adding SortAndPaginateNode

gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/visitor/ExecutionContextBuilderVisitor.java

@@ -58,6 +58,13 @@ public Void visit(DataFetcherNode dataFetcherNode) {
     executionContext.removePendingSelectionSource(source);
     // TODO: Currently, assumes that the order by attribute is also present in the selection set
     executionContext.removePendingSelectionSourceForOrderBy(source);
+    // TODO: Remove redundant attributes for metric aggregation source for order by
+    // The current metric aggregation source is only QS
+
+    // The order by on metric aggregations will also be added in the selections list of
+    // DataFetcherNode, so that the order by metric aggregations can be fetched before
+    // and only the required data set is paginated
+    executionContext.removePendingMetricAggregationSourceForOrderBy(source);
 
     // set of attributes which were fetched from the source
     Map<String, Set<String>> sourceToSelectionAttributeMap =

gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/visitor/ExecutionVisitor.java

@@ -6,6 +6,7 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.LinkedList;
@@ -16,6 +17,7 @@
 import java.util.concurrent.Executors;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
+import java.util.stream.Stream;
 
 import com.google.common.collect.Sets;
 import org.hypertrace.gateway.service.common.datafetcher.EntityFetcherResponse;
@@ -39,6 +41,7 @@
 import org.hypertrace.gateway.service.v1.common.Filter;
 import org.hypertrace.gateway.service.v1.common.LiteralConstant;
 import org.hypertrace.gateway.service.v1.common.Operator;
+import org.hypertrace.gateway.service.v1.common.OrderByExpression;
 import org.hypertrace.gateway.service.v1.common.Value;
 import org.hypertrace.gateway.service.v1.common.ValueType;
 import org.hypertrace.gateway.service.v1.entity.EntitiesRequest;
@@ -143,6 +146,23 @@ public EntityResponse visit(DataFetcherNode dataFetcherNode) {
             executionContext.getTimestampAttributeId(),
             executionContext.getRequestHeaders());
 
+    // fetching both attribute selections and metric order by selections for
+    // optimized pagination
+    List<Expression> attributeSelections =
+        executionContext
+            .getSourceToSelectionExpressionMap()
+            .getOrDefault(source, executionContext.getEntityIdExpressions());
+    List<Expression> metricOrderBySelections =
+        executionContext
+            .getSourceToMetricOrderByExpressionMap()
+            .getOrDefault(source, Collections.emptyList())
+            .stream()
+            .map(OrderByExpression::getExpression)
+            .collect(Collectors.toList());
+    List<Expression> selections = Stream.of(attributeSelections, metricOrderBySelections)
+        .flatMap(Collection::stream)
+        .collect(Collectors.toList());
+
     EntitiesRequest.Builder requestBuilder =
         EntitiesRequest.newBuilder(entitiesRequest)
             .clearSelection()
@@ -151,10 +171,7 @@ public EntityResponse visit(DataFetcherNode dataFetcherNode) {
             .clearOrderBy()
             .clearLimit()
             .clearOffset()
-            .addAllSelection(
-                executionContext
-                    .getSourceToSelectionExpressionMap()
-                    .getOrDefault(source, executionContext.getEntityIdExpressions()))
+            .addAllSelection(selections)
             .setFilter(dataFetcherNode.getFilter());
 
     if (dataFetcherNode.getLimit() != null) {