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,11 @@
 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 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 com.google.common.base.Supplier;
 import com.google.common.base.Suppliers;
 import java.time.Duration;
@@ -15,7 +21,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 +142,10 @@ 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());
-        }
+        String alias = getAlias(expression).orElseThrow(IllegalArgumentException::new);
+        builder.setColumnName(alias);
         builder.setValueType(ValueType.STRING);
         builder.setIsRepeated(false);
         break;
@@ -172,8 +161,10 @@ 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:
+        String logicalColumnName =
+            getLogicalColumnName(expression).orElseThrow(IllegalArgumentException::new);
+        columns.add(logicalColumnName);
         break;
       case LITERAL:
         // no columns
@@ -233,7 +224,7 @@ 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())
+    return getLogicalColumnName(filter.getLhs()).map(column::equals).orElse(false)
         && (operators.stream()
             .anyMatch(operator -> Objects.equals(operator, filter.getOperator())));
   }

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

@@ -2,7 +2,9 @@
 
 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 java.util.Optional;
 import org.hypertrace.core.query.service.api.AttributeExpression;
 import org.hypertrace.core.query.service.api.ColumnIdentifier;
 import org.hypertrace.core.query.service.api.Expression;
@@ -124,19 +126,38 @@ public static boolean isDateTimeFunction(Expression expression) {
         && expression.getFunction().getFunctionName().equals("dateTimeConvert");
   }
 
