Allow Postgres group-by on nested function expressions

document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java

@@ -308,6 +308,17 @@ public Stream<Arguments> provideArguments(final ExtensionContext context) {
     }
   }
 
+  private static class AllWithFlatProvider implements ArgumentsProvider {
+
+    @Override
+    public Stream<Arguments> 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<Map<String, Object>> actual = new ArrayList<>();
+    try (CloseableIterator<Document> resultDocs = collection.aggregate(query)) {
+      while (resultDocs.hasNext()) {
+        actual.add(normalizeNumbers(Utils.convertDocumentToMap(resultDocs.next())));
+      }
+    }
+
+    List<Map<String, Object>> 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<String, Object> normalizeNumbers(Map<String, Object> doc) {
+    Map<String, Object> 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 {

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
+  }
+]

document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java

@@ -6,17 +6,20 @@
 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.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);
@@ -30,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
@@ -84,10 +88,43 @@ private Collector getCollectorForFunctionOperator(FunctionOperator operator) {
   }
 
   private String getParsedExpression(final SelectTypeExpression expression) {
-    Optional<String> identifier =
-        Optional.ofNullable(expression.accept(identifierExpressionVisitor));
-    return identifier
-        .map(v -> getPostgresQueryParser().getPgSelections().get(v))
-        .orElse(expression.accept(selectTypeExpressionVisitor));
+    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;
+    }
+
+    /** 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));
+      }
+      return value.toString();
+    }
+
+    @Override
+    public String visit(final IdentifierExpression expression) {
+      Optional<String> selection =
+          Optional.ofNullable(identifierExpressionVisitor.visit(expression))
+              .map(identifier -> getPostgresQueryParser().getPgSelections().get(identifier));
+      return selection.orElseGet(() -> super.visit(expression));
+    }
   }
 }

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

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();
+  }
 }