Allow Mongo group-by on an aliased function expression

document-store/src/main/java/org/hypertrace/core/documentstore/expression/impl/FunctionExpression.java

@@ -48,13 +48,21 @@ public class FunctionExpression
 
   FunctionOperator operator;
 
+  /**
+   * Name of this expression when it is used as a group key. Mongo materializes the function in
+   * {@code $addFields} under this alias and groups by that field. A selection may use a different
+   * alias; that name is the projected column, read from {@code $_id.<this alias>}. A null or blank
+   * alias is not a group key.
+   */
+  String alias;
+
   public static class FunctionExpressionBuilder {
     public FunctionExpression build() {
       Preconditions.checkArgument(!operands.isEmpty(), "operands is empty");
       Preconditions.checkArgument(
           operands.stream().noneMatch(Objects::isNull), "One or more operands is null");
       Preconditions.checkArgument(operator != null, "operator is null");
-      return new FunctionExpression(operands, operator);
+      return new FunctionExpression(operands, operator, alias);
     }
   }
 

document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoGroupTypeExpressionParser.java

@@ -35,11 +35,15 @@ public final class MongoGroupTypeExpressionParser implements GroupTypeExpression
   @SuppressWarnings("unchecked")
   @Override
   public Map<String, Object> visit(final FunctionExpression expression) {
-    // To support this, we need to take an alias for GroupingExpressions
-    throw new UnsupportedOperationException(
-        String.format(
-            "Grouping by a function ($%s) is not yet supported by this library for MongoDB",
-            expression));
+    String alias = expression.getAlias();
+    if (alias == null || alias.isBlank()) {
+      throw new UnsupportedOperationException(
+          String.format(
+              "Grouping by a function ($%s) is not yet supported by this library for MongoDB",
+              expression));
+    }
+    // $addFields already computed this alias. Group by that field.
+    return Map.of(encodeKey(alias), PREFIX + alias);
   }
 
   @SuppressWarnings("unchecked")
