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 01e00053..1c815bcf 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 00000000..3177513f --- /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 + } +] 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(); + } }