optimize metric aggregations order by

gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/QueryRequestContext.java

@@ -24,6 +24,10 @@ public void mapAliasToFunctionExpression(String alias, FunctionExpression functi
     aliasToFunctionExpressionMap.put(alias, functionExpression);
   }
 
+  public boolean containsFunctionExpression(String alias) {
+    return aliasToFunctionExpressionMap.containsKey(alias);
+  }
+
   public FunctionExpression getFunctionExpressionByAlias(String alias) {
     return aliasToFunctionExpressionMap.get(alias);
   }

gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/EntityDataServiceEntityFetcher.java

@@ -169,12 +169,6 @@ public EntityFetcherResponse getEntities(
     return new EntityFetcherResponse(entityBuilders);
   }
 
-  @Override
-  public EntityFetcherResponse getAggregatedMetrics(
-      EntitiesRequestContext requestContext, EntitiesRequest entitiesRequest) {
-    throw new UnsupportedOperationException("Fetching aggregated metrics not supported by EDS");
-  }
-
   @Override
   public EntityFetcherResponse getTimeAggregatedMetrics(
       EntitiesRequestContext requestContext, EntitiesRequest entitiesRequest) {

gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/IEntityFetcher.java

@@ -27,16 +27,6 @@ public interface IEntityFetcher {
   EntityFetcherResponse getEntities(
       EntitiesRequestContext requestContext, EntitiesRequest entitiesRequest);
 
-  /**
-   * Get aggregated metrics
-   *
-   * @param requestContext Additional context for the incoming request
-   * @param entitiesRequest encapsulates the aggregated metrics query (selection, filter, order)
-   * @return Map of the Entity Builders keyed by the EntityId
-   */
-  EntityFetcherResponse getAggregatedMetrics(
-      EntitiesRequestContext requestContext, EntitiesRequest entitiesRequest);
-
   /**
    * Get time series data
    *

gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/QueryServiceEntityFetcher.java

@@ -157,7 +157,10 @@ public EntityFetcherResponse getEntities(
             i++) {
           ColumnMetadata metadata = chunk.getResultSetMetadata().getColumnMetadata(i);
           org.hypertrace.core.query.service.api.Value columnValue = row.getColumn(i);
-          addEntityAttribute(entityBuilder,
+          buildEntity(
+              entityBuilder,
+              requestContext,
+              entitiesRequest,
               metadata,
               columnValue,
               attributeMetadataMap,
@@ -169,97 +172,6 @@ public EntityFetcherResponse getEntities(
     return new EntityFetcherResponse(entityBuilders);
   }
 
-  @Override
-  public EntityFetcherResponse getAggregatedMetrics(
-      EntitiesRequestContext requestContext, EntitiesRequest entitiesRequest) {
-    // Only supported filter is entityIds IN ["id1", "id2", "id3"]
-    Map<String, AttributeMetadata> attributeMetadataMap =
-        attributeMetadataProvider.getAttributesMetadata(
-            requestContext, entitiesRequest.getEntityType());
-    entitiesRequestValidator.validate(entitiesRequest, attributeMetadataMap);
-
-    List<org.hypertrace.gateway.service.v1.common.Expression> aggregates =
-        ExpressionReader.getFunctionExpressions(entitiesRequest.getSelectionList().stream());
-    if (aggregates.isEmpty()) {
-      return new EntityFetcherResponse();
-    }
-
-    List<String> entityIdAttributes =
-        AttributeMetadataUtil.getIdAttributeIds(
-            attributeMetadataProvider, entityIdColumnsConfigs, requestContext, entitiesRequest.getEntityType());
-
-    QueryRequest.Builder builder =
-        constructSelectionQuery(requestContext, entitiesRequest, entityIdAttributes, aggregates);
-    adjustLimitAndOffset(builder, entitiesRequest.getLimit(), entitiesRequest.getOffset());
-
-    QueryRequest request = builder.build();
-
-    if (LOG.isDebugEnabled()) {
-      LOG.debug("Sending Aggregated Metrics Request to Query Service ======== \n {}", request);
-    }
-
-    Iterator<ResultSetChunk> resultSetChunkIterator =
-        queryServiceClient.executeQuery(request, requestContext.getHeaders(),
-            requestTimeout);
-
-    // We want to retain the order as returned from the respective source. Hence using a
-    // LinkedHashMap
-    Map<EntityKey, Builder> entityMap = new LinkedHashMap<>();
-
-    while (resultSetChunkIterator.hasNext()) {
-      ResultSetChunk chunk = resultSetChunkIterator.next();
-      if (LOG.isDebugEnabled()) {
-        LOG.debug("Received chunk: " + chunk.toString());
-      }
-
-      if (chunk.getRowCount() < 1) {
-        break;
-      }
-
-      if (!chunk.hasResultSetMetadata()) {
-        LOG.warn("Chunk doesn't have result metadata so couldn't process the response.");
-        break;
-      }
-
-      for (Row row : chunk.getRowList()) {
-        // Construct the EntityKey from the EntityId attributes columns
-        EntityKey entityKey =
-            EntityKey.of(
-                IntStream.range(0, entityIdAttributes.size())
-                    .mapToObj(value -> row.getColumn(value).getString())
-                    .toArray(String[]::new));
-        Builder entityBuilder = entityMap.computeIfAbsent(entityKey, k -> Entity.newBuilder());
-        entityBuilder.setEntityType(entitiesRequest.getEntityType());
-        entityBuilder.setId(entityKey.toString());
-        // Always include the id in entity since that's needed to make follow up queries in
-        // optimal fashion. If this wasn't really requested by the client, it should be removed
-        // as post processing.
-        for (int i = 0; i < entityIdAttributes.size(); i++) {
-          entityBuilder.putAttribute(
-              entityIdAttributes.get(i),
-              Value.newBuilder()
-                  .setString(entityKey.getAttributes().get(i))
-                  .setValueType(ValueType.STRING)
-                  .build());
-        }
-
-        for (int i = entityIdAttributes.size();
-            i < chunk.getResultSetMetadata().getColumnMetadataCount();
-            i++) {
-          ColumnMetadata metadata = chunk.getResultSetMetadata().getColumnMetadata(i);
-          org.hypertrace.core.query.service.api.Value columnValue = row.getColumn(i);
-          addAggregateMetric(entityBuilder,
-              requestContext,
-              entitiesRequest,
-              metadata,
-              columnValue,
-              attributeMetadataMap);
-        }
-      }
-    }
-    return new EntityFetcherResponse(entityMap);
-  }
-
   private void adjustLimitAndOffset(QueryRequest.Builder builder, int limit, int offset) {
     // If there is more than one groupBy column, we cannot set the same limit that came
     // in the request since that might return less entities than needed when the same
@@ -336,28 +248,51 @@ private QueryRequest.Builder constructSelectionQuery(EntitiesRequestContext requ
     return builder;
   }
 
-  private void addEntityAttribute(Entity.Builder entityBuilder,
+  private void buildEntity(
+      Entity.Builder entityBuilder,
+      QueryRequestContext requestContext,
+      EntitiesRequest entitiesRequest,
       ColumnMetadata metadata,
       org.hypertrace.core.query.service.api.Value columnValue,
       Map<String, AttributeMetadata> attributeMetadataMap,
       boolean isSkipCountColumn) {
 
     // Ignore the count column since we introduced that ourselves into the query
-    if (isSkipCountColumn &&
-        StringUtils.equalsIgnoreCase(COUNT_COLUMN_NAME, metadata.getColumnName())) {
+    if (isSkipCountColumn
+        && StringUtils.equalsIgnoreCase(COUNT_COLUMN_NAME, metadata.getColumnName())) {
       return;
     }
 
+    // aggregate
+    if (requestContext.containsFunctionExpression(metadata.getColumnName())) {
+      addAggregateMetric(
+          entityBuilder,
+          requestContext,
+          entitiesRequest,
+          metadata,
+          columnValue,
+          attributeMetadataMap);
+    } else {
+      // attribute
+      addEntityAttribute(entityBuilder, metadata, columnValue, attributeMetadataMap);
+    }
+  }
+
+  private void addEntityAttribute(
+      Entity.Builder entityBuilder,
+      ColumnMetadata metadata,
+      org.hypertrace.core.query.service.api.Value columnValue,
+      Map<String, AttributeMetadata> attributeMetadataMap) {
+
     String attributeName = metadata.getColumnName();
     entityBuilder.putAttribute(
         attributeName,
         QueryAndGatewayDtoConverter.convertToGatewayValue(
-            attributeName,
-            columnValue,
-            attributeMetadataMap));
+            attributeName, columnValue, attributeMetadataMap));
   }
 
-  private void addAggregateMetric(Entity.Builder entityBuilder,
+  private void addAggregateMetric(
+      Entity.Builder entityBuilder,
       QueryRequestContext requestContext,
       EntitiesRequest entitiesRequest,
       ColumnMetadata metadata,

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,18 +6,18 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
-import java.util.concurrent.Future;
 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;
@@ -41,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;
@@ -145,17 +146,32 @@ 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()
+            .clearTimeAggregation()
             .clearFilter()
             .clearOrderBy()
             .clearLimit()
             .clearOffset()
-            .addAllSelection(
-                executionContext
-                    .getSourceToSelectionExpressionMap()
-                    .getOrDefault(source, executionContext.getEntityIdExpressions()))
+            .addAllSelection(selections)
             .setFilter(dataFetcherNode.getFilter());
 
     if (dataFetcherNode.getLimit() != null) {
@@ -243,6 +259,7 @@ public EntityResponse visit(SelectionNode selectionNode) {
               EntitiesRequest request =
                   EntitiesRequest.newBuilder(executionContext.getEntitiesRequest())
                       .clearSelection()
+                      .clearTimeAggregation()
                       .clearFilter()
                       // TODO: Should we push order by, limit and offet down to the data source?
                       // If we want to push the order by down, we would also have to divide order by into
@@ -273,6 +290,7 @@ public EntityResponse visit(SelectionNode selectionNode) {
                   EntitiesRequest request =
                       EntitiesRequest.newBuilder(executionContext.getEntitiesRequest())
                           .clearSelection()
+                          .clearTimeAggregation()
                           .clearFilter()
                           .clearOrderBy()
                           .clearOffset()
@@ -290,7 +308,7 @@ public EntityResponse visit(SelectionNode selectionNode) {
                           request.getEntityType(),
                           executionContext.getTimestampAttributeId(),
                           executionContext.getRequestHeaders());
-                  return entityFetcher.getAggregatedMetrics(context, request);
+                  return entityFetcher.getEntities(context, request);
                 })
             .collect(Collectors.toList()));
     resultMapList.addAll(
@@ -299,6 +317,7 @@ public EntityResponse visit(SelectionNode selectionNode) {
                 source -> {
                   EntitiesRequest request =
                       EntitiesRequest.newBuilder(executionContext.getEntitiesRequest())
+                          .clearSelection()
                           .clearTimeAggregation()
                           .clearFilter()
                           .clearOrderBy()

gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/EntitiesRequestAndResponseUtils.java

@@ -115,6 +115,13 @@ public static AggregatedMetricValue getAggregatedMetricValue(FunctionType functi
         .build();
   }
 
+  public static AggregatedMetricValue getAggregatedMetricValue(FunctionType functionType, long value) {
+    return AggregatedMetricValue.newBuilder()
+        .setFunction(functionType)
+        .setValue(Value.newBuilder().setLong(value).setValueType(ValueType.LONG))
+        .build();
+  }
+
   public static Expression getLiteralExpression(long value) {
     return Expression.newBuilder()
         .setLiteral(

gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/datafetcher/EntityDataServiceEntityFetcherTests.java

@@ -176,17 +176,6 @@ public void test_getEntities_WithoutPagination() {
         entityDataServiceEntityFetcher.getEntities(entitiesRequestContext, entitiesRequest).size());
   }
 
-  @Test
-  public void test_getAggregatedMetrics() {
-    assertThrows(
-        UnsupportedOperationException.class,
-        () -> {
-          entityDataServiceEntityFetcher.getAggregatedMetrics(
-              new EntitiesRequestContext(TENANT_ID, 0, 1, "API", "API.startTime", Map.of()),
-              EntitiesRequest.newBuilder().build());
-        });
-  }
-
   @Test
   public void test_getTimeAggregatedMetrics() {
     assertThrows(