feat : Addition of API support for receiving attribute expressions

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

@@ -1,5 +1,8 @@
 package org.hypertrace.core.query.service;
 
+import static org.hypertrace.core.query.service.QueryRequestUtil.getAlias;
+import static org.hypertrace.core.query.service.QueryRequestUtil.getLogicalColumnName;
+
 import com.google.common.base.Supplier;
 import com.google.common.base.Suppliers;
 import java.time.Duration;
@@ -15,7 +18,6 @@
 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;
@@ -137,26 +139,9 @@ private ColumnMetadata toColumnMetadata(Expression expression) {
     ValueCase valueCase = expression.getValueCase();
     switch (valueCase) {
       case COLUMNIDENTIFIER:
-        ColumnIdentifier columnIdentifier = expression.getColumnIdentifier();
-        String alias = columnIdentifier.getAlias();
-        if (alias != null && alias.trim().length() > 0) {
-          builder.setColumnName(alias);
-        } else {
-          builder.setColumnName(columnIdentifier.getColumnName());
-        }
-        builder.setValueType(ValueType.STRING);
-        builder.setIsRepeated(false);
-        break;
+      case ATTRIBUTE_EXPRESSION:
       case FUNCTION:
-        Function function = expression.getFunction();
-        alias = function.getAlias();
-        if (alias != null && alias.trim().length() > 0) {
-          builder.setColumnName(alias);
-        } else {
-          // todo: handle recursive functions max(rollup(time,50)
-          // workaround is to use alias for now
-          builder.setColumnName(function.getFunctionName());
-        }
+        builder.setColumnName(getAlias(expression));
         builder.setValueType(ValueType.STRING);
         builder.setIsRepeated(false);
         break;
@@ -172,8 +157,8 @@ private void extractColumns(List<String> columns, Expression expression) {
     ValueCase valueCase = expression.getValueCase();
     switch (valueCase) {
       case COLUMNIDENTIFIER:
-        ColumnIdentifier columnIdentifier = expression.getColumnIdentifier();
-        columns.add(columnIdentifier.getColumnName());
+      case ATTRIBUTE_EXPRESSION:
+        columns.add(getLogicalColumnName(expression));
         break;
       case LITERAL:
         // no columns
@@ -233,7 +218,13 @@ private Optional<QueryTimeRange> buildQueryTimeRange(Filter filter, String timeF
   }
 
   private boolean isMatchingFilter(Filter filter, String column, Collection<Operator> operators) {
-    return column.equals(filter.getLhs().getColumnIdentifier().getColumnName())
+
+    // if not leaf filter, then return false
+    if (!filter.hasLhs()) {
+      return false;
+    }
+
+    return column.equals(getLogicalColumnName(filter.getLhs()))
         && (operators.stream()
             .anyMatch(operator -> Objects.equals(operator, filter.getOperator())));
   }

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

@@ -2,6 +2,7 @@
 
 import static org.hypertrace.core.query.service.api.Expression.ValueCase.ATTRIBUTE_EXPRESSION;
 import static org.hypertrace.core.query.service.api.Expression.ValueCase.COLUMNIDENTIFIER;
+import static org.hypertrace.core.query.service.api.Expression.ValueCase.FUNCTION;
 
 import org.hypertrace.core.query.service.api.AttributeExpression;
 import org.hypertrace.core.query.service.api.ColumnIdentifier;
@@ -139,4 +140,32 @@ public static String getLogicalColumnName(Expression expression) {
                 + " expression type only");
     }
   }
+
+  public static String getAlias(Expression expression) {
+    switch (expression.getValueCase()) {
+      case COLUMNIDENTIFIER:
+        return expression.getColumnIdentifier().getAlias().isBlank()
+            ? getLogicalColumnName(expression)
+            : expression.getColumnIdentifier().getAlias();
+      case ATTRIBUTE_EXPRESSION:
+        return expression.getAttributeExpression().getAlias().isBlank()
+            ? getLogicalColumnName(expression)
+            : expression.getAttributeExpression().getAlias();
+      case FUNCTION:
+        // todo: handle recursive functions max(rollup(time,50)
+        // workaround is to use alias for now
+        return expression.getFunction().getAlias().isBlank()
+            ? expression.getFunction().getFunctionName()
+            : expression.getFunction().getAlias();
+      default:
+        throw new IllegalArgumentException(
+            "Supports "
+                + ATTRIBUTE_EXPRESSION
+                + " , "
+                + FUNCTION
+                + " and "
+                + COLUMNIDENTIFIER
+                + " expression type only");
+    }
+  }
 }

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

@@ -1,6 +1,7 @@
 package org.hypertrace.core.query.service.pinot;
 
 import static org.hypertrace.core.query.service.ConfigUtils.optionallyGet;
+import static org.hypertrace.core.query.service.QueryRequestUtil.getLogicalColumnName;
 
 import com.google.common.base.Preconditions;
 import com.google.common.base.Stopwatch;
@@ -225,10 +226,12 @@ && rhsHasLongValue(filter.getRhs())) {
   }
 
   private boolean lhsIsStartTimeAttribute(Expression lhs) {
-    return lhs.hasColumnIdentifier()
-        && startTimeAttributeName
-            .map(attributeName -> attributeName.equals(lhs.getColumnIdentifier().getColumnName()))
-            .orElse(false);
+    if (lhs.hasColumnIdentifier() || lhs.hasAttributeExpression()) {
+      return startTimeAttributeName
+          .map(attributeName -> attributeName.equals(getLogicalColumnName(lhs)))
+          .orElse(false);
+    }
+    return false;
   }
 
   private boolean rhsHasLongValue(Expression rhs) {
@@ -245,7 +248,7 @@ private Set<String> getMatchingViewFilterColumns(
     // return it.
     if (filter.getChildFilterCount() == 0) {
       return doesSingleViewFilterMatchLeafQueryFilter(viewFilterMap, filter)
-          ? Set.of(filter.getLhs().getColumnIdentifier().getColumnName())
+          ? Set.of(getLogicalColumnName(filter.getLhs()))
           : Set.of();
     } else {
       // 2. Internal filter node. Recursively get the matching nodes from children.
@@ -274,14 +277,15 @@ private Set<String> getMatchingViewFilterColumns(
    */
   private boolean doesSingleViewFilterMatchLeafQueryFilter(
       Map<String, ViewColumnFilter> viewFilterMap, Filter queryFilter) {
-    if (queryFilter.getLhs().getValueCase() != ValueCase.COLUMNIDENTIFIER) {
+    if (queryFilter.getLhs().getValueCase() != ValueCase.COLUMNIDENTIFIER
+        && queryFilter.getLhs().getValueCase() != ValueCase.ATTRIBUTE_EXPRESSION) {
       return false;
     }
     if (queryFilter.getOperator() != Operator.IN && queryFilter.getOperator() != Operator.EQ) {
       return false;
     }
 
-    String columnName = queryFilter.getLhs().getColumnIdentifier().getColumnName();
+    String columnName = getLogicalColumnName(queryFilter.getLhs());
     ViewColumnFilter viewColumnFilter = viewFilterMap.get(columnName);
     if (viewColumnFilter == null) {
       return false;
@@ -469,7 +473,7 @@ private Filter removeViewColumnFilter(
   private Filter rewriteLeafFilter(
       Filter queryFilter, Map<String, ViewColumnFilter> columnFilterMap) {
     ViewColumnFilter viewColumnFilter =
-        columnFilterMap.get(queryFilter.getLhs().getColumnIdentifier().getColumnName());
+        columnFilterMap.get(getLogicalColumnName(queryFilter.getLhs()));
     // If the RHS of both the view filter and query filter match, return empty filter.
     if (viewColumnFilter != null && isEquals(viewColumnFilter.getValues(), queryFilter.getRhs())) {
       return Filter.getDefaultInstance();

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

@@ -30,6 +30,7 @@
 import org.hypertrace.core.query.service.QueryFunctionConstants;
 import org.hypertrace.core.query.service.QueryRequestUtil;
 import org.hypertrace.core.query.service.QueryTransformation;
+import org.hypertrace.core.query.service.api.AttributeExpression;
 import org.hypertrace.core.query.service.api.ColumnIdentifier;
 import org.hypertrace.core.query.service.api.Expression;
 import org.hypertrace.core.query.service.api.Filter;
@@ -79,6 +80,8 @@ private Single<Expression> transformExpression(Expression expression) {
     switch (expression.getValueCase()) {
       case COLUMNIDENTIFIER:
         return this.transformColumnIdentifier(expression.getColumnIdentifier());
+      case ATTRIBUTE_EXPRESSION:
+        return this.transformAttributeExpression(expression.getAttributeExpression());
       case FUNCTION:
         return this.transformFunction(expression.getFunction())
             .map(expression.toBuilder()::setFunction)
@@ -96,10 +99,19 @@ private Single<Expression> transformExpression(Expression expression) {
 
   private Single<Expression> transformColumnIdentifier(ColumnIdentifier columnIdentifier) {
     return this.projectAttributeIfPossible(columnIdentifier.getColumnName())
-        .map(expression -> this.aliasToMatchOriginal(columnIdentifier, expression))
+        .map(expression -> this.aliasToMatchOriginal(getOriginalKey(columnIdentifier), expression))
         .defaultIfEmpty(Expression.newBuilder().setColumnIdentifier(columnIdentifier).build());
   }
 
+  private Single<Expression> transformAttributeExpression(AttributeExpression attributeExpression) {
+    return this.projectAttributeIfPossible(attributeExpression.getAttributeId())
+        .map(
+            expression ->
+                this.aliasToMatchOriginal(getOriginalKey(attributeExpression), expression))
+        .defaultIfEmpty(
+            Expression.newBuilder().setAttributeExpression(attributeExpression).build());
+  }
+
   private Single<Function> transformFunction(Function function) {
     return this.transformExpressionList(function.getArgumentsList())
         .map(expressions -> function.toBuilder().clearArguments().addAllArguments(expressions))
@@ -267,15 +279,18 @@ private Single<String> convertOperator(ProjectionOperator operator) {
     }
   }
 
-  private Expression aliasToMatchOriginal(ColumnIdentifier original, Expression newExpression) {
-    String originalKey =
-        original.getAlias().isEmpty() ? original.getColumnName() : original.getAlias();
+  private Expression aliasToMatchOriginal(String originalKey, Expression newExpression) {
     switch (newExpression.getValueCase()) {
       case COLUMNIDENTIFIER:
         return newExpression.toBuilder()
             .setColumnIdentifier(
                 newExpression.getColumnIdentifier().toBuilder().setAlias(originalKey))
             .build();
+      case ATTRIBUTE_EXPRESSION:
+        return newExpression.toBuilder()
+            .setAttributeExpression(
+                newExpression.getAttributeExpression().toBuilder().setAlias(originalKey))
+            .build();
       case FUNCTION:
         return newExpression.toBuilder()
             .setFunction(newExpression.getFunction().toBuilder().setAlias(originalKey))
@@ -417,4 +432,14 @@ private List<Filter> createFilterForComplexAttributeExpressionFromOrderBy(
         .map(QueryRequestUtil::createContainsKeyFilter)
         .collect(Collectors.toList());
   }
+
+  private String getOriginalKey(AttributeExpression attributeExpression) {
+    String alias = attributeExpression.getAlias();
+    return alias.isEmpty() ? attributeExpression.getAttributeId() : alias;
+  }
+
+  private String getOriginalKey(ColumnIdentifier columnIdentifier) {
+    String alias = columnIdentifier.getAlias();
+    return alias.isEmpty() ? columnIdentifier.getColumnName() : alias;
+  }
 }

query-service-impl/src/test/java/org/hypertrace/core/query/service/QueryRequestBuilderUtils.java

@@ -157,6 +157,13 @@ public static Filter createSimpleAttributeFilter(
     return createFilter(createSimpleAttributeExpression(column).build(), operator, expression);
   }
 
+  public static Filter createSimpleAttributeFilter(String column, Operator operator, String value) {
+    return createFilter(
+        createSimpleAttributeExpression(column).build(),
+        operator,
+        createStringLiteralValueExpression(value));
+  }
+
   public static Filter createFilter(
       Expression columnExpression, Operator operator, Expression expression) {
     return Filter.newBuilder()

query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/ExecutionContextTest.java

@@ -3,6 +3,7 @@
 import static org.hypertrace.core.query.service.QueryRequestBuilderUtils.createFilter;
 import static org.hypertrace.core.query.service.QueryRequestBuilderUtils.createTimeColumnGroupByExpression;
 import static org.hypertrace.core.query.service.QueryRequestBuilderUtils.createTimeFilter;
+import static org.hypertrace.core.query.service.QueryRequestUtil.createSimpleAttributeExpression;
 import static org.hypertrace.core.query.service.QueryRequestUtil.createStringLiteralValueExpression;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -320,6 +321,78 @@ public void testSelectionsLinkedHashSet() {
         Expression.newBuilder().setFunction(minFunction).build(), selectionsIterator.next());
   }
 
+  @Test
+  public void testSelectionsLinkedHashSetWithAttributeExpression() {
+    Builder builder = QueryRequest.newBuilder();
+    // agg function with alias
+    Function count =
+        Function.newBuilder()
+            .setFunctionName("Count")
+            .setAlias("myCountAlias")
+            .addArguments(createSimpleAttributeExpression("Trace.id"))
+            .build();
+    builder.addAggregation(Expression.newBuilder().setFunction(count));
+
+    // agg function without alias
+    Function minFunction =
+        Function.newBuilder()
+            .setFunctionName("MIN")
+            .addArguments(createSimpleAttributeExpression("Trace.duration"))
+            .build();
+    builder.addAggregation(Expression.newBuilder().setFunction(minFunction));
+
+    // Add some selections
+    builder.addSelection(createSimpleAttributeExpression("Trace.transaction_name"));
+    builder.addSelection(createSimpleAttributeExpression("Trace.id"));
+
+    // A function added into selections list is treated as a selection
+    Function avg =
+        Function.newBuilder()
+            .setFunctionName("AVG")
+            .setAlias("myAvgAlias")
+            .addArguments(createSimpleAttributeExpression("Trace.duration"))
+            .build();
+    builder.addSelection(Expression.newBuilder().setFunction(avg));
+
+    // Add some group bys
+    builder.addGroupBy(createSimpleAttributeExpression("Trace.api_name"));
+    builder.addGroupBy(createSimpleAttributeExpression("Trace.service_name"));
+    QueryRequest queryRequest = builder.build();
+
+    ExecutionContext context = new ExecutionContext("test", queryRequest);
+
+    // The order in resultSetMetadata.getColumnMetadataList() and selections is group bys,
+    // selections then aggregations
+    ResultSetMetadata resultSetMetadata = context.getResultSetMetadata();
+
+    assertNotNull(resultSetMetadata);
+    assertEquals(7, resultSetMetadata.getColumnMetadataCount());
+    assertEquals("Trace.api_name", resultSetMetadata.getColumnMetadata(0).getColumnName());
+    assertEquals("Trace.service_name", resultSetMetadata.getColumnMetadata(1).getColumnName());
+    assertEquals("Trace.transaction_name", resultSetMetadata.getColumnMetadata(2).getColumnName());
+    assertEquals("Trace.id", resultSetMetadata.getColumnMetadata(3).getColumnName());
+    assertEquals("myAvgAlias", resultSetMetadata.getColumnMetadata(4).getColumnName());
+    assertEquals("myCountAlias", resultSetMetadata.getColumnMetadata(5).getColumnName());
+    assertEquals("MIN", resultSetMetadata.getColumnMetadata(6).getColumnName());
+
+    // Selections should correspond in size and order to the
+    // resultSetMetadata.getColumnMetadataList()
+    assertEquals(7, context.getAllSelections().size());
+    Iterator<Expression> selectionsIterator = context.getAllSelections().iterator();
+    assertEquals(
+        createSimpleAttributeExpression("Trace.api_name").build(), selectionsIterator.next());
+    assertEquals(
+        createSimpleAttributeExpression("Trace.service_name").build(), selectionsIterator.next());
+    assertEquals(
+        createSimpleAttributeExpression("Trace.transaction_name").build(),
+        selectionsIterator.next());
+    assertEquals(createSimpleAttributeExpression("Trace.id").build(), selectionsIterator.next());
+    assertEquals(Expression.newBuilder().setFunction(avg).build(), selectionsIterator.next());
+    assertEquals(Expression.newBuilder().setFunction(count).build(), selectionsIterator.next());
+    assertEquals(
+        Expression.newBuilder().setFunction(minFunction).build(), selectionsIterator.next());
+  }
+
   @Test
   public void testSetTimeSeriesPeriod() {