chore | adding the support pagination in explore (#202)

* chore | adding the support for pagination in explore
hypertrace · Jun 18, 2024 · 99ef811 · 99ef811
1 parent 961268a
commit 99ef811
Show file tree

Hide file tree

Showing 10 changed files with 53 additions and 62 deletions.
diff --git a/...hypertrace/gateway/service/common/converters/EntityServiceAndGatewayServiceConverter.java b/...hypertrace/gateway/service/common/converters/EntityServiceAndGatewayServiceConverter.java
@@ -11,6 +11,7 @@
 import org.hypertrace.entity.query.service.v1.ColumnMetadata;
 import org.hypertrace.entity.query.service.v1.EntityQueryRequest;
 import org.hypertrace.entity.query.service.v1.Expression;
+import org.hypertrace.entity.query.service.v1.Expression.Builder;
 import org.hypertrace.entity.query.service.v1.Filter;
 import org.hypertrace.entity.query.service.v1.Function;
 import org.hypertrace.entity.query.service.v1.LiteralConstant;
@@ -133,6 +134,19 @@ public static Operator convertOperator(
     return Operator.valueOf(operator.name());
   }
 
+  public static OrderByExpression.Builder convertToEntityServiceOrderByExpression(
+      org.hypertrace.gateway.service.v1.common.OrderByExpression orderByExpression) {
+    Builder expressionBuilder = convertToEntityServiceExpression(orderByExpression.getExpression());
+    return OrderByExpression.newBuilder()
+        .setExpression(expressionBuilder)
+        .setOrder(convertToEntityServiceSortOrder(orderByExpression.getOrder()));
+  }
+
+  public static SortOrder convertToEntityServiceSortOrder(
+      org.hypertrace.gateway.service.v1.common.SortOrder sortOrder) {
+    return SortOrder.valueOf(sortOrder.name());
+  }
+
   public static Expression.Builder convertToEntityServiceExpression(
       org.hypertrace.gateway.service.v1.common.Expression expression) {
     Expression.Builder builder = Expression.newBuilder();

diff --git a/...way-service-impl/src/main/java/org/hypertrace/gateway/service/explore/ExploreService.java b/...way-service-impl/src/main/java/org/hypertrace/gateway/service/explore/ExploreService.java
@@ -93,7 +93,6 @@ private void initMetrics() {
   }
 
   public ExploreResponse explore(RequestContext requestContext, ExploreRequest request) {
-
     final Instant start = Instant.now();
     try {
       ExploreRequestContext exploreRequestContext =

diff --git a/...way-service-impl/src/main/java/org/hypertrace/gateway/service/explore/RequestHandler.java b/...way-service-impl/src/main/java/org/hypertrace/gateway/service/explore/RequestHandler.java
@@ -107,7 +107,6 @@ public ExploreResponse.Builder handleRequest(
         buildQueryRequest(requestContext, request, attributeMetadataProvider);
 
     Iterator<ResultSetChunk> resultSetChunkIterator = executeQuery(requestContext, queryRequest);
-
     return handleQueryServiceResponse(
         request, requestContext, resultSetChunkIterator, requestContext, attributeMetadataProvider);
   }
@@ -340,7 +339,11 @@ private void addSortLimitAndOffset(
       ExploreRequestContext requestContext,
       QueryRequest.Builder queryBuilder) {
     if (request.getOrderByCount() > 0) {
-      List<OrderByExpression> orderByExpressions = request.getOrderByList();
+      List<OrderByExpression> orderByExpressions =
+          OrderByUtil.matchOrderByExpressionsAliasToSelectionAlias(
+              request.getOrderByList(),
+              request.getSelectionList(),
+              request.getTimeAggregationList());
       queryBuilder.addAllOrderBy(
           QueryAndGatewayDtoConverter.convertToQueryOrderByExpressions(orderByExpressions));
     }

diff --git a/...ava/org/hypertrace/gateway/service/explore/TimeAggregationsWithGroupByRequestHandler.java b/...ava/org/hypertrace/gateway/service/explore/TimeAggregationsWithGroupByRequestHandler.java
@@ -11,6 +11,7 @@
 import org.hypertrace.gateway.service.common.AttributeMetadataProvider;
 import org.hypertrace.gateway.service.common.util.AttributeMetadataUtil;
 import org.hypertrace.gateway.service.common.util.ExpressionReader;
+import org.hypertrace.gateway.service.common.util.OrderByUtil;
 import org.hypertrace.gateway.service.v1.common.Expression;
 import org.hypertrace.gateway.service.v1.common.Filter;
 import org.hypertrace.gateway.service.v1.common.LiteralConstant;
@@ -111,7 +112,12 @@ private ExploreRequest buildGroupByRequest(ExploreRequest originalRequest) {
         originalRequest.getOrderByList().stream()
             .filter(orderByExpression -> !containsIntervalOrdering(orderByExpression))
             .collect(Collectors.toList());
-    requestBuilder.addAllOrderBy(orderByExpressionList);
+    List<OrderByExpression> orderByExpressions =
+        OrderByUtil.matchOrderByExpressionsAliasToSelectionAlias(
+            orderByExpressionList,
+            originalRequest.getSelectionList(),
+            originalRequest.getTimeAggregationList());
+    requestBuilder.addAllOrderBy(orderByExpressions);
 
     return requestBuilder.build();
   }

diff --git a/...mpl/src/main/java/org/hypertrace/gateway/service/explore/entity/EntityRequestHandler.java b/...mpl/src/main/java/org/hypertrace/gateway/service/explore/entity/EntityRequestHandler.java
@@ -1,18 +1,14 @@
 package org.hypertrace.gateway.service.explore.entity;
 
 import java.util.HashSet;
-import java.util.List;
 import java.util.Optional;
 import java.util.Set;
 import org.hypertrace.gateway.service.common.AttributeMetadataProvider;
 import org.hypertrace.gateway.service.common.datafetcher.QueryServiceEntityFetcher;
-import org.hypertrace.gateway.service.common.util.DataCollectionUtil;
 import org.hypertrace.gateway.service.common.util.QueryServiceClient;
 import org.hypertrace.gateway.service.entity.config.EntityIdColumnsConfig;
 import org.hypertrace.gateway.service.explore.ExploreRequestContext;
 import org.hypertrace.gateway.service.explore.RequestHandler;
-import org.hypertrace.gateway.service.explore.RowComparator;
-import org.hypertrace.gateway.service.v1.common.OrderByExpression;
 import org.hypertrace.gateway.service.v1.explore.EntityOption;
 import org.hypertrace.gateway.service.v1.explore.ExploreRequest;
 import org.hypertrace.gateway.service.v1.explore.ExploreResponse;
@@ -35,8 +31,6 @@
  */
 public class EntityRequestHandler extends RequestHandler {
   private final EntityServiceEntityFetcher entityServiceEntityFetcher;
-  private final AttributeMetadataProvider attributeMetadataProvider;
-  private final QueryServiceEntityFetcher queryServiceEntityFetcher;
 
   public EntityRequestHandler(
       AttributeMetadataProvider attributeMetadataProvider,
@@ -50,9 +44,7 @@ public EntityRequestHandler(
         entityIdColumnsConfig,
         queryServiceEntityFetcher,
         entityServiceEntityFetcher);
-    this.attributeMetadataProvider = attributeMetadataProvider;
     this.entityServiceEntityFetcher = entityServiceEntityFetcher;
-    this.queryServiceEntityFetcher = queryServiceEntityFetcher;
   }
 
   @Override
@@ -67,7 +59,8 @@ public ExploreResponse.Builder handleRequest(
     ExploreResponse.Builder builder = ExploreResponse.newBuilder();
     Set<String> entityIds = new HashSet<>();
     Optional<EntityOption> maybeEntityOption = getEntityOption(exploreRequest);
-    if (requestOnLiveEntities(maybeEntityOption)) {
+    boolean requestOnLiveEntities = requestOnLiveEntities(maybeEntityOption);
+    if (requestOnLiveEntities) {
       entityIds.addAll(getEntityIdsInTimeRangeFromQueryService(requestContext, exploreRequest));
       if (entityIds.isEmpty()) {
         return builder;
@@ -76,16 +69,6 @@ public ExploreResponse.Builder handleRequest(
 
     builder.addAllRow(
         entityServiceEntityFetcher.getResults(requestContext, exploreRequest, entityIds));
-
-    // If there's a Group By in the request, we need to do the sorting and pagination ourselves.
-    if (requestContext.hasGroupBy()) {
-      sortAndPaginatePostProcess(
-          builder,
-          requestContext.getOrderByExpressions(),
-          requestContext.getRowLimitBeforeRest(),
-          requestContext.getOffset());
-    }
-
     if (requestContext.hasGroupBy() && requestContext.getIncludeRestGroup()) {
       getTheRestGroupRequestHandler()
           .getRowsForTheRestGroup(requestContext, exploreRequest, builder);
@@ -94,38 +77,8 @@ public ExploreResponse.Builder handleRequest(
     return builder;
   }
 
-  @Override
-  public void sortAndPaginatePostProcess(
-      ExploreResponse.Builder builder,
-      List<OrderByExpression> orderByExpressions,
-      int limit,
-      int offset) {
-    List<org.hypertrace.gateway.service.v1.common.Row.Builder> rowBuilders =
-        builder.getRowBuilderList();
-
-    List<org.hypertrace.gateway.service.v1.common.Row.Builder> sortedRowBuilders =
-        sortAndPaginateRowBuilders(rowBuilders, orderByExpressions, limit, offset);
-
-    builder.clearRow();
-    sortedRowBuilders.forEach(builder::addRow);
-  }
-
-  protected List<org.hypertrace.gateway.service.v1.common.Row.Builder> sortAndPaginateRowBuilders(
-      List<org.hypertrace.gateway.service.v1.common.Row.Builder> rowBuilders,
-      List<OrderByExpression> orderByExpressions,
-      int limit,
-      int offset) {
-    RowComparator rowComparator = new RowComparator(orderByExpressions);
-
-    return DataCollectionUtil.limitAndSort(
-        rowBuilders.stream(), limit, offset, orderByExpressions.size(), rowComparator);
-  }
-
   private boolean requestOnLiveEntities(Optional<EntityOption> entityOption) {
-    if (entityOption.isEmpty()) {
-      return true;
-    }
-    return !entityOption.get().getIncludeNonLiveEntities();
+    return entityOption.map(option -> !option.getIncludeNonLiveEntities()).orElse(true);
   }
 
   private Optional<EntityOption> getEntityOption(ExploreRequest exploreRequest) {

diff --git a/...c/main/java/org/hypertrace/gateway/service/explore/entity/EntityServiceEntityFetcher.java b/...c/main/java/org/hypertrace/gateway/service/explore/entity/EntityServiceEntityFetcher.java
@@ -25,17 +25,18 @@
 import org.hypertrace.gateway.service.common.util.AttributeMetadataUtil;
 import org.hypertrace.gateway.service.common.util.ExpressionReader;
 import org.hypertrace.gateway.service.common.util.MetricAggregationFunctionUtil;
+import org.hypertrace.gateway.service.common.util.OrderByUtil;
 import org.hypertrace.gateway.service.entity.config.EntityIdColumnsConfig;
 import org.hypertrace.gateway.service.explore.ExploreRequestContext;
 import org.hypertrace.gateway.service.v1.common.FunctionExpression;
+import org.hypertrace.gateway.service.v1.common.OrderByExpression;
 import org.hypertrace.gateway.service.v1.common.Row;
 import org.hypertrace.gateway.service.v1.explore.ExploreRequest;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 public class EntityServiceEntityFetcher {
   private static final Logger LOGGER = LoggerFactory.getLogger(EntityServiceEntityFetcher.class);
-  private static final int DEFAULT_ENTITY_REQUEST_LIMIT = 10_000;
   private final AttributeMetadataProvider attributeMetadataProvider;
   private final EntityIdColumnsConfig entityIdColumnsConfig;
   private final EntityQueryServiceClient entityQueryServiceClient;
@@ -166,13 +167,28 @@ private EntityQueryRequest buildRequest(
         EntityQueryRequest.newBuilder()
             .setEntityType(entityType)
             .setFilter(buildFilter(exploreRequest, entityIdAttributeIds, entityIds));
-
-    addGroupBys(exploreRequest, builder);
     addSelections(requestContext, exploreRequest, builder);
-    builder.setLimit(DEFAULT_ENTITY_REQUEST_LIMIT);
+    addGroupBys(exploreRequest, builder);
+    addSortBy(exploreRequest, builder);
+    builder.setLimit(exploreRequest.getLimit());
+    builder.setOffset(exploreRequest.getOffset());
     return builder.build();
   }
 
+  private void addSortBy(ExploreRequest exploreRequest, EntityQueryRequest.Builder builder) {
+    List<OrderByExpression> orderByExpressions =
+        OrderByUtil.matchOrderByExpressionsAliasToSelectionAlias(
+            exploreRequest.getOrderByList(),
+            exploreRequest.getSelectionList(),
+            exploreRequest.getTimeAggregationList());
+    orderByExpressions.forEach(
+        orderBy ->
+            builder.addOrderBy(
+                EntityServiceAndGatewayServiceConverter.convertToEntityServiceOrderByExpression(
+                        orderBy)
+                    .build()));
+  }
+
   private void addGroupBys(ExploreRequest exploreRequest, EntityQueryRequest.Builder builder) {
     List<org.hypertrace.gateway.service.v1.common.Expression> groupBys =
         ExpressionReader.getAttributeExpressions(exploreRequest.getGroupByList());

diff --git a/...ice-requests-and-responses/explore/aggregations-with-groupby-and-the-rest-deprecated.json b/...ice-requests-and-responses/explore/aggregations-with-groupby-and-the-rest-deprecated.json
@@ -112,7 +112,7 @@
                   }
                 }
               ],
-              "alias": "COUNT#COUNT:Api.Trace|apiTraceId"
+              "alias": "COUNT#results/countTraces:Api.Trace|apiTraceId"
             }
           }
         }

diff --git a/...est/resources/query-service-requests-and-responses/explore/aggregations-with-groupby.json b/...est/resources/query-service-requests-and-responses/explore/aggregations-with-groupby.json
@@ -112,7 +112,7 @@
                   }
                 }
               ],
-              "alias": "COUNT#COUNT:Api.Trace|apiTraceId"
+              "alias": "COUNT#results/countTraces:Api.Trace|apiTraceId"
             }
           }
         }

diff --git a/...ts-and-responses/explore/aggregations-with-multiple-groupbys-and-the-rest-deprecated.json b/...ts-and-responses/explore/aggregations-with-multiple-groupbys-and-the-rest-deprecated.json
@@ -122,7 +122,7 @@
                   }
                 }
               ],
-              "alias": "COUNT#COUNT:Api.Trace|apiTraceId"
+              "alias": "COUNT#results/countTraces:Api.Trace|apiTraceId"
             }
           }
         }

diff --git a/.../src/test/resources/query-service-requests-and-responses/explore/simple-aggregations.json b/.../src/test/resources/query-service-requests-and-responses/explore/simple-aggregations.json
@@ -130,7 +130,7 @@
                   }
                 }
               ],
-              "alias": "AVG#AVG:Api.Trace|duration"
+              "alias": "AVG#results/avgLatency:Api.Trace|duration"
             }
           },
           "order": "DESC"
-Original file line number
+Diff line change
@@ Expand Up / @@ -93,7 +93,6 @@ private void initMetrics() { @@
       }
       public ExploreResponse explore(RequestContext requestContext, ExploreRequest request) {
         final Instant start = Instant.now();
         try {
           ExploreRequestContext exploreRequestContext =
@@ Expand Down @@