Adding average rate logic in query service

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

@@ -1,17 +1,26 @@
 package org.hypertrace.core.query.service;
 
+import com.google.common.base.Supplier;
+import com.google.common.base.Suppliers;
+import java.time.Duration;
+import java.time.temporal.ChronoUnit;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.TimeUnit;
 import org.hypertrace.core.query.service.api.ColumnIdentifier;
 import org.hypertrace.core.query.service.api.ColumnMetadata;
 import org.hypertrace.core.query.service.api.Expression;
 import org.hypertrace.core.query.service.api.Expression.ValueCase;
 import org.hypertrace.core.query.service.api.Filter;
 import org.hypertrace.core.query.service.api.Function;
+import org.hypertrace.core.query.service.api.Operator;
 import org.hypertrace.core.query.service.api.OrderByExpression;
 import org.hypertrace.core.query.service.api.QueryRequest;
 import org.hypertrace.core.query.service.api.ResultSetMetadata;
@@ -27,6 +36,10 @@ public class ExecutionContext {
   private Set<String> referencedColumns;
   private final LinkedHashSet<String> selectedColumns;
   private ResultSetMetadata resultSetMetadata;
+  // default value null
+  private String timeFilterColumn = null;
+  private Supplier<Optional<Duration>> timeRangeDurationSupplier;
+
   // Contains all selections to be made in the DB: selections on group by, single columns and
   // aggregations in that order.
   // There should be a one-to-one mapping between this and the columnMetadataSet in
@@ -35,14 +48,41 @@ public class ExecutionContext {
   // while the selectedColumns
   // is a set of column names.
   private final LinkedHashSet<Expression> allSelections;
+  private final Optional<Duration> timeSeriesPeriod;
+  private final Filter queryRequestFilter;
 
   public ExecutionContext(String tenantId, QueryRequest request) {
     this.tenantId = tenantId;
     this.selectedColumns = new LinkedHashSet<>();
     this.allSelections = new LinkedHashSet<>();
+    this.timeSeriesPeriod = calculateTimeSeriesPeriod(request);
+    this.queryRequestFilter = request.getFilter();
+    timeRangeDurationSupplier =
+        Suppliers.memoize(
+            () -> findTimeRangeDuration(this.queryRequestFilter, this.timeFilterColumn));
     analyze(request);
   }
 
+  private Optional<Duration> calculateTimeSeriesPeriod(QueryRequest request) {
+    if (request.getGroupByCount() > 0) {
+      for (Expression expression : request.getGroupByList()) {
+        if (expression.getValueCase() == ValueCase.FUNCTION
+            && expression.getFunction().getFunctionName().equals("dateTimeConvert")) {
+          String timeSeriesPeriod =
+              expression
+                  .getFunction()
+                  .getArgumentsList()
+                  .get(3)
+                  .getLiteral()
+                  .getValue()
+                  .getString();
+          return Optional.of(parseDuration(timeSeriesPeriod));
+        }
+      }
+    }
+    return Optional.empty();
+  }
+
   private void analyze(QueryRequest request) {
     List<String> filterColumns = new ArrayList<>();
     LinkedList<Filter> filterQueue = new LinkedList<>();
@@ -154,6 +194,58 @@ private void extractColumns(List<String> columns, Expression expression) {
     }
   }
 
+  private Duration parseDuration(String timeSeriesPeriod) {
+    String[] splitPeriodString = timeSeriesPeriod.split(":");
+    long amount = Long.parseLong(splitPeriodString[0]);
+    ChronoUnit unit = TimeUnit.valueOf(splitPeriodString[1]).toChronoUnit();
+    return Duration.of(amount, unit);
+  }
+
+  private Optional<Duration> findTimeRangeDuration(Filter filter, String timeFilterColumn) {
+
+    // time filter will always be present with AND operator
+    if (filter.getOperator() != Operator.AND) {
+      return Optional.empty();
+    }
+
+    Optional<Long> timeRangeStart =
+        filter.getChildFilterList().stream()
+            .filter(
+                childFilter ->
+                    this.isMatchingFilter(
+                        childFilter, timeFilterColumn, List.of(Operator.GE, Operator.GT)))
+            .map(matchingFilter -> matchingFilter.getRhs().getLiteral().getValue().getLong())
+            .findFirst();
+
+    Optional<Long> timeRangeEnd =
+        filter.getChildFilterList().stream()
+            .filter(
+                childFilter ->
+                    this.isMatchingFilter(
+                        childFilter, timeFilterColumn, List.of(Operator.LT, Operator.LE)))
+            .map(matchingFilter -> matchingFilter.getRhs().getLiteral().getValue().getLong())
+            .findFirst();
+
+    if (timeRangeStart.isPresent() && timeRangeEnd.isPresent()) {
+      return Optional.of(Duration.ofMillis(timeRangeEnd.get() - timeRangeStart.get()));
+    }
+
+    return filter.getChildFilterList().stream()
+        .map(childFilter -> this.findTimeRangeDuration(childFilter, timeFilterColumn))
+        .flatMap(Optional::stream)
+        .findFirst();
+  }
+
+  private boolean isMatchingFilter(Filter filter, String column, Collection<Operator> operators) {
+    return column.equals(filter.getLhs().getColumnIdentifier().getColumnName())
+        && (operators.stream()
+            .anyMatch(operator -> Objects.equals(operator, filter.getOperator())));
+  }
+
+  public void setTimeFilterColumn(String timeFilterColumn) {
+    this.timeFilterColumn = timeFilterColumn;
+  }
+
   public String getTenantId() {
     return this.tenantId;
   }
@@ -173,4 +265,12 @@ public LinkedHashSet<String> getSelectedColumns() {
   public LinkedHashSet<Expression> getAllSelections() {
     return this.allSelections;
   }
+
+  public Optional<Duration> getTimeSeriesPeriod() {
+    return this.timeSeriesPeriod;
+  }
+
+  public Optional<Duration> getTimeRangeDuration() {
+    return timeRangeDurationSupplier.get();
+  }
 }

query-service-impl/src/main/java/org/hypertrace/core/query/service/QueryFunctionConstants.java

@@ -10,6 +10,7 @@ public interface QueryFunctionConstants {
   String QUERY_FUNCTION_MIN = "MIN";
   String QUERY_FUNCTION_MAX = "MAX";
   String QUERY_FUNCTION_COUNT = "COUNT";
+  String QUERY_FUNCTION_AVG_RATE = "AVG_RATE";
   String QUERY_FUNCTION_PERCENTILE = "PERCENTILE";
   String QUERY_FUNCTION_DISTINCTCOUNT = "DISTINCTCOUNT";
   String QUERY_FUNCTION_CONCAT = "CONCAT";

query-service-impl/src/main/java/org/hypertrace/core/query/service/QueryServiceImpl.java

@@ -66,7 +66,11 @@ private Observable<ResultSetChunk> executeTransformedRequest(
                 Status.FAILED_PRECONDITION
                     .withDescription("No handler available matching request")
                     .asException()))
-        .flatMapObservable(handler -> handler.handleRequest(transformedRequest, context))
+        .flatMapObservable(
+            handler -> {
+              handler.getTimeFilterColumn().ifPresent(context::setTimeFilterColumn);
+              return handler.handleRequest(transformedRequest, context);
+            })
         .lift(chunkRows(context.getResultSetMetadata()));
   }
 }

