From 0722ddd95808baa0f5b214da338ec41d6cd019a4 Mon Sep 17 00:00:00 2001 From: bhubam Date: Wed, 23 Sep 2026 17:49:51 +0530 Subject: [PATCH 1/2] Allow Mongo group-by on an aliased function expression. Interval buckets are computed in $addFields and grouped by that field instead of the function itself. --- .../expression/impl/FunctionExpression.java | 8 ++- .../MongoGroupTypeExpressionParser.java | 14 ++-- .../parser/GroupByAliasGetter.java | 6 +- .../parser/MongoFunctionGroupByTest.java | 70 +++++++++++++++++++ 4 files changed, 91 insertions(+), 7 deletions(-) create mode 100644 document-store/src/test/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoFunctionGroupByTest.java 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..15f0a82c2 --- /dev/null +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoFunctionGroupByTest.java @@ -0,0 +1,70 @@ +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")); + } +} From 869676908be6f3a94c82ae9e1c9d6aa3c0128604 Mon Sep 17 00:00:00 2001 From: bhubam Date: Wed, 23 Sep 2026 21:11:42 +0530 Subject: [PATCH 2/2] fix --- .../mongo/query/parser/MongoFunctionGroupByTest.java | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) 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 index 15f0a82c2..10f43f9f5 100644 --- 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 @@ -60,9 +60,7 @@ void groupsByAliasedArithmeticFunction() { 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")); + assertEquals("$attributes.score_category", id.get("attributes\\u002escore_category")); BasicDBObject projection = MongoSelectTypeExpressionParser.getSelections(query); assertEquals("$_id.INTERVAL_START_TIME", projection.get("INTERVAL_START_TIME"));