fix: handle group_by identifier more generically

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

@@ -743,6 +743,30 @@ void testQueryQ1DistinctCountAggregationWithOnlyFilter(String dataStoreName) thr
     }
   }
 
+  @ParameterizedTest
+  @MethodSource("databaseContextBoth")
+  void testQueryQ1DistinctCountAggregationWithMatchingSelectionAndGroupBy(String dataStoreName)
+      throws IOException {
+    Datastore datastore = datastoreMap.get(dataStoreName);
+    Collection collection = datastore.getCollection(COLLECTION_NAME);
+    org.hypertrace.core.documentstore.query.Query query =
+        org.hypertrace.core.documentstore.query.Query.builder()
+            .addSelection(
+                AggregateExpression.of(DISTINCT_COUNT, IdentifierExpression.of("quantity")),
+                "qty_count")
+            .addSelection(IdentifierExpression.of("item"))
+            .addSelection(IdentifierExpression.of("price"))
+            .setFilter(
+                RelationalExpression.of(
+                    IdentifierExpression.of("price"), LTE, ConstantExpression.of(10)))
+            .addAggregation(IdentifierExpression.of("item"))
+            .build();
+    try (CloseableIterator<Document> resultDocs = collection.aggregate(query)) {
+      assertDocsAndSizeEqualWithoutOrder(
+          dataStoreName, resultDocs, 3, "mongo/test_aggr_with_match_selection_and_groupby.json");
+    }
+  }
+
   @ParameterizedTest
   @MethodSource("databaseContextBoth")
   public void testQueryV1ForSimpleWhereClause(String dataStoreName) throws IOException {

document-store/src/integrationTest/resources/mongo/test_aggr_with_match_selection_and_groupby.json

@@ -0,0 +1,14 @@
+[
+  {
+    "item":"Soap",
+    "qty_count":2
+  },
+  {
+    "item":"Comb",
+    "qty_count":2
+  },
+  {
+    "item":"Shampoo",
+    "qty_count":2
+  }
+]

document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresSelectionQueryTransformer.java

@@ -1,11 +1,16 @@
 package org.hypertrace.core.documentstore.postgres.query.v1.transformer;
 
 import java.util.List;
+import java.util.Set;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import org.hypertrace.core.documentstore.expression.impl.AggregateExpression;
 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.type.GroupTypeExpression;
+import org.hypertrace.core.documentstore.expression.type.SelectTypeExpression;
+import org.hypertrace.core.documentstore.parser.GroupTypeExpressionVisitor;
 import org.hypertrace.core.documentstore.parser.SelectTypeExpressionVisitor;
 import org.hypertrace.core.documentstore.query.Query;
 import org.hypertrace.core.documentstore.query.SelectionSpec;
@@ -29,42 +34,110 @@
  *
  * This is the similar behavior supported in our other document store implementation (e.g Mongo)
  * */
-public class PostgresSelectionQueryTransformer
-    implements QueryTransformer, SelectTypeExpressionVisitor {
+public class PostgresSelectionQueryTransformer implements QueryTransformer {
 
   @Override
   public Query transform(Query query) {
-    // no-op if group by clause exits
-    if (!query.getAggregations().isEmpty()) return query;
+    LocalSelectTypeAggregationExpressionSelector aggregationExpressionSelector =
+        new LocalSelectTypeAggregationExpressionSelector();
+    Boolean noAggregation =
+        query.getSelections().stream()
+            .filter(
+                selectionSpec ->
+                    selectionSpec.getExpression().accept(aggregationExpressionSelector))
+            .findAny()
+            .isEmpty();
+    if (noAggregation) return query;
+
+    // get all identifier of group by clause
+    LocalGroupByIdentifierExpressionSelector groupByIdentifierExpressionSelector =
+        new LocalGroupByIdentifierExpressionSelector();
+    Set<GroupTypeExpression> groupByIdentifierExpressions =
+        query.getAggregations().stream()
+            .filter(exp -> exp.accept(groupByIdentifierExpressionSelector))
+            .collect(Collectors.toUnmodifiableSet());
 
-    // check for all selections, remove non-aggregated selections.
+    // keep only matching identifier expression with group by along with rest.
+    LocalSelectTypeIdentifierExpressionSelector identifierExpressionSelector =
+        new LocalSelectTypeIdentifierExpressionSelector();
     List<SelectionSpec> finalSelectionSpecs =
         query.getSelections().stream()
-            .filter(selectionSpec -> selectionSpec.getExpression().accept(this))
+            .filter(
+                getSelectionSpecPredicate(
+                    groupByIdentifierExpressions, identifierExpressionSelector))
             .collect(Collectors.toUnmodifiableList());
 
-    return finalSelectionSpecs.size() > 0
+    return finalSelectionSpecs.size() != query.getSelections().size()
         ? new TransformedQueryBuilder(query).setSelections(finalSelectionSpecs).build()
         : query;
   }
 
-  @Override
-  public Boolean visit(AggregateExpression expression) {
-    return true;
+  private Predicate<SelectionSpec> getSelectionSpecPredicate(
+      Set<GroupTypeExpression> groupByIdentifierExpressions,
+      LocalSelectTypeIdentifierExpressionSelector identifierExpressionSelector) {
+    return selectionSpec -> {
+      SelectTypeExpression expression = selectionSpec.getExpression();
+      return (!(boolean) expression.accept(identifierExpressionSelector))
+          || groupByIdentifierExpressions.contains(expression);
+    };
   }
 
-  @Override
-  public Boolean visit(ConstantExpression expression) {
-    return false;
+  private static class LocalSelectTypeAggregationExpressionSelector
+      implements SelectTypeExpressionVisitor {
+    @Override
+    public Boolean visit(AggregateExpression expression) {
+      return true;
+    }
+
+    @Override
+    public Boolean visit(ConstantExpression expression) {
+      return false;
+    }
+
+    @Override
+    public Boolean visit(FunctionExpression expression) {
+      return false;
+    }
+
+    @Override
+    public Boolean visit(IdentifierExpression expression) {
+      return false;
+    }
   }
 
-  @Override
-  public Boolean visit(FunctionExpression expression) {
-    return false;
+  private static class LocalSelectTypeIdentifierExpressionSelector
+      implements SelectTypeExpressionVisitor {
+    @Override
+    public Boolean visit(AggregateExpression expression) {
+      return false;
+    }
+
+    @Override
+    public Boolean visit(ConstantExpression expression) {
+      return false;
+    }
+
+    @Override
+    public Boolean visit(FunctionExpression expression) {
+      return false;
+    }
+
+    @Override
+    public Boolean visit(IdentifierExpression expression) {
+      return true;
+    }
   }
 
-  @Override
-  public Boolean visit(IdentifierExpression expression) {
-    return false;
+  private static class LocalGroupByIdentifierExpressionSelector
+      implements GroupTypeExpressionVisitor {
+    @Override
+    public Boolean visit(FunctionExpression expression) {
+      return false;
+    }
+
+    @Override
+    public Boolean visit(IdentifierExpression expression) {
+      return true;
+    }
   }
 }

document-store/src/test/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParserTest.java

@@ -917,4 +917,35 @@ void testQueryWithFunctionalLhsInRelationalFilter() {
     assertEquals(1, params.getObjectParams().size());
     assertEquals(50, params.getObjectParams().get(1));
   }
+
+  @Test
+  void testQueryQ1DistinctCountAggregationWithMatchingSelectionAndGroupBy() {
+    org.hypertrace.core.documentstore.query.Query query =
+        org.hypertrace.core.documentstore.query.Query.builder()
+            .addSelection(
+                AggregateExpression.of(DISTINCT_COUNT, IdentifierExpression.of("quantity")),
+                "qty_count")
+            .addSelection(IdentifierExpression.of("item"))
+            .addSelection(IdentifierExpression.of("price"))
+            .setFilter(
+                RelationalExpression.of(
+                    IdentifierExpression.of("price"), LTE, ConstantExpression.of(10)))
+            .addAggregation(IdentifierExpression.of("item"))
+            .build();
+
+    PostgresQueryParser postgresQueryParser =
+        new PostgresQueryParser(TEST_COLLECTION, PostgresQueryTransformer.transform(query));
+    String sql = postgresQueryParser.parse();
+
+    assertEquals(
+        "SELECT COUNT(DISTINCT document->>'quantity' ) AS \"qty_count\", "
+            + "document->'item' AS item "
+            + "FROM testCollection WHERE CAST (document->>'price' AS NUMERIC) <= ? "
+            + "GROUP BY document->'item'",
+        sql);
+
+    Params params = postgresQueryParser.getParamsBuilder().build();
+    assertEquals(1, params.getObjectParams().size());
+    assertEquals(10, params.getObjectParams().get(1));
+  }
 }