query-service-impl/src/main/java/org/hypertrace/core/query/service/RequestHandler.java

@@ -1,6 +1,7 @@
 package org.hypertrace.core.query.service;
 
 import io.reactivex.rxjava3.core.Observable;
+import java.util.Optional;
 import org.hypertrace.core.query.service.api.QueryRequest;
 import org.hypertrace.core.query.service.api.Row;
 
@@ -20,6 +21,8 @@ public interface RequestHandler {
 
   String getName();
 
+  Optional<String> getTimeFilterColumn();
+
   QueryCost canHandle(QueryRequest request, ExecutionContext context);
 
   /** Handle the request and add rows to the collector. */

query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandler.java

@@ -109,6 +109,11 @@ public String getName() {
     return name;
   }
 
+  @Override
+  public Optional<String> getTimeFilterColumn() {
+    return this.startTimeAttributeName;
+  }
+
   private void processConfig(Config config) {
 
     if (!config.hasPath(TENANT_COLUMN_NAME_CONFIG_KEY)) {

query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/QueryRequestToPinotSQLConverter.java

@@ -19,6 +19,7 @@
 import org.hypertrace.core.query.service.api.SortOrder;
 import org.hypertrace.core.query.service.api.Value;
 import org.hypertrace.core.query.service.api.ValueType;
+import org.hypertrace.core.query.service.pinot.Params.Builder;
 import org.hypertrace.core.query.service.pinot.converters.DestinationColumnValueConverter;
 import org.hypertrace.core.query.service.pinot.converters.PinotFunctionConverter;
 import org.slf4j.Logger;
@@ -63,7 +64,7 @@ Entry<String, Params> toSQL(
     // how it is created.
     for (Expression expr : allSelections) {
       pqlBuilder.append(delim);
-      pqlBuilder.append(convertExpression2String(expr, paramsBuilder));
+      pqlBuilder.append(convertExpression2String(expr, paramsBuilder, executionContext));
       delim = ", ";
     }
 
@@ -75,7 +76,8 @@ Entry<String, Params> toSQL(
 
     if (request.hasFilter()) {
       pqlBuilder.append(" AND ");
-      String filterClause = convertFilter2String(request.getFilter(), paramsBuilder);
+      String filterClause =
+          convertFilter2String(request.getFilter(), paramsBuilder, executionContext);
       pqlBuilder.append(filterClause);
     }
 
@@ -84,7 +86,8 @@ Entry<String, Params> toSQL(
       delim = "";
       for (Expression groupByExpression : request.getGroupByList()) {
         pqlBuilder.append(delim);
-        pqlBuilder.append(convertExpression2String(groupByExpression, paramsBuilder));
+        pqlBuilder.append(
+            convertExpression2String(groupByExpression, paramsBuilder, executionContext));
         delim = ", ";
       }
     }
@@ -93,7 +96,9 @@ Entry<String, Params> toSQL(
       delim = "";
       for (OrderByExpression orderByExpression : request.getOrderByList()) {
         pqlBuilder.append(delim);
-        String orderBy = convertExpression2String(orderByExpression.getExpression(), paramsBuilder);
+        String orderBy =
+            convertExpression2String(
+                orderByExpression.getExpression(), paramsBuilder, executionContext);
         pqlBuilder.append(orderBy);
         if (SortOrder.DESC.equals(orderByExpression.getOrder())) {
           pqlBuilder.append(" desc ");
@@ -119,15 +124,16 @@ Entry<String, Params> toSQL(
     return new SimpleEntry<>(pqlBuilder.toString(), paramsBuilder.build());
   }
 
-  private String convertFilter2String(Filter filter, Params.Builder paramsBuilder) {
+  private String convertFilter2String(
+      Filter filter, Builder paramsBuilder, ExecutionContext executionContext) {
     StringBuilder builder = new StringBuilder();
     String operator = convertOperator2String(filter.getOperator());
     if (filter.getChildFilterCount() > 0) {
       String delim = "";
       builder.append("( ");
       for (Filter childFilter : filter.getChildFilterList()) {
         builder.append(delim);
-        builder.append(convertFilter2String(childFilter, paramsBuilder));
+        builder.append(convertFilter2String(childFilter, paramsBuilder, executionContext));
         builder.append(" ");
         delim = operator + " ";
       }
@@ -140,9 +146,10 @@ private String convertFilter2String(Filter filter, Params.Builder paramsBuilder)
               handleValueConversionForLiteralExpression(filter.getLhs(), filter.getRhs());
           builder.append(operator);
           builder.append("(");
-          builder.append(convertExpression2String(filter.getLhs(), paramsBuilder));
+          builder.append(
+              convertExpression2String(filter.getLhs(), paramsBuilder, executionContext));
           builder.append(",");
-          builder.append(convertExpression2String(rhs, paramsBuilder));
+          builder.append(convertExpression2String(rhs, paramsBuilder, executionContext));
           builder.append(")");
           break;
         case CONTAINS_KEY:
@@ -175,11 +182,12 @@ private String convertFilter2String(Filter filter, Params.Builder paramsBuilder)
           break;
         default:
           rhs = handleValueConversionForLiteralExpression(filter.getLhs(), filter.getRhs());
-          builder.append(convertExpression2String(filter.getLhs(), paramsBuilder));
+          builder.append(
+              convertExpression2String(filter.getLhs(), paramsBuilder, executionContext));
           builder.append(" ");
           builder.append(operator);
           builder.append(" ");
-          builder.append(convertExpression2String(rhs, paramsBuilder));
+          builder.append(convertExpression2String(rhs, paramsBuilder, executionContext));
       }
     }
     return builder.toString();
@@ -251,7 +259,8 @@ private String convertOperator2String(Operator operator) {
     }
   }
 
-  private String convertExpression2String(Expression expression, Params.Builder paramsBuilder) {
+  private String convertExpression2String(
+      Expression expression, Builder paramsBuilder, ExecutionContext executionContext) {
     switch (expression.getValueCase()) {
       case COLUMNIDENTIFIER:
         String logicalColumnName = expression.getColumnIdentifier().getColumnName();
@@ -262,11 +271,13 @@ private String convertExpression2String(Expression expression, Params.Builder pa
         return convertLiteralToString(expression.getLiteral(), paramsBuilder);
       case FUNCTION:
         return this.functionConverter.convert(
+            executionContext,
             expression.getFunction(),
-            argExpression -> convertExpression2String(argExpression, paramsBuilder));
+            argExpression ->
+                convertExpression2String(argExpression, paramsBuilder, executionContext));
       case ORDERBY:
         OrderByExpression orderBy = expression.getOrderBy();
-        return convertExpression2String(orderBy.getExpression(), paramsBuilder);
+        return convertExpression2String(orderBy.getExpression(), paramsBuilder, executionContext);
       case VALUE_NOT_SET:
         break;
     }