Adding average rate logic in query service

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

@@ -1,11 +1,15 @@
 package org.hypertrace.core.query.service;
 
+import java.time.Duration;
+import java.time.temporal.ChronoUnit;
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.LinkedList;
 import java.util.List;
+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;
@@ -35,14 +39,62 @@ public class ExecutionContext {
   // while the selectedColumns
   // is a set of column names.
   private final LinkedHashSet<Expression> allSelections;
+  private final Optional<Double> timeSeriesPeriod;
+  private final double timeRangeDuration;
 
   public ExecutionContext(String tenantId, QueryRequest request) {
     this.tenantId = tenantId;
     this.selectedColumns = new LinkedHashSet<>();
     this.allSelections = new LinkedHashSet<>();
+    this.timeSeriesPeriod = setTimeSeriesPeriod(request);
+    this.timeRangeDuration = setTimeRangeDuration(request);
     analyze(request);
   }
 
+  private Optional<Double> setTimeSeriesPeriod(QueryRequest request) {
+    Optional<Double> period = Optional.empty();
+    if (request.getGroupByCount() > 0) {
+      for (Expression expression : request.getGroupByList()) {
+        if (expression.getValueCase() == ValueCase.FUNCTION
+            && expression.getFunction().getFunctionName().equals("dateTimeConvert")) {
+          String periodInIso =
+              expression
+                  .getFunction()
+                  .getArgumentsList()
+                  .get(3)
+                  .getLiteral()
+                  .getValue()
+                  .getString();
+          period = Optional.of((double) isoDurationToSeconds(periodInIso));
+        }
+      }
+    }
+    return period;
+  }
+
+  private double setTimeRangeDuration(QueryRequest request) {
+    Optional<Long> duration = Optional.empty();
+    if (request.getFilter().getChildFilterCount() > 0) {
+      for (Filter filter : request.getFilter().getChildFilterList()) {
+        // will extract columnName as a part of follow up PR
+        String columnName = filter.getLhs().getColumnIdentifier().getColumnName();
+        int colLength = columnName.length();
+        if (colLength >= 9 && columnName.substring(colLength - 9).equals("startTime")) {
+          Long val = filter.getRhs().getLiteral().getValue().getLong();
+          if (duration.isPresent()) {
+            val = val - duration.get();
+          }
+          duration = Optional.of(Math.abs(val));
+        }
+      }
+    }
+
+    if (duration.isPresent()) {
+      return (double) TimeUnit.SECONDS.convert(duration.get(), TimeUnit.MILLISECONDS);
+    }
+    return 0;
+  }
+
   private void analyze(QueryRequest request) {
     List<String> filterColumns = new ArrayList<>();
     LinkedList<Filter> filterQueue = new LinkedList<>();
@@ -154,6 +206,11 @@ private void extractColumns(List<String> columns, Expression expression) {
     }
   }
 
+  private static long isoDurationToSeconds(String duration) {
+    Duration d = java.time.Duration.parse(duration);
+    return d.get(ChronoUnit.SECONDS);
+  }
+
   public String getTenantId() {
     return this.tenantId;
   }
@@ -173,4 +230,12 @@ public LinkedHashSet<String> getSelectedColumns() {
   public LinkedHashSet<Expression> getAllSelections() {
     return this.allSelections;
   }
+
+  public Optional<Double> getTimeSeriesPeriod() {
+    return this.timeSeriesPeriod;
+  }
+
+  public double getTimeRangeDuration() {
+    return this.timeRangeDuration;
+  }
 }

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/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;
     }

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

@@ -1,13 +1,16 @@
 package org.hypertrace.core.query.service.pinot.converters;
 
+import static org.hypertrace.core.query.service.QueryFunctionConstants.QUERY_FUNCTION_AVG_RATE;
 import static org.hypertrace.core.query.service.QueryFunctionConstants.QUERY_FUNCTION_CONCAT;
 import static org.hypertrace.core.query.service.QueryFunctionConstants.QUERY_FUNCTION_COUNT;
 import static org.hypertrace.core.query.service.QueryFunctionConstants.QUERY_FUNCTION_PERCENTILE;
 
 import java.util.Optional;
 import java.util.stream.Collectors;
