Allow Postgres group-by on nested function expressions

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<String> identifier =
         Optional.ofNullable(expression.accept(identifierExpressionVisitor));
     return identifier

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