diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/expression/impl/FunctionExpression.java b/document-store/src/main/java/org/hypertrace/core/documentstore/expression/impl/FunctionExpression.java index 1200920e9..7382a73ef 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/expression/impl/FunctionExpression.java +++ b/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); } } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoGroupTypeExpressionParser.java b/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoGroupTypeExpressionParser.java index 520520ce2..ff5a049e6 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoGroupTypeExpressionParser.java +++ b/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 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") diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/parser/GroupByAliasGetter.java b/document-store/src/main/java/org/hypertrace/core/documentstore/parser/GroupByAliasGetter.java index defb7e86b..6c91241dc 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/parser/GroupByAliasGetter.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/parser/GroupByAliasGetter.java @@ -9,7 +9,11 @@ public class GroupByAliasGetter implements GroupTypeExpressionVisitor { @Override public Optional visit(FunctionExpression expression) { - return Optional.empty(); + String alias = expression.getAlias(); + if (alias == null || alias.isBlank()) { + return Optional.empty(); + } + return Optional.of(alias); } @Override diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoFunctionGroupByTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoFunctionGroupByTest.java new file mode 100644 index 000000000..10f43f9f5 --- /dev/null +++ b/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 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")); + } +}