From 0722ddd95808baa0f5b214da338ec41d6cd019a4 Mon Sep 17 00:00:00 2001 From: bhubam Date: Wed, 23 Sep 2026 17:49:51 +0530 Subject: [PATCH 1/3] 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/3] 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")); From 3952cbdb0131592d8c1c4e4d6e2e388f50b2bb6d Mon Sep 17 00:00:00 2001 From: bhubam Date: Wed, 30 Sep 2026 11:51:00 +0530 Subject: [PATCH 3/3] Cover null and mismatched aliases for Mongo function group-by. A null or blank function alias still fails. When the selection alias differs, group by the function alias and project that _id field under the selection name. Co-authored-by: Cursor --- .../expression/impl/FunctionExpression.java | 4 +- .../MongoGroupTypeExpressionParser.java | 38 ++++++-- .../MongoSelectTypeExpressionParser.java | 5 +- .../parser/MongoFunctionGroupByTest.java | 92 +++++++++++++++---- 4 files changed, 110 insertions(+), 29 deletions(-) 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 7382a73ef..6fefcd825 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 @@ -50,7 +50,9 @@ public class FunctionExpression /** * 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. + * {@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.}. A null or blank + * alias is not a group key. */ String 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 ff5a049e6..5f3211a1a 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 @@ -60,8 +60,9 @@ public static List getGroupClauses(final Query query) { final List basicDBObjects = new ArrayList<>(); + final List groupByAliases = getGroupByAliases(expressions); final List functionExpressionSelectionWithGroupBys = - getFunctionExpressionSelectionWithGroupBys(selectionSpecs, expressions); + getFunctionExpressionSelectionWithGroupBys(selectionSpecs, groupByAliases); if (!functionExpressionSelectionWithGroupBys.isEmpty()) { MongoSelectTypeExpressionParser parser = @@ -69,7 +70,12 @@ public static List getGroupClauses(final Query query) { new MongoIdentifierExpressionParser(new MongoFunctionExpressionParser())); Map 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) -> { @@ -138,9 +144,7 @@ private Map parse(final GroupTypeExpression expression) { } private static List getFunctionExpressionSelectionWithGroupBys( - final List selectionSpecs, final List expressions) { - List groupByAliases = getGroupByAliases(expressions); - + final List selectionSpecs, final List groupByAliases) { return selectionSpecs.stream() .filter( selectionSpec -> @@ -148,11 +152,29 @@ private static List getFunctionExpressionSelectionWithGroupBys( .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 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 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") diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoSelectTypeExpressionParser.java b/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoSelectTypeExpressionParser.java index 2cf002623..76d785ccf 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoSelectTypeExpressionParser.java +++ b/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)) 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 10f43f9f5..e38cd8f63 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 @@ -5,6 +5,8 @@ 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; @@ -16,30 +18,15 @@ 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() { - 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(); + FunctionExpression bucket = bucket("INTERVAL_START_TIME"); Query query = Query.builder() @@ -65,4 +52,71 @@ void groupsByAliasedArithmeticFunction() { 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 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(); + } }