From 8e3e8958d7cc4f4b1a66e120c4d731e656d57f4c Mon Sep 17 00:00:00 2001 From: bhubam Date: Mon, 19 Oct 2026 14:09:33 +0530 Subject: [PATCH 1/3] Allow Postgres group-by on nested function expressions Interval buckets such as MULTIPLY(FLOOR(DIVIDE(col, 60000)), 60000) failed on Postgres because function operands could only be columns. Nested functions and numeric constants are now rendered inline, and a function group-by repeats the selected expression. Co-authored-by: Cursor --- .../PostgresFunctionExpressionVisitor.java | 16 ++++ .../PostgresGroupTypeExpressionVisitor.java | 9 +- .../query/v1/PostgresQueryParserTest.java | 90 +++++++++++++++++++ 3 files changed, 113 insertions(+), 2 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java index f1b0f2ee4..6133f8f7e 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java @@ -6,6 +6,7 @@ import java.util.stream.Collectors; import lombok.NoArgsConstructor; import org.apache.commons.lang3.StringUtils; +import org.hypertrace.core.documentstore.expression.impl.ConstantExpression; import org.hypertrace.core.documentstore.expression.impl.FunctionExpression; import org.hypertrace.core.documentstore.expression.operators.FunctionOperator; import org.hypertrace.core.documentstore.expression.type.SelectTypeExpression; @@ -84,6 +85,21 @@ private Collector getCollectorForFunctionOperator(FunctionOperator operator) { } private String getParsedExpression(final SelectTypeExpression expression) { + if (expression instanceof FunctionExpression) { + String nested = visit((FunctionExpression) expression); + return nested != null ? String.format("(%s)", nested) : null; + } + + if (expression instanceof ConstantExpression) { + // Inlined rather than bound so a GROUP BY repeats the selected expression verbatim. + Object value = ((ConstantExpression) expression).getValue(); + if (!(value instanceof Number)) { + throw new UnsupportedOperationException( + String.format("Only numeric constants are supported in functions: %s", expression)); + } + return value.toString(); + } + Optional identifier = Optional.ofNullable(expression.accept(identifierExpressionVisitor)); return identifier diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresGroupTypeExpressionVisitor.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresGroupTypeExpressionVisitor.java index c994b0752..f80082895 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresGroupTypeExpressionVisitor.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresGroupTypeExpressionVisitor.java @@ -14,15 +14,20 @@ public class PostgresGroupTypeExpressionVisitor implements GroupTypeExpressionVisitor { private PostgresSelectTypeExpressionVisitor selectTypeExpressionVisitor; + private PostgresFunctionExpressionVisitor functionExpressionVisitor; public PostgresGroupTypeExpressionVisitor(PostgresQueryParser postgresQueryParser) { selectTypeExpressionVisitor = new PostgresFieldIdentifierExpressionVisitor(postgresQueryParser); + functionExpressionVisitor = new PostgresFunctionExpressionVisitor(postgresQueryParser); } + /** + * Renders the function the same way as its selection, so Postgres matches the GROUP BY expression + * against the selected one. + */ @Override public String visit(final FunctionExpression expression) { - throw new UnsupportedOperationException( - "FunctionalExpression on Group by is not yet supported"); + return functionExpressionVisitor.visit(expression); } @Override diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParserTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParserTest.java index 090e10b61..46ca7eb94 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParserTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParserTest.java @@ -2505,4 +2505,94 @@ void testArrayOperatorsRejectEmptyValueList() { // can never reach the ALL/EXACTLY_ONE parsers assertThrows(IllegalArgumentException.class, () -> ConstantExpression.ofStrings(List.of())); } + + @Test + void testNestedFunctionSelection() { + Query query = + Query.builder() + .addSelection( + FunctionExpression.builder() + .operator(MULTIPLY) + .operand( + FunctionExpression.builder() + .operator(FunctionOperator.ADD) + .operand(IdentifierExpression.of("price")) + .operand(ConstantExpression.of(5)) + .build()) + .operand(IdentifierExpression.of("quantity")) + .build(), + "total") + .build(); + + PostgresQueryParser postgresQueryParser = + new PostgresQueryParser(TEST_TABLE, PostgresQueryTransformer.transform(query)); + + assertEquals( + "SELECT (CAST (document->>'price' AS NUMERIC) + 5) * CAST (document->>'quantity' AS NUMERIC) " + + "AS \"total\" FROM \"testCollection\"", + postgresQueryParser.parse()); + } + + @Test + void testGroupByIntervalFunction() { + Query query = intervalQuery(); + + PostgresQueryParser postgresQueryParser = + new PostgresQueryParser(TEST_TABLE, PostgresQueryTransformer.transform(query)); + + assertEquals( + "SELECT document->'category' AS \"category\", " + + "(FLOOR( (CAST (document->>'lastActivity' AS NUMERIC) / 60000) )) * 60000 AS \"bucket\", " + + "COUNT( document->>'category' ) AS \"count\" " + + "FROM \"testCollection\" " + + "GROUP BY document->'category'," + + "(FLOOR( (CAST (document->>'lastActivity' AS NUMERIC) / 60000) )) * 60000", + postgresQueryParser.parse()); + } + + @Test + void testGroupByIntervalFunctionOnFlatCollection() { + Query query = intervalQuery(); + + PostgresQueryParser postgresQueryParser = + new PostgresQueryParser( + TEST_TABLE, + PostgresQueryTransformer.transform(query), + new FlatPostgresFieldTransformer()); + + assertEquals( + "SELECT \"category\" AS \"category\", " + + "(FLOOR( (\"lastActivity\" / 60000) )) * 60000 AS \"bucket\", " + + "COUNT( \"category\" ) AS \"count\" " + + "FROM \"testCollection\" " + + "GROUP BY \"category\",(FLOOR( (\"lastActivity\" / 60000) )) * 60000", + postgresQueryParser.parse()); + } + + private static Query intervalQuery() { + FunctionExpression bucket = + FunctionExpression.builder() + .operator(MULTIPLY) + .operand( + FunctionExpression.builder() + .operator(FunctionOperator.FLOOR) + .operand( + FunctionExpression.builder() + .operator(FunctionOperator.DIVIDE) + .operand(IdentifierExpression.of("lastActivity")) + .operand(ConstantExpression.of(60000)) + .build()) + .build()) + .operand(ConstantExpression.of(60000)) + .alias("bucket") + .build(); + + return Query.builder() + .addSelection(IdentifierExpression.of("category")) + .addSelection(bucket, "bucket") + .addSelection(AggregateExpression.of(COUNT, IdentifierExpression.of("category")), "count") + .addAggregation(IdentifierExpression.of("category")) + .addAggregation(bucket) + .build(); + } } From 48a20cef34de1d092c4b81c8fd22fcf82644ee48 Mon Sep 17 00:00:00 2001 From: bhubam Date: Wed, 21 Oct 2026 11:26:00 +0530 Subject: [PATCH 2/3] Add Mongo and Postgres parity test for interval function group-by Co-authored-by: Cursor --- .../documentstore/DocStoreQueryV1Test.java | 82 +++++++++++++++++++ .../query/interval_group_by_response.json | 32 ++++++++ 2 files changed, 114 insertions(+) create mode 100644 document-store/src/integrationTest/resources/query/interval_group_by_response.json diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java index 01e000537..1c815bcf5 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java @@ -308,6 +308,17 @@ public Stream provideArguments(final ExtensionContext context) { } } + private static class AllWithFlatProvider implements ArgumentsProvider { + + @Override + public Stream provideArguments(final ExtensionContext context) { + return Stream.of( + Arguments.of(MONGO_STORE, DocumentType.NESTED), + Arguments.of(POSTGRES_STORE, DocumentType.NESTED), + Arguments.of(POSTGRES_STORE, DocumentType.FLAT)); + } + } + /** * Provides arguments for testing array operations with different expression types. Returns: * (datastoreName, expressionType) - "WITH_TYPE": ArrayIdentifierExpression WITH ArrayType @@ -6883,6 +6894,77 @@ public void testMongoFunctionExpressionGroupBy(String dataStoreName) throws Exce testCountApi(dataStoreName, query, "query/function_expression_group_by_response.json"); } + /** + * Buckets quantity into intervals of 10 and groups by item and bucket. The flat collection holds + * two extra items, which are filtered out so every store sees the same rows. + */ + @ParameterizedTest + @ArgumentsSource(AllWithFlatProvider.class) + public void testIntervalFunctionGroupBy(String dataStoreName, DocumentType documentType) + throws IOException { + Collection collection = + documentType == DocumentType.FLAT + ? getFlatCollection(dataStoreName) + : getCollection(dataStoreName); + + FunctionExpression bucket = + FunctionExpression.builder() + .operator(MULTIPLY) + .operand( + FunctionExpression.builder() + .operator(FLOOR) + .operand( + FunctionExpression.builder() + .operator(DIVIDE) + .operand(IdentifierExpression.of("quantity")) + .operand(ConstantExpression.of(10)) + .build()) + .build()) + .operand(ConstantExpression.of(10)) + .alias("bucket") + .build(); + + Query query = + Query.builder() + .addSelection(IdentifierExpression.of("item")) + .addSelection(bucket, "bucket") + .addSelection(AggregateExpression.of(COUNT, IdentifierExpression.of("item")), "count") + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("item"), + NOT_IN, + ConstantExpression.ofStrings(List.of("Bottle", "Cup")))) + .addAggregation(IdentifierExpression.of("item")) + .addAggregation(bucket) + .build(); + + List> actual = new ArrayList<>(); + try (CloseableIterator resultDocs = collection.aggregate(query)) { + while (resultDocs.hasNext()) { + actual.add(normalizeNumbers(Utils.convertDocumentToMap(resultDocs.next()))); + } + } + + List> expected = + Utils.convertJsonToMap( + readFileFromResource("query/interval_group_by_response.json").orElseThrow()) + .stream() + .map(DocStoreQueryV1Test::normalizeNumbers) + .collect(toUnmodifiableList()); + + assertEquals(expected.size(), actual.size(), "Unexpected rows: " + actual); + assertTrue(actual.containsAll(expected), "Unexpected rows: " + actual); + } + + /** Stores return integral results as either integers or doubles, so compare them as doubles. */ + private static Map normalizeNumbers(Map doc) { + Map normalized = new HashMap<>(); + doc.forEach( + (key, value) -> + normalized.put(key, value instanceof Number ? ((Number) value).doubleValue() : value)); + return normalized; + } + @ParameterizedTest @ArgumentsSource(MongoProvider.class) public void testToLowerCaseMongoFunctionOperator(String dataStoreName) throws Exception { diff --git a/document-store/src/integrationTest/resources/query/interval_group_by_response.json b/document-store/src/integrationTest/resources/query/interval_group_by_response.json new file mode 100644 index 000000000..3177513fe --- /dev/null +++ b/document-store/src/integrationTest/resources/query/interval_group_by_response.json @@ -0,0 +1,32 @@ +[ + { + "item": "Soap", + "bucket": 0, + "count": 3 + }, + { + "item": "Mirror", + "bucket": 0, + "count": 1 + }, + { + "item": "Shampoo", + "bucket": 10, + "count": 1 + }, + { + "item": "Shampoo", + "bucket": 20, + "count": 1 + }, + { + "item": "Comb", + "bucket": 0, + "count": 1 + }, + { + "item": "Comb", + "bucket": 10, + "count": 1 + } +] From 96edd60ad2d2172b42a061da015a749d0badd363 Mon Sep 17 00:00:00 2001 From: bhubam Date: Thu, 22 Oct 2026 11:07:59 +0530 Subject: [PATCH 3/3] Use a visitor for Postgres function operands instead of instanceof checks Co-authored-by: Cursor --- .../PostgresFunctionExpressionVisitor.java | 49 +++++++++++++------ 1 file changed, 35 insertions(+), 14 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java index 6133f8f7e..5113a4ca1 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java @@ -8,16 +8,18 @@ import org.apache.commons.lang3.StringUtils; 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.expression.operators.FunctionOperator; import org.hypertrace.core.documentstore.expression.type.SelectTypeExpression; import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser; +import org.hypertrace.core.documentstore.postgres.utils.PostgresUtils.Type; @NoArgsConstructor public class PostgresFunctionExpressionVisitor extends PostgresSelectTypeExpressionVisitor { private static final int ARRAY_DIMENSION = 1; private PostgresIdentifierExpressionVisitor identifierExpressionVisitor; - private PostgresSelectTypeExpressionVisitor selectTypeExpressionVisitor; + private PostgresFunctionOperandVisitor operandVisitor; public PostgresFunctionExpressionVisitor(PostgresSelectTypeExpressionVisitor baseVisitor) { super(baseVisitor); @@ -31,9 +33,10 @@ public PostgresFunctionExpressionVisitor(PostgresQueryParser postgresQueryParser private void initExpressionVisitor() { identifierExpressionVisitor = new PostgresIdentifierExpressionVisitor(getPostgresQueryParser()); - selectTypeExpressionVisitor = - new PostgresDataAccessorIdentifierExpressionVisitor( - new PostgresConstantExpressionVisitor(getPostgresQueryParser())); + operandVisitor = + new PostgresFunctionOperandVisitor( + new PostgresDataAccessorIdentifierExpressionVisitor( + getPostgresQueryParser(), Type.NUMERIC)); } @Override @@ -85,14 +88,30 @@ private Collector getCollectorForFunctionOperator(FunctionOperator operator) { } private String getParsedExpression(final SelectTypeExpression expression) { - if (expression instanceof FunctionExpression) { - String nested = visit((FunctionExpression) expression); + return expression.accept(operandVisitor); + } + + private class PostgresFunctionOperandVisitor extends PostgresSelectTypeExpressionVisitor { + + private PostgresFunctionOperandVisitor(final PostgresSelectTypeExpressionVisitor baseVisitor) { + super(baseVisitor); + } + + @Override + public PostgresQueryParser getPostgresQueryParser() { + return PostgresFunctionExpressionVisitor.this.getPostgresQueryParser(); + } + + @Override + public String visit(final FunctionExpression expression) { + String nested = PostgresFunctionExpressionVisitor.this.visit(expression); return nested != null ? String.format("(%s)", nested) : null; } - if (expression instanceof ConstantExpression) { - // Inlined rather than bound so a GROUP BY repeats the selected expression verbatim. - Object value = ((ConstantExpression) expression).getValue(); + /** Inlined rather than bound so a GROUP BY repeats the selected expression verbatim. */ + @Override + public String visit(final ConstantExpression expression) { + Object value = expression.getValue(); if (!(value instanceof Number)) { throw new UnsupportedOperationException( String.format("Only numeric constants are supported in functions: %s", expression)); @@ -100,10 +119,12 @@ private String getParsedExpression(final SelectTypeExpression expression) { return value.toString(); } - Optional identifier = - Optional.ofNullable(expression.accept(identifierExpressionVisitor)); - return identifier - .map(v -> getPostgresQueryParser().getPgSelections().get(v)) - .orElse(expression.accept(selectTypeExpressionVisitor)); + @Override + public String visit(final IdentifierExpression expression) { + Optional selection = + Optional.ofNullable(identifierExpressionVisitor.visit(expression)) + .map(identifier -> getPostgresQueryParser().getPgSelections().get(identifier)); + return selection.orElseGet(() -> super.visit(expression)); + } } }