From 63fd747899f5a98c7abc550a8450e8261a49ab24 Mon Sep 17 00:00:00 2001 From: Ronak Date: Fri, 2 Sep 2022 16:34:53 +0530 Subject: [PATCH 1/4] fix: adding group by checkers --- .../documentstore/DocStoreQueryV1Test.java | 21 +++++ .../PostgresSelectionQueryTransformer.java | 89 +++++++++++++++---- .../query/v1/PostgresQueryParserTest.java | 29 ++++++ 3 files changed, 121 insertions(+), 18 deletions(-) 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 c39275f3..19c635c4 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 @@ -743,6 +743,27 @@ 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 resultDocs = collection.aggregate(query)) { + assertDocsAndSizeEqualWithoutOrder( + dataStoreName, resultDocs, 1, "mongo/test_aggr_only_with_fliter_response.json"); + } + } + @ParameterizedTest @MethodSource("databaseContextBoth") public void testQueryV1ForSimpleWhereClause(String dataStoreName) throws IOException { diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresSelectionQueryTransformer.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresSelectionQueryTransformer.java index 89a0589b..954395b5 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresSelectionQueryTransformer.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresSelectionQueryTransformer.java @@ -1,11 +1,14 @@ package org.hypertrace.core.documentstore.postgres.query.v1.transformer; +import com.google.common.base.Strings; import java.util.List; +import java.util.Set; 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.parser.GroupTypeExpressionVisitor; import org.hypertrace.core.documentstore.parser.SelectTypeExpressionVisitor; import org.hypertrace.core.documentstore.query.Query; import org.hypertrace.core.documentstore.query.SelectionSpec; @@ -29,18 +32,32 @@ * * 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; + SelectTypeAggregationExpressionChecker selectTypeAggregationExpressionChecker = new SelectTypeAggregationExpressionChecker(); + SelectTypeIdentifierExpressionSelector selectTypeIdentifierExpressionSelector = new SelectTypeIdentifierExpressionSelector(); + + Boolean hasAggregationFunction = query.getSelections().stream() + .filter(selectionSpec -> selectionSpec.getExpression().accept(selectTypeAggregationExpressionChecker)) + .findAny() + .isEmpty(); + + if (!hasAggregationFunction) return query; + + // get all identifier of group by clause + LocalGroupByExpressionVisitor groupByExpressionVisitor = new LocalGroupByExpressionVisitor(); + Set groupByExpressions = query.getAggregations().stream() + .map(gte -> (String) gte.accept(groupByExpressionVisitor)) + .filter(s -> !Strings.isNullOrEmpty(s)) + .collect(Collectors.toUnmodifiableSet()); // check for all selections, remove non-aggregated selections. List finalSelectionSpecs = query.getSelections().stream() - .filter(selectionSpec -> selectionSpec.getExpression().accept(this)) + .filter(selectionSpec -> (boolean) selectionSpec.getExpression().accept(selectTypeAggregationExpressionChecker) + && groupByExpressions.contains((String)selectionSpec.getExpression().accept(selectTypeIdentifierExpressionSelector))) .collect(Collectors.toUnmodifiableList()); return finalSelectionSpecs.size() > 0 @@ -48,23 +65,59 @@ public Query transform(Query query) { : query; } - @Override - public Boolean visit(AggregateExpression expression) { - return true; - } + private static class SelectTypeAggregationExpressionChecker implements SelectTypeExpressionVisitor { + @Override + public Boolean visit(AggregateExpression expression) { + return true; + } - @Override - public Boolean visit(ConstantExpression 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 false; + } } - @Override - public Boolean visit(FunctionExpression expression) { - return false; + private static class SelectTypeIdentifierExpressionSelector implements SelectTypeExpressionVisitor { + @Override + public String visit(AggregateExpression expression) { + return null; + } + + @Override + public String visit(ConstantExpression expression) { + return null; + } + + @Override + public String visit(FunctionExpression expression) { + return null; + } + + @Override + public String visit(IdentifierExpression expression) { + return expression.getName(); + } } - @Override - public Boolean visit(IdentifierExpression expression) { - return false; + private static class LocalGroupByExpressionVisitor implements GroupTypeExpressionVisitor{ + @Override + public String visit(FunctionExpression expression) { + return null; + } + + @Override + public String visit(IdentifierExpression expression) { + return expression.getName(); + } } } 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 077d9817..07bc46d3 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 @@ -917,4 +917,33 @@ 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, document->'price' AS price " + + "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)); + } } From 4d5bd20718c61cf89f72e5a348e33e63376055df Mon Sep 17 00:00:00 2001 From: Ronak Date: Mon, 5 Sep 2022 09:58:27 +0530 Subject: [PATCH 2/4] fix: handle group by more generically --- .../documentstore/DocStoreQueryV1Test.java | 9 +- ...aggr_with_match_selection_and_groupby.json | 14 ++++ .../PostgresSelectionQueryTransformer.java | 82 +++++++++++-------- .../query/v1/PostgresQueryParserTest.java | 6 +- 4 files changed, 71 insertions(+), 40 deletions(-) create mode 100644 document-store/src/integrationTest/resources/mongo/test_aggr_with_match_selection_and_groupby.json 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 19c635c4..84dab285 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 @@ -745,7 +745,8 @@ void testQueryQ1DistinctCountAggregationWithOnlyFilter(String dataStoreName) thr @ParameterizedTest @MethodSource("databaseContextBoth") - void testQueryQ1DistinctCountAggregationWithMatchingSelectionAndGroupBy(String dataStoreName) throws IOException { + void testQueryQ1DistinctCountAggregationWithMatchingSelectionAndGroupBy(String dataStoreName) + throws IOException { Datastore datastore = datastoreMap.get(dataStoreName); Collection collection = datastore.getCollection(COLLECTION_NAME); org.hypertrace.core.documentstore.query.Query query = @@ -755,12 +756,14 @@ void testQueryQ1DistinctCountAggregationWithMatchingSelectionAndGroupBy(String d "qty_count") .addSelection(IdentifierExpression.of("item")) .addSelection(IdentifierExpression.of("price")) - .setFilter(RelationalExpression.of(IdentifierExpression.of("price"), LTE, ConstantExpression.of(10))) + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("price"), LTE, ConstantExpression.of(10))) .addAggregation(IdentifierExpression.of("item")) .build(); try (CloseableIterator resultDocs = collection.aggregate(query)) { assertDocsAndSizeEqualWithoutOrder( - dataStoreName, resultDocs, 1, "mongo/test_aggr_only_with_fliter_response.json"); + dataStoreName, resultDocs, 3, "mongo/test_aggr_with_match_selection_and_groupby.json"); } } diff --git a/document-store/src/integrationTest/resources/mongo/test_aggr_with_match_selection_and_groupby.json b/document-store/src/integrationTest/resources/mongo/test_aggr_with_match_selection_and_groupby.json new file mode 100644 index 00000000..ef642f68 --- /dev/null +++ b/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 + } +] \ No newline at end of file diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresSelectionQueryTransformer.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresSelectionQueryTransformer.java index 954395b5..d15d5126 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresSelectionQueryTransformer.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresSelectionQueryTransformer.java @@ -1,6 +1,5 @@ package org.hypertrace.core.documentstore.postgres.query.v1.transformer; -import com.google.common.base.Strings; import java.util.List; import java.util.Set; import java.util.stream.Collectors; @@ -8,6 +7,8 @@ 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; @@ -36,36 +37,45 @@ public class PostgresSelectionQueryTransformer implements QueryTransformer { @Override public Query transform(Query query) { - SelectTypeAggregationExpressionChecker selectTypeAggregationExpressionChecker = new SelectTypeAggregationExpressionChecker(); - SelectTypeIdentifierExpressionSelector selectTypeIdentifierExpressionSelector = new SelectTypeIdentifierExpressionSelector(); - - Boolean hasAggregationFunction = query.getSelections().stream() - .filter(selectionSpec -> selectionSpec.getExpression().accept(selectTypeAggregationExpressionChecker)) - .findAny() - .isEmpty(); - - if (!hasAggregationFunction) 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 - LocalGroupByExpressionVisitor groupByExpressionVisitor = new LocalGroupByExpressionVisitor(); - Set groupByExpressions = query.getAggregations().stream() - .map(gte -> (String) gte.accept(groupByExpressionVisitor)) - .filter(s -> !Strings.isNullOrEmpty(s)) - .collect(Collectors.toUnmodifiableSet()); - - // check for all selections, remove non-aggregated selections. + LocalGroupByIdentifierExpressionSelector groupByIdentifierExpressionSelector = + new LocalGroupByIdentifierExpressionSelector(); + Set groupByIdentifierExpressions = + query.getAggregations().stream() + .filter(exp -> exp.accept(groupByIdentifierExpressionSelector)) + .collect(Collectors.toUnmodifiableSet()); + + // keep only matching identifier expression with group by along with rest. + LocalSelectTypeIdentifierExpressionSelector identifierExpressionSelector = + new LocalSelectTypeIdentifierExpressionSelector(); List finalSelectionSpecs = query.getSelections().stream() - .filter(selectionSpec -> (boolean) selectionSpec.getExpression().accept(selectTypeAggregationExpressionChecker) - && groupByExpressions.contains((String)selectionSpec.getExpression().accept(selectTypeIdentifierExpressionSelector))) + .filter( + selectionSpec -> { + SelectTypeExpression expression = selectionSpec.getExpression(); + return !((boolean) expression.accept(identifierExpressionSelector)) + || groupByIdentifierExpressions.contains(expression); + }) .collect(Collectors.toUnmodifiableList()); - return finalSelectionSpecs.size() > 0 + return finalSelectionSpecs.size() != query.getSelections().size() ? new TransformedQueryBuilder(query).setSelections(finalSelectionSpecs).build() : query; } - private static class SelectTypeAggregationExpressionChecker implements SelectTypeExpressionVisitor { + private static class LocalSelectTypeAggregationExpressionSelector + implements SelectTypeExpressionVisitor { @Override public Boolean visit(AggregateExpression expression) { return true; @@ -87,37 +97,39 @@ public Boolean visit(IdentifierExpression expression) { } } - private static class SelectTypeIdentifierExpressionSelector implements SelectTypeExpressionVisitor { + private static class LocalSelectTypeIdentifierExpressionSelector + implements SelectTypeExpressionVisitor { @Override - public String visit(AggregateExpression expression) { - return null; + public Boolean visit(AggregateExpression expression) { + return false; } @Override - public String visit(ConstantExpression expression) { - return null; + public Boolean visit(ConstantExpression expression) { + return false; } @Override - public String visit(FunctionExpression expression) { - return null; + public Boolean visit(FunctionExpression expression) { + return false; } @Override - public String visit(IdentifierExpression expression) { - return expression.getName(); + public Boolean visit(IdentifierExpression expression) { + return true; } } - private static class LocalGroupByExpressionVisitor implements GroupTypeExpressionVisitor{ + private static class LocalGroupByIdentifierExpressionSelector + implements GroupTypeExpressionVisitor { @Override - public String visit(FunctionExpression expression) { - return null; + public Boolean visit(FunctionExpression expression) { + return false; } @Override - public String visit(IdentifierExpression expression) { - return expression.getName(); + public Boolean visit(IdentifierExpression expression) { + return true; } } } 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 07bc46d3..e88b02a5 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 @@ -927,7 +927,9 @@ void testQueryQ1DistinctCountAggregationWithMatchingSelectionAndGroupBy() { "qty_count") .addSelection(IdentifierExpression.of("item")) .addSelection(IdentifierExpression.of("price")) - .setFilter(RelationalExpression.of(IdentifierExpression.of("price"), LTE, ConstantExpression.of(10))) + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("price"), LTE, ConstantExpression.of(10))) .addAggregation(IdentifierExpression.of("item")) .build(); @@ -937,7 +939,7 @@ void testQueryQ1DistinctCountAggregationWithMatchingSelectionAndGroupBy() { assertEquals( "SELECT COUNT(DISTINCT document->>'quantity' ) AS \"qty_count\", " - + "document->'item' AS item, document->'price' AS price " + + "document->'item' AS item " + "FROM testCollection WHERE CAST (document->>'price' AS NUMERIC) <= ? " + "GROUP BY document->'item'", sql); From 5db4d5c7c3e6d306a305552b7bd12e9245ef478b Mon Sep 17 00:00:00 2001 From: Ronak Date: Mon, 5 Sep 2022 10:52:06 +0530 Subject: [PATCH 3/4] addressed comments --- .../PostgresSelectionQueryTransformer.java | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresSelectionQueryTransformer.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresSelectionQueryTransformer.java index d15d5126..25c1fb15 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresSelectionQueryTransformer.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresSelectionQueryTransformer.java @@ -2,6 +2,7 @@ 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; @@ -62,11 +63,8 @@ public Query transform(Query query) { List finalSelectionSpecs = query.getSelections().stream() .filter( - selectionSpec -> { - SelectTypeExpression expression = selectionSpec.getExpression(); - return !((boolean) expression.accept(identifierExpressionSelector)) - || groupByIdentifierExpressions.contains(expression); - }) + getSelectionSpecPredicate( + groupByIdentifierExpressions, identifierExpressionSelector)) .collect(Collectors.toUnmodifiableList()); return finalSelectionSpecs.size() != query.getSelections().size() @@ -74,6 +72,16 @@ public Query transform(Query query) { : query; } + private Predicate getSelectionSpecPredicate( + Set groupByIdentifierExpressions, + LocalSelectTypeIdentifierExpressionSelector identifierExpressionSelector) { + return selectionSpec -> { + SelectTypeExpression expression = selectionSpec.getExpression(); + return ((boolean) expression.accept(identifierExpressionSelector)) + || groupByIdentifierExpressions.contains(expression); + }; + } + private static class LocalSelectTypeAggregationExpressionSelector implements SelectTypeExpressionVisitor { @Override From d5e3360e5f54dfa78e4594945eaec4b5cbfef4ce Mon Sep 17 00:00:00 2001 From: Ronak Date: Mon, 5 Sep 2022 11:09:29 +0530 Subject: [PATCH 4/4] fixed the broken tests due to refactoring --- .../query/v1/transformer/PostgresSelectionQueryTransformer.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresSelectionQueryTransformer.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresSelectionQueryTransformer.java index 25c1fb15..a0d52d00 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresSelectionQueryTransformer.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresSelectionQueryTransformer.java @@ -77,7 +77,7 @@ private Predicate getSelectionSpecPredicate( LocalSelectTypeIdentifierExpressionSelector identifierExpressionSelector) { return selectionSpec -> { SelectTypeExpression expression = selectionSpec.getExpression(); - return ((boolean) expression.accept(identifierExpressionSelector)) + return (!(boolean) expression.accept(identifierExpressionSelector)) || groupByIdentifierExpressions.contains(expression); }; }