-  public static String getLogicalColumnName(Expression expression) {
+  public static Optional<String> getLogicalColumnName(Expression expression) {
     switch (expression.getValueCase()) {
       case COLUMNIDENTIFIER:
-        return expression.getColumnIdentifier().getColumnName();
+        return Optional.of(expression.getColumnIdentifier().getColumnName());
       case ATTRIBUTE_EXPRESSION:
-        return expression.getAttributeExpression().getAttributeId();
+        return Optional.of(expression.getAttributeExpression().getAttributeId());
       default:
-        throw new IllegalArgumentException(
-            "Supports "
-                + ATTRIBUTE_EXPRESSION
-                + " and "
-                + COLUMNIDENTIFIER
-                + " expression type only");
+        return Optional.empty();
+    }
+  }
+
+  public static Optional<String> getAlias(Expression expression) {
+    switch (expression.getValueCase()) {
+      case COLUMNIDENTIFIER:
+        return Optional.of(
+            expression.getColumnIdentifier().getAlias().isBlank()
+                ? getLogicalColumnName(expression).get()
+                : expression.getColumnIdentifier().getAlias());
+      case ATTRIBUTE_EXPRESSION:
+        return Optional.of(
+            expression.getAttributeExpression().getAlias().isBlank()
+                ? getLogicalColumnName(expression).get()
+                : expression.getAttributeExpression().getAlias());
+      case FUNCTION:
+        // todo: handle recursive functions max(rollup(time,50)
+        // workaround is to use alias for now
+        return Optional.of(
+            expression.getFunction().getAlias().isBlank()
+                ? expression.getFunction().getFunctionName()
+                : expression.getFunction().getAlias());
+      default:
+        return Optional.empty();
     }
   }
 }

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,8 @@ && rhsHasLongValue(filter.getRhs())) {
   }
 
   private boolean lhsIsStartTimeAttribute(Expression lhs) {
-    return lhs.hasColumnIdentifier()
-        && startTimeAttributeName
-            .map(attributeName -> attributeName.equals(lhs.getColumnIdentifier().getColumnName()))
-            .orElse(false);
+    return startTimeAttributeName.isPresent()
+        && startTimeAttributeName.equals(getLogicalColumnName(lhs));
   }
 
   private boolean rhsHasLongValue(Expression rhs) {
@@ -245,7 +244,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()).orElseThrow(IllegalArgumentException::new))
           : Set.of();
     } else {
       // 2. Internal filter node. Recursively get the matching nodes from children.
@@ -274,15 +273,13 @@ private Set<String> getMatchingViewFilterColumns(
    */
   private boolean doesSingleViewFilterMatchLeafQueryFilter(
       Map<String, ViewColumnFilter> viewFilterMap, Filter queryFilter) {
-    if (queryFilter.getLhs().getValueCase() != ValueCase.COLUMNIDENTIFIER) {
-      return false;
-    }
+
     if (queryFilter.getOperator() != Operator.IN && queryFilter.getOperator() != Operator.EQ) {
       return false;
     }
 
-    String columnName = queryFilter.getLhs().getColumnIdentifier().getColumnName();
-    ViewColumnFilter viewColumnFilter = viewFilterMap.get(columnName);
+    ViewColumnFilter viewColumnFilter =
+        viewFilterMap.get(getLogicalColumnName(queryFilter.getLhs()).orElse(null));
     if (viewColumnFilter == null) {
       return false;
     }
@@ -469,7 +466,8 @@ 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()).orElseThrow(IllegalArgumentException::new));
     // 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/pinot/QueryRequestToPinotSQLConverter.java

@@ -211,7 +211,7 @@ private Expression handleValueConversionForLiteralExpression(Expression lhs, Exp
       return rhs;
     }
 
-    String lhsColumnName = getLogicalColumnName(lhs);
+    String lhsColumnName = getLogicalColumnName(lhs).orElseThrow(IllegalArgumentException::new);
     try {
       Value value =
           DestinationColumnValueConverter.INSTANCE.convert(
@@ -272,7 +272,8 @@ private String convertExpressionToString(
       case COLUMNIDENTIFIER:
         // this takes care of the Map Type where it's split into 2 columns
         List<String> columnNames =
-            viewDefinition.getPhysicalColumnNames(getLogicalColumnName(expression));
+            viewDefinition.getPhysicalColumnNames(
+                getLogicalColumnName(expression).orElseThrow(IllegalArgumentException::new));
         return joiner.join(columnNames);
       case ATTRIBUTE_EXPRESSION:
         if (isAttributeExpressionWithSubpath(expression)) {
@@ -292,7 +293,9 @@ private String convertExpressionToString(
               valCol);
         } else {
           // this takes care of the Map Type where it's split into 2 columns
-          columnNames = viewDefinition.getPhysicalColumnNames(getLogicalColumnName(expression));
+          columnNames =
+              viewDefinition.getPhysicalColumnNames(
+                  getLogicalColumnName(expression).orElseThrow(IllegalArgumentException::new));
           return joiner.join(columnNames);
         }
       case LITERAL:
@@ -313,15 +316,19 @@ private String convertExpressionToString(
   }
 
   private String convertExpressionToMapKeyColumn(Expression expression) {
-    String col = viewDefinition.getKeyColumnNameForMap(getLogicalColumnName(expression));
+    String col =
+        viewDefinition.getKeyColumnNameForMap(
+            getLogicalColumnName(expression).orElseThrow(IllegalArgumentException::new));
     if (col != null && col.length() > 0) {
       return col;
     }
     throw new IllegalArgumentException("operator supports multi value column only");
   }
 
   private String convertExpressionToMapValueColumn(Expression expression) {
-    String col = viewDefinition.getValueColumnNameForMap(getLogicalColumnName(expression));
+    String col =
+        viewDefinition.getValueColumnNameForMap(
+            getLogicalColumnName(expression).orElseThrow(IllegalArgumentException::new));
     if (col != null && col.length() > 0) {
       return col;
     }

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

@@ -17,6 +17,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import javax.inject.Inject;
 import org.hypertrace.core.attribute.service.cachingclient.CachingAttributeClient;
 import org.hypertrace.core.attribute.service.projection.AttributeProjection;
@@ -30,6 +31,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 +81,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 +100,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 +280,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))
@@ -329,7 +345,8 @@ private QueryRequest rebuildRequest(
       List<OrderByExpression> orderBys) {
 
     QueryRequest.Builder builder = original.toBuilder();
-    Filter updatedFilter = rebuildFilterForComplexAttributeExpression(originalFilter, orderBys);
+    Filter updatedFilter =
+        rebuildFilterForComplexAttributeExpression(originalFilter, orderBys, selections);
 
     if (Filter.getDefaultInstance().equals(updatedFilter)) {
       builder.clearFilter();
@@ -350,16 +367,20 @@ private QueryRequest rebuildRequest(
   }
 
   /*
-   * We need the CONTAINS_KEY filter in all filters and order bys dealing with complex
+   * We need the CONTAINS_KEY filter in all filters, selections and order bys dealing with complex
    * attribute expressions as Pinot gives error if particular key is absent. Rest all work fine.
-   * To handle order bys, we add the corresponding filter at the top and 'AND' it with the main filter.
+   * To handle order bys and selections, we add the corresponding filter at the top and 'AND' it with the main filter.
    * To handle filter, we modify each filter (say filter1) as : "CONTAINS_KEY AND filter1".
    */
   private Filter rebuildFilterForComplexAttributeExpression(
-      Filter originalFilter, List<OrderByExpression> orderBys) {
+      Filter originalFilter, List<OrderByExpression> orderBys, List<Expression> selections) {
 
     Filter updatedFilter = updateFilterForComplexAttributeExpressionFromFilter(originalFilter);
-    List<Filter> filterList = createFilterForComplexAttributeExpressionFromOrderBy(orderBys);
+    List<Filter> filterList =
+        Stream.concat(
+                createFilterForComplexAttributeExpressionFromOrderBy(orderBys),
+                createFilterForComplexAttributeExpressionFromSelection(selections))
+            .collect(Collectors.toList());
 
     if (filterList.isEmpty()) {
       return updatedFilter;
@@ -408,13 +429,40 @@ private Filter updateFilterForComplexAttributeExpressionFromFilter(Filter origin
     }
   }
 
-  private List<Filter> createFilterForComplexAttributeExpressionFromOrderBy(
+  private Stream<Filter> createFilterForComplexAttributeExpressionFromOrderBy(
       List<OrderByExpression> orderByExpressionList) {
     return orderByExpressionList.stream()
         .map(OrderByExpression::getExpression)
         .filter(QueryRequestUtil::isAttributeExpressionWithSubpath)
         .map(Expression::getAttributeExpression)
-        .map(QueryRequestUtil::createContainsKeyFilter)
-        .collect(Collectors.toList());
+        .map(QueryRequestUtil::createContainsKeyFilter);
+  }
+
+  private Stream<Filter> createFilterForComplexAttributeExpressionFromSelection(
+      List<Expression> selections) {
+    return selections.stream()
+        .flatMap(this::getAnyAttributeExpression)
+        .map(QueryRequestUtil::createContainsKeyFilter);
+  }
+
+  private Stream<AttributeExpression> getAnyAttributeExpression(Expression selection) {
+    if (selection.hasFunction()) {
+      return selection.getFunction().getArgumentsList().stream()
+          .flatMap(this::getAnyAttributeExpression);
+    } else {
+      return Stream.of(selection)
+          .filter(QueryRequestUtil::isAttributeExpressionWithSubpath)
+          .map(Expression::getAttributeExpression);
+    }
+  }
+
+  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/main/java/org/hypertrace/core/query/service/prometheus/FilterToPromqlConverter.java

@@ -28,7 +28,8 @@ void convertFilterToString(
       }
     } else {
       if (QueryRequestUtil.isSimpleAttributeExpression(filter.getLhs())
-          && timeFilterColumn.equals(getLogicalColumnName(filter.getLhs()))) {
+          && timeFilterColumn.equals(
+              getLogicalColumnName(filter.getLhs()).orElseThrow(IllegalArgumentException::new))) {
         return;
       }
       StringBuilder builder = new StringBuilder();