@@ -56,16 +60,22 @@ public static List<BasicDBObject> getGroupClauses(final Query query) {
 
     final List<BasicDBObject> basicDBObjects = new ArrayList<>();
 
+    final List<String> groupByAliases = getGroupByAliases(expressions);
     final List<SelectionSpec> functionExpressionSelectionWithGroupBys =
-        getFunctionExpressionSelectionWithGroupBys(selectionSpecs, expressions);
+        getFunctionExpressionSelectionWithGroupBys(selectionSpecs, groupByAliases);
 
     if (!functionExpressionSelectionWithGroupBys.isEmpty()) {
       MongoSelectTypeExpressionParser parser =
           new MongoIdentifierPrefixingParser(
               new MongoIdentifierExpressionParser(new MongoFunctionExpressionParser()));
       Map<String, Object> addFields =
           functionExpressionSelectionWithGroupBys.stream()
-              .map(spec -> MongoGroupTypeExpressionParser.parse(parser, spec))
+              .map(
+                  spec ->
+                      MongoGroupTypeExpressionParser.parse(
+                          parser,
+                          SelectionSpec.of(
+                              spec.getExpression(), groupedFunctionField(spec, groupByAliases))))
               .reduce(
                   new LinkedHashMap<>(),
                   (first, second) -> {
@@ -134,21 +144,37 @@ private Map<String, Object> parse(final GroupTypeExpression expression) {
   }
 
   private static List<SelectionSpec> getFunctionExpressionSelectionWithGroupBys(
-      final List<SelectionSpec> selectionSpecs, final List<GroupTypeExpression> expressions) {
-    List<String> groupByAliases = getGroupByAliases(expressions);
-
+      final List<SelectionSpec> selectionSpecs, final List<String> groupByAliases) {
     return selectionSpecs.stream()
         .filter(
             selectionSpec ->
                 isFunctionExpressionSelectionWithGroupBy(selectionSpec, groupByAliases))
         .collect(Collectors.toUnmodifiableList());
   }
 
+  /**
+   * Field written by {@code $addFields} and stored under {@code $group._id}. This is the function
+   * alias when that alias is itself a group key, so a selection can expose the same value under a
+   * different name.
+   */
+  static String groupedFunctionField(
+      final SelectionSpec selectionSpec, final List<String> groupByAliases) {
+    if (selectionSpec.getExpression() instanceof FunctionExpression) {
+      String alias = ((FunctionExpression) selectionSpec.getExpression()).getAlias();
+      if (alias != null && !alias.isBlank() && groupByAliases.contains(alias)) {
+        return alias;
+      }
+    }
+    return selectionSpec.getAlias();
+  }
+
   public static boolean isFunctionExpressionSelectionWithGroupBy(
       final SelectionSpec selectionSpec, final List<String> groupByAliases) {
-    return selectionSpec.getAlias() != null
-        && groupByAliases.contains(selectionSpec.getAlias())
-        && (Boolean) selectionSpec.getExpression().accept(FUNCTION_EXPRESSION_CHECKER);
+    if (!(Boolean) selectionSpec.getExpression().accept(FUNCTION_EXPRESSION_CHECKER)) {
+      return false;
+    }
+    String field = groupedFunctionField(selectionSpec, groupByAliases);
+    return field != null && groupByAliases.contains(field);
   }
 
   @SuppressWarnings("unchecked")

document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoSelectTypeExpressionParser.java

@@ -3,6 +3,7 @@
 import static java.util.stream.Collectors.toMap;
 import static org.hypertrace.core.documentstore.mongo.MongoCollection.ID_KEY;
 import static org.hypertrace.core.documentstore.mongo.query.parser.MongoGroupTypeExpressionParser.getGroupByAliases;
+import static org.hypertrace.core.documentstore.mongo.query.parser.MongoGroupTypeExpressionParser.groupedFunctionField;
 import static org.hypertrace.core.documentstore.mongo.query.parser.MongoGroupTypeExpressionParser.isFunctionExpressionSelectionWithGroupBy;
 
 import com.google.common.base.Joiner;
@@ -79,7 +80,9 @@ public static BasicDBObject getSelections(final Query query) {
                 spec ->
                     isFunctionExpressionSelectionWithGroupBy(spec, groupByAliases)
                         ? SelectionSpec.of(
-                            IdentifierExpression.of(DOT_JOINER.join(ID_KEY, spec.getAlias())),
+                            IdentifierExpression.of(
+                                DOT_JOINER.join(
+                                    ID_KEY, groupedFunctionField(spec, groupByAliases))),
                             spec.getAlias())
                         : spec)
             .map(spec -> MongoSelectTypeExpressionParser.parse(parser, spec))

document-store/src/main/java/org/hypertrace/core/documentstore/parser/GroupByAliasGetter.java

@@ -9,7 +9,11 @@ public class GroupByAliasGetter implements GroupTypeExpressionVisitor {
 
   @Override
   public Optional<String> visit(FunctionExpression expression) {
-    return Optional.empty();
+    String alias = expression.getAlias();
+    if (alias == null || alias.isBlank()) {
+      return Optional.empty();
+    }
+    return Optional.of(alias);
   }
 
   @Override

document-store/src/test/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoFunctionGroupByTest.java

@@ -0,0 +1,122 @@
+package org.hypertrace.core.documentstore.mongo.query.parser;
+
+import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.COUNT;
+import static org.hypertrace.core.documentstore.expression.operators.FunctionOperator.DIVIDE;
+import static org.hypertrace.core.documentstore.expression.operators.FunctionOperator.FLOOR;
+import static org.hypertrace.core.documentstore.expression.operators.FunctionOperator.MULTIPLY;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import com.mongodb.BasicDBObject;
+import java.util.List;
+import java.util.Map;
+import org.hypertrace.core.documentstore.expression.impl.AggregateExpression;
+import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
+import org.hypertrace.core.documentstore.expression.impl.FunctionExpression;
+import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression;
+import org.hypertrace.core.documentstore.query.Query;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.NullSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+class MongoFunctionGroupByTest {
+
+  @Test
+  void groupsByAliasedArithmeticFunction() {
+    FunctionExpression bucket = bucket("INTERVAL_START_TIME");
+
+    Query query =
+        Query.builder()
+            .addSelection(bucket, "INTERVAL_START_TIME")
+            .addSelection(AggregateExpression.of(COUNT, IdentifierExpression.of("id")), "count")
+            .addAggregation(bucket)
+            .addAggregation(IdentifierExpression.of("attributes.score_category"))
+            .build();
+
+    List<BasicDBObject> clauses = MongoGroupTypeExpressionParser.getGroupClauses(query);
+    assertEquals(2, clauses.size());
+    assertTrue(clauses.get(0).containsKey("$addFields"));
+    assertTrue(clauses.get(1).containsKey("$group"));
+
+    Map<?, ?> addFields = (Map<?, ?>) clauses.get(0).get("$addFields");
+    assertTrue(addFields.containsKey("INTERVAL_START_TIME"));
+
+    Map<?, ?> group = (Map<?, ?>) clauses.get(1).get("$group");
+    Map<?, ?> id = (Map<?, ?>) group.get("_id");
+    assertEquals("$INTERVAL_START_TIME", id.get("INTERVAL_START_TIME"));
+    assertEquals("$attributes.score_category", id.get("attributes\\u002escore_category"));
+
+    BasicDBObject projection = MongoSelectTypeExpressionParser.getSelections(query);
+    assertEquals("$_id.INTERVAL_START_TIME", projection.get("INTERVAL_START_TIME"));
+  }
+
+  @ParameterizedTest
+  @NullSource
+  @ValueSource(strings = {"", " "})
+  void rejectsFunctionGroupByWhenAliasIsMissing(final String alias) {
+    FunctionExpression bucket = bucket(alias);
+    Query query =
+        Query.builder()
+            .addSelection(bucket, "foo")
+            .addSelection(AggregateExpression.of(COUNT, IdentifierExpression.of("id")), "count")
+            .addAggregation(bucket)
+            .build();
+
+    UnsupportedOperationException exception =
+        assertThrows(
+            UnsupportedOperationException.class,
+            () -> MongoGroupTypeExpressionParser.getGroupClauses(query));
+    assertTrue(exception.getMessage().contains("not yet supported"));
+  }
+
+  @Test
+  void projectsSelectionAliasWhenItDiffersFromFunctionAlias() {
+    FunctionExpression bucket = bucket("bar");
+    Query query =
+        Query.builder()
+            .addSelection(bucket, "foo")
+            .addSelection(AggregateExpression.of(COUNT, IdentifierExpression.of("id")), "count")
+            .addAggregation(bucket)
+            .build();
+
+    List<BasicDBObject> clauses = MongoGroupTypeExpressionParser.getGroupClauses(query);
+    assertEquals(2, clauses.size());
+
+    Map<?, ?> addFields = (Map<?, ?>) clauses.get(0).get("$addFields");
+    assertTrue(addFields.containsKey("bar"));
+    assertFalse(addFields.containsKey("foo"));
+
+    Map<?, ?> group = (Map<?, ?>) clauses.get(1).get("$group");
+    Map<?, ?> id = (Map<?, ?>) group.get("_id");
+    assertEquals("$bar", id.get("bar"));
+    assertFalse(id.containsKey("foo"));
+
+    BasicDBObject projection = MongoSelectTypeExpressionParser.getSelections(query);
+    assertEquals("$_id.bar", projection.get("foo"));
+    assertFalse(projection.containsKey("bar"));
+  }
+
+  private static FunctionExpression bucket(final String alias) {
+    IdentifierExpression timestamp =
+        IdentifierExpression.of("attributes.last_activity_timestamp.value.long");
+    ConstantExpression interval = ConstantExpression.of(86_400_000L);
+    return FunctionExpression.builder()
+        .alias(alias)
+        .operator(MULTIPLY)
+        .operand(
+            FunctionExpression.builder()
+                .operator(FLOOR)
+                .operand(
+                    FunctionExpression.builder()
+                        .operator(DIVIDE)
+                        .operand(timestamp)
+                        .operand(interval)
+                        .build())
+                .build())
+        .operand(interval)
+        .build();
+  }
+}