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,19 @@ 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.
+   */
+  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")

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,68 @@
+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.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;
+
+class MongoFunctionGroupByTest {
+
+  @Test
+  void groupsByAliasedArithmeticFunction() {
+    IdentifierExpression timestamp =
+        IdentifierExpression.of("attributes.last_activity_timestamp.value.long");
+    ConstantExpression interval = ConstantExpression.of(86_400_000L);
+    FunctionExpression bucket =
+        FunctionExpression.builder()
+            .alias("INTERVAL_START_TIME")
+            .operator(MULTIPLY)
+            .operand(
+                FunctionExpression.builder()
+                    .operator(FLOOR)
+                    .operand(
+                        FunctionExpression.builder()
+                            .operator(DIVIDE)
+                            .operand(timestamp)
+                            .operand(interval)
+                            .build())
+                    .build())
+            .operand(interval)
+            .build();
+
+    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"));
+  }
+}