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 f1b0f2ee..6133f8f7 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 c994b075..f8008289 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 090e10b6..46ca7eb9 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(); + } }