Multi value attribute grouping support

entity-service-impl/build.gradle.kts

@@ -11,7 +11,7 @@ dependencies {
   annotationProcessor("org.projectlombok:lombok:1.18.22")
   compileOnly("org.projectlombok:lombok:1.18.18")
 
-  implementation("org.hypertrace.core.documentstore:document-store:0.6.12")
+  implementation("org.hypertrace.core.documentstore:document-store:0.6.13")
   implementation("org.hypertrace.core.grpcutils:grpc-context-utils:0.7.0")
   implementation("org.hypertrace.core.grpcutils:grpc-client-utils:0.7.0")
   implementation("org.hypertrace.core.attribute.service:caching-attribute-service-client:0.12.3")

entity-service-impl/src/main/java/org/hypertrace/entity/query/service/converter/ConverterModule.java

@@ -6,6 +6,7 @@
 import java.util.List;
 import lombok.AllArgsConstructor;
 import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
+import org.hypertrace.core.documentstore.expression.type.FromTypeExpression;
 import org.hypertrace.core.documentstore.query.Query;
 import org.hypertrace.core.documentstore.query.Sort;
 import org.hypertrace.entity.query.service.EntityAttributeMapping;
@@ -16,6 +17,7 @@
 import org.hypertrace.entity.query.service.converter.response.ResponseModule;
 import org.hypertrace.entity.query.service.converter.selection.SelectionConverterModule;
 import org.hypertrace.entity.query.service.v1.EntityQueryRequest;
+import org.hypertrace.entity.query.service.v1.Expression;
 import org.hypertrace.entity.query.service.v1.LiteralConstant;
 import org.hypertrace.entity.query.service.v1.OrderByExpression;
 
@@ -36,6 +38,8 @@ protected void configure() {
         .to(ConstantExpressionConverter.class);
     bind(new TypeLiteral<Converter<List<OrderByExpression>, Sort>>() {}).to(OrderByConverter.class);
     bind(new TypeLiteral<Converter<EntityQueryRequest, Query>>() {}).to(QueryConverter.class);
+    bind(new TypeLiteral<Converter<List<Expression>, List<FromTypeExpression>>>() {})
+        .to(FromClauseConverter.class);
   }
 
   @Provides

entity-service-impl/src/main/java/org/hypertrace/entity/query/service/converter/FromClauseConverter.java

@@ -0,0 +1,58 @@
+package org.hypertrace.entity.query.service.converter;
+
+import static java.util.Optional.empty;
+import static java.util.stream.Collectors.toUnmodifiableList;
+import static org.hypertrace.entity.query.service.v1.Expression.ValueCase.COLUMNIDENTIFIER;
+
+import java.util.HashSet;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import javax.inject.Inject;
+import lombok.AllArgsConstructor;
+import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression;
+import org.hypertrace.core.documentstore.expression.impl.UnnestExpression;
+import org.hypertrace.core.documentstore.expression.type.FromTypeExpression;
+import org.hypertrace.core.grpcutils.context.RequestContext;
+import org.hypertrace.entity.query.service.EntityAttributeMapping;
+import org.hypertrace.entity.query.service.converter.accessor.OneOfAccessor;
+import org.hypertrace.entity.query.service.v1.ColumnIdentifier;
+import org.hypertrace.entity.query.service.v1.Expression;
+import org.hypertrace.entity.query.service.v1.Expression.ValueCase;
+
+@AllArgsConstructor(onConstructor_ = {@Inject})
+public class FromClauseConverter implements Converter<List<Expression>, List<FromTypeExpression>> {
+  private final OneOfAccessor<Expression, ValueCase> expressionAccessor;
+  private final Converter<ColumnIdentifier, IdentifierExpression> identifierExpressionConverter;
+  private final EntityAttributeMapping entityAttributeMapping;
+
+  @Override
+  public List<FromTypeExpression> convert(
+      final List<Expression> expressions, final RequestContext requestContext)
+      throws ConversionException {
+    final Set<FromTypeExpression> set = new HashSet<>();
+
+    for (final Expression expression : expressions) {
+      final Optional<FromTypeExpression> optionalExpression = convert(expression, requestContext);
+      optionalExpression.ifPresent(set::add);
+    }
+
+    return set.stream().collect(toUnmodifiableList());
+  }
+
+  private Optional<FromTypeExpression> convert(
+      final Expression expression, final RequestContext requestContext) throws ConversionException {
+    final ColumnIdentifier identifier =
+        expressionAccessor.access(expression, expression.getValueCase(), Set.of(COLUMNIDENTIFIER));
+
+    if (!entityAttributeMapping.isMultiValued(requestContext, identifier.getColumnName())) {
+      return empty();
+    }
+
+    final IdentifierExpression identifierExpression =
+        identifierExpressionConverter.convert(identifier, requestContext);
+    final FromTypeExpression fromTypeExpression = UnnestExpression.of(identifierExpression, false);
+
+    return Optional.of(fromTypeExpression);
+  }
+}

entity-service-impl/src/main/java/org/hypertrace/entity/query/service/converter/QueryConverter.java

@@ -1,10 +1,14 @@
 package org.hypertrace.entity.query.service.converter;
 
+import static java.util.Collections.unmodifiableList;
+
 import com.google.inject.Inject;
 import com.google.inject.Singleton;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Function;
 import lombok.AllArgsConstructor;
+import org.hypertrace.core.documentstore.expression.type.FromTypeExpression;
 import org.hypertrace.core.documentstore.query.Aggregation;
 import org.hypertrace.core.documentstore.query.Filter;
 import org.hypertrace.core.documentstore.query.Pagination;
@@ -14,6 +18,7 @@
 import org.hypertrace.core.documentstore.query.Selection;
 import org.hypertrace.core.documentstore.query.Sort;
 import org.hypertrace.core.grpcutils.context.RequestContext;
+import org.hypertrace.entity.query.service.converter.aggregation.AggregationColumnProvider;
 import org.hypertrace.entity.query.service.v1.EntityQueryRequest;
 import org.hypertrace.entity.query.service.v1.Expression;
 import org.hypertrace.entity.query.service.v1.OrderByExpression;
@@ -24,6 +29,8 @@ public class QueryConverter implements Converter<EntityQueryRequest, Query> {
   private final Converter<List<Expression>, Selection> selectionConverter;
   private final Converter<EntityQueryRequest, Filter> filterConverter;
 
+  private final AggregationColumnProvider aggregationColumnProvider;
+  private final Converter<List<Expression>, List<FromTypeExpression>> fromClauseConverter;
   private final Converter<List<Expression>, Aggregation> groupByConverter;
 
   private final Converter<List<OrderByExpression>, Sort> orderByConverter;
@@ -40,6 +47,12 @@ public Query convert(final EntityQueryRequest request, final RequestContext requ
     final Filter filter = filterConverter.convert(request, requestContext);
     builder.setFilter(filter);
 
+    setFieldIfNotEmpty(
+        getExpressionsForFromClause(request),
+        builder::addFromClauses,
+        fromClauseConverter,
+        requestContext);
+
     setFieldIfNotEmpty(
         request.getGroupByList(), builder::setAggregation, groupByConverter, requestContext);
     setFieldIfNotEmpty(
@@ -67,4 +80,20 @@ private <T, U> void setFieldIfNotEmpty(
     final T converted = converter.convert(list, requestContext);
     setter.apply(converted);
   }
+
+  private List<Expression> getExpressionsForFromClause(final EntityQueryRequest request)
+      throws ConversionException {
+    final List<Expression> list = new ArrayList<>(request.getGroupByList());
+
+    for (Expression expression : request.getSelectionList()) {
+      if (expression.hasFunction()) {
+        org.hypertrace.entity.query.service.v1.Function function = expression.getFunction();
+        final Expression aggregationColumn =
+            aggregationColumnProvider.getAggregationColumn(function);
+        list.add(aggregationColumn);
+      }
+    }
+
+    return unmodifiableList(list);
+  }
 }

entity-service-impl/src/main/java/org/hypertrace/entity/query/service/converter/aggregation/AggregationAliasProvider.java

@@ -4,7 +4,6 @@
 
 import com.google.inject.Inject;
 import com.google.inject.Singleton;
-import java.util.List;
 import java.util.Set;
 import lombok.AllArgsConstructor;
 import org.hypertrace.entity.query.service.converter.AliasProvider;
@@ -20,20 +19,15 @@
 @AllArgsConstructor(onConstructor_ = {@Inject})
 public class AggregationAliasProvider implements AliasProvider<Function> {
   private final AliasProvider<ColumnIdentifier> identifierAliasProvider;
+  private final AggregationColumnProvider aggregationColumnProvider;
   private final OneOfAccessor<Expression, ValueCase> expressionAccessor;
 
   @Override
   public String getAlias(final Function aggregateExpression) throws ConversionException {
-    final List<Expression> innerExpressions = aggregateExpression.getArgumentsList();
-
-    if (innerExpressions.size() != 1) {
-      throw new ConversionException("Aggregation function should have exactly one argument");
-    }
-
-    final Expression innerExpression = innerExpressions.get(0);
+    final Expression expression =
+        aggregationColumnProvider.getAggregationColumn(aggregateExpression);
     final ColumnIdentifier containingIdentifier =
-        expressionAccessor.access(
-            innerExpression, innerExpression.getValueCase(), Set.of(COLUMNIDENTIFIER));
+        expressionAccessor.access(expression, expression.getValueCase(), Set.of(COLUMNIDENTIFIER));
     final String alias = aggregateExpression.getAlias();
 
     if (StringUtils.isNotBlank(alias)) {

entity-service-impl/src/main/java/org/hypertrace/entity/query/service/converter/aggregation/AggregationColumnProvider.java

@@ -0,0 +1,19 @@
+package org.hypertrace.entity.query.service.converter.aggregation;
+
+import java.util.List;
+import org.hypertrace.entity.query.service.converter.ConversionException;
+import org.hypertrace.entity.query.service.v1.Expression;
+import org.hypertrace.entity.query.service.v1.Function;
+
+public class AggregationColumnProvider {
+
+  public final Expression getAggregationColumn(final Function function) throws ConversionException {
+    final List<Expression> innerExpressions = function.getArgumentsList();
+
+    if (innerExpressions.size() != 1) {
+      throw new ConversionException("Aggregation function should have exactly one argument");
+    }
+
+    return innerExpressions.get(0);
+  }
+}

entity-service-impl/src/main/java/org/hypertrace/entity/query/service/converter/identifier/ArrayElementSuffixAddingIdentifierConverter.java

@@ -1,12 +1,14 @@
 package org.hypertrace.entity.query.service.converter.identifier;
 
+import static org.hypertrace.entity.query.service.converter.ValueHelper.VALUES_KEY;
+import static org.hypertrace.entity.query.service.converter.ValueHelper.VALUE_KEY;
 import static org.hypertrace.entity.query.service.converter.ValueHelper.VALUE_LIST_KEY;
 import static org.hypertrace.entity.query.service.converter.filter.FilteringExpressionConverterBase.ARRAY_OPERATORS;
 
 import com.google.inject.Inject;
 import com.google.inject.Singleton;
+import org.hypertrace.entity.query.service.converter.ConversionException;
 import org.hypertrace.entity.query.service.converter.ValueHelper;
-import org.hypertrace.entity.query.service.v1.Operator;
 
 /**
  * Adds suffix .valueList.values.%d.value.<type> for element-by-element comparison.
@@ -15,7 +17,8 @@
  */
 @Singleton
 public class ArrayElementSuffixAddingIdentifierConverter extends SuffixAddingIdentifierConverter {
-  private static final String ARRAY_ELEMENT_SUFFIX = "." + VALUE_LIST_KEY + ".values.%d.value.";
+  private static final String ARRAY_ELEMENT_SUFFIX =
+      "." + VALUE_LIST_KEY + "." + VALUES_KEY + ".%d." + VALUE_KEY + ".";
   private final ArraySuffixAddingIdentifierConverter arraySuffixAddingIdentifierConverter;
 
   @Inject
@@ -27,12 +30,13 @@ public ArrayElementSuffixAddingIdentifierConverter(
   }
 
   @Override
-  protected String getSuffix(final Operator operator) {
-    if (ARRAY_OPERATORS.contains(operator)) {
+  protected String getSuffix(final IdentifierConversionMetadata metadata)
+      throws ConversionException {
+    if (ARRAY_OPERATORS.contains(metadata.getOperator())) {
       // If the operator is an array operator, fall-back to array suffix
-      return arraySuffixAddingIdentifierConverter.getSuffix(operator);
+      return arraySuffixAddingIdentifierConverter.getSuffix(metadata);
     }
 
-    return ARRAY_ELEMENT_SUFFIX;
+    return ARRAY_ELEMENT_SUFFIX + getTypeName(metadata);
   }
 }

entity-service-impl/src/main/java/org/hypertrace/entity/query/service/converter/identifier/ArrayPathSuffixAddingIdentifierConverter.java

@@ -0,0 +1,24 @@
+package org.hypertrace.entity.query.service.converter.identifier;
+
+import static org.hypertrace.entity.query.service.converter.ValueHelper.VALUES_KEY;
+import static org.hypertrace.entity.query.service.converter.ValueHelper.VALUE_LIST_KEY;
+
+import javax.inject.Inject;
+import javax.inject.Singleton;
+import org.hypertrace.entity.query.service.converter.ValueHelper;
+
+/** Adds suffix .valueList.values for unwinding */
+@Singleton
+public class ArrayPathSuffixAddingIdentifierConverter extends SuffixAddingIdentifierConverter {
+  private static final String ARRAY_PATH_SUFFIX = "." + VALUE_LIST_KEY + "." + VALUES_KEY;
+
+  @Inject
+  public ArrayPathSuffixAddingIdentifierConverter(final ValueHelper valueHelper) {
+    super(valueHelper);
+  }
+
+  @Override
+  protected String getSuffix(final IdentifierConversionMetadata metadata) {
+    return ARRAY_PATH_SUFFIX;
+  }
+}

entity-service-impl/src/main/java/org/hypertrace/entity/query/service/converter/identifier/ArraySuffixAddingIdentifierConverter.java

@@ -1,24 +1,28 @@
 package org.hypertrace.entity.query.service.converter.identifier;
 
+import static org.hypertrace.entity.query.service.converter.ValueHelper.VALUES_KEY;
+import static org.hypertrace.entity.query.service.converter.ValueHelper.VALUE_KEY;
 import static org.hypertrace.entity.query.service.converter.ValueHelper.VALUE_LIST_KEY;
 
 import com.google.inject.Inject;
 import com.google.inject.Singleton;
+import org.hypertrace.entity.query.service.converter.ConversionException;
 import org.hypertrace.entity.query.service.converter.ValueHelper;
-import org.hypertrace.entity.query.service.v1.Operator;
 
 /** Adds suffix .valueList.values.value.<type> for direct comparison */
 @Singleton
 public class ArraySuffixAddingIdentifierConverter extends SuffixAddingIdentifierConverter {
-  private static final String ARRAY_SUFFIX = "." + VALUE_LIST_KEY + ".values.value.";
+  private static final String ARRAY_SUFFIX =
+      "." + VALUE_LIST_KEY + "." + VALUES_KEY + "." + VALUE_KEY + ".";
 
   @Inject
   public ArraySuffixAddingIdentifierConverter(final ValueHelper valueHelper) {
     super(valueHelper);
   }
 
   @Override
-  protected String getSuffix(final Operator operator) {
-    return ARRAY_SUFFIX;
+  protected String getSuffix(final IdentifierConversionMetadata metadata)
+      throws ConversionException {
+    return ARRAY_SUFFIX + getTypeName(metadata);
   }
 }

entity-service-impl/src/main/java/org/hypertrace/entity/query/service/converter/identifier/IdentifierExpressionConverter.java

@@ -17,13 +17,26 @@
 public class IdentifierExpressionConverter
     implements Converter<ColumnIdentifier, IdentifierExpression> {
   private final EntityAttributeMapping attributeMapping;
+  private final ArrayPathSuffixAddingIdentifierConverter arrayPathSuffixAddingIdentifierConverter;
 
   @Override
   public IdentifierExpression convert(
       final ColumnIdentifier identifier, final RequestContext requestContext)
       throws ConversionException {
-    final String subDocPath =
-        getSubDocPathById(attributeMapping, identifier.getColumnName(), requestContext);
-    return IdentifierExpression.of(subDocPath);
+    final String columnId = identifier.getColumnName();
+    final String subDocPath = getSubDocPathById(attributeMapping, columnId, requestContext);
+
+    final String suffixedSubDocPath;
+
+    if (attributeMapping.isMultiValued(requestContext, columnId)) {
+      suffixedSubDocPath =
+          arrayPathSuffixAddingIdentifierConverter.convert(
+              IdentifierConversionMetadata.builder().subDocPath(subDocPath).build(),
+              requestContext);
+    } else {
+      suffixedSubDocPath = subDocPath;
+    }
+
+    return IdentifierExpression.of(suffixedSubDocPath);
   }
 }

entity-service-impl/src/main/java/org/hypertrace/entity/query/service/converter/identifier/MapSuffixAddingIdentifierConverter.java

@@ -1,11 +1,13 @@
 package org.hypertrace.entity.query.service.converter.identifier;
 
+import static org.hypertrace.entity.query.service.converter.ValueHelper.VALUES_KEY;
+import static org.hypertrace.entity.query.service.converter.ValueHelper.VALUE_KEY;
 import static org.hypertrace.entity.query.service.converter.ValueHelper.VALUE_MAP_KEY;
 
 import com.google.inject.Inject;
 import com.google.inject.Singleton;
+import org.hypertrace.entity.query.service.converter.ConversionException;
 import org.hypertrace.entity.query.service.converter.ValueHelper;
-import org.hypertrace.entity.query.service.v1.Operator;
 
 /**
  * Adds suffix .valueMap.values.%s.value.<type> for element-by-element comparison.
@@ -14,15 +16,17 @@
  */
 @Singleton
 public class MapSuffixAddingIdentifierConverter extends SuffixAddingIdentifierConverter {
-  private static final String SUFFIX = "." + VALUE_MAP_KEY + ".values.%s.value.";
+  private static final String SUFFIX =
+      "." + VALUE_MAP_KEY + "." + VALUES_KEY + ".%s." + VALUE_KEY + ".";
 
   @Inject
   public MapSuffixAddingIdentifierConverter(final ValueHelper valueHelper) {
     super(valueHelper);
   }
 
   @Override
-  protected String getSuffix(final Operator operator) {
-    return SUFFIX;
+  protected String getSuffix(final IdentifierConversionMetadata metadata)
+      throws ConversionException {
+    return SUFFIX + getTypeName(metadata);
   }
 }