+import org.hypertrace.core.query.service.ExecutionContext;
 import org.hypertrace.core.query.service.api.Expression;
 import org.hypertrace.core.query.service.api.Function;
+import org.hypertrace.core.query.service.api.Function.Builder;
 import org.hypertrace.core.query.service.api.LiteralConstant;
 import org.hypertrace.core.query.service.api.Value;
 import org.hypertrace.core.query.service.api.ValueType;
@@ -16,6 +19,9 @@ public class PinotFunctionConverter {
   /**
    * Computing PERCENTILE in Pinot is resource intensive. T-Digest calculation is much faster and
    * reasonably accurate, hence use that as the default.
+   *
+   * <p>AVG_RATE not supported directly in Pinot. So AVG_RATE is computed by summing over all values
+   * and then dividing by a constant.
    */
   private static final String DEFAULT_PERCENTILE_AGGREGATION_FUNCTION = "PERCENTILETDIGEST";
 
@@ -34,24 +40,73 @@ public PinotFunctionConverter() {
   }
 
   public String convert(
-      Function function, java.util.function.Function<Expression, String> argumentConverter) {
+      ExecutionContext executionContext,
+      Function function,
+      java.util.function.Function<Expression, String> argumentConverter) {
     switch (function.getFunctionName().toUpperCase()) {
       case QUERY_FUNCTION_COUNT:
         return this.convertCount();
       case QUERY_FUNCTION_PERCENTILE:
         return this.functionToString(this.toPinotPercentile(function), argumentConverter);
       case QUERY_FUNCTION_CONCAT:
         return this.functionToString(this.toPinotConcat(function), argumentConverter);
+      case QUERY_FUNCTION_AVG_RATE:
+        return this.functionToStringForAvgRate(function, argumentConverter, executionContext);
       default:
         // TODO remove once pinot-specific logic removed from gateway - this normalization reverts
         // that logic
         if (this.isHardcodedPercentile(function)) {
-          return this.convert(this.normalizeHardcodedPercentile(function), argumentConverter);
+          return this.convert(
+              executionContext, this.normalizeHardcodedPercentile(function), argumentConverter);
         }
         return this.functionToString(function, argumentConverter);
     }
   }
 
+  private String functionToStringForAvgRate(
+      Function function,
+      java.util.function.Function<Expression, String> argumentConverter,
+      ExecutionContext executionContext) {
+    function = updateFunctionForAvgRate(function, executionContext);
+    String argumentString =
+        function.getArgumentsList().stream()
+            .map(argumentConverter::apply)
+            .collect(Collectors.joining(","));
+
+    return "SUM(DIV(" + argumentString + "))";
+  }
+
+  private Function updateFunctionForAvgRate(Function function, ExecutionContext executionContext) {
+
+    Expression columnName = function.getArgumentsList().get(0);
+    Expression literal = function.getArgumentsList().get(1);
+
+    Builder builder = function.toBuilder();
+    builder.clearArguments();
+    builder.addArguments(columnName);
+    builder.addArguments(getUpdatedLiteral(literal, executionContext));
+
+    return builder.build();
+  }
+
+  private Expression getUpdatedLiteral(Expression literal, ExecutionContext executionContext) {
+
+    long rateIntervalInSeconds = literal.getLiteral().getValue().getLong();
+    double aggregateIntervalInSeconds =
+        executionContext.getTimeSeriesPeriod().orElse(executionContext.getTimeRangeDuration());
+
+    return Expression.newBuilder()
+        .setLiteral(
+            LiteralConstant.newBuilder()
+                .setValue(
+                    Value.newBuilder()
+                        .setDouble(aggregateIntervalInSeconds / rateIntervalInSeconds)
+                        .setValueType(ValueType.DOUBLE)
+                        .build())
+                .build())
+        .build();
+  }
+
   private String functionToString(
       Function function, java.util.function.Function<Expression, String> argumentConverter) {
     String argumentString =