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 a23c0a7a..2b8d33e8 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 @@ -1249,6 +1249,38 @@ public void testQueryV1AggregationFilter(String dataStoreName) throws IOExceptio dataStoreName, resultDocs, "query/distinct_count_response.json", 4); } + @ParameterizedTest + @ArgumentsSource(AllProvider.class) + public void testQueryV1AggregationWithInFilterWithPrimitiveLhs(final String dataStoreName) throws IOException { + final Collection collection = getCollection(dataStoreName); + final Query query = Query.builder() + .addSelection(IdentifierExpression.of("item")) + .addSelection(IdentifierExpression.of("quantity")) + .setFilter( + RelationalExpression.of(IdentifierExpression.of("item"), IN, ConstantExpression.ofStrings(List.of("Comb", "Shampoo"))) + ) + .build(); + + final Iterator resultDocs = collection.aggregate(query); + assertDocsAndSizeEqualWithoutOrder(dataStoreName, resultDocs, "query/test_primitive_lhs_in_filter_aggr_response.json", 4); + } + + @ParameterizedTest + @ArgumentsSource(AllProvider.class) + public void testQueryV1AggregationWithInFilterWithArrayLhs(final String dataStoreName) throws IOException { + final Collection collection = getCollection(dataStoreName); + final Query query = Query.builder() + .addSelection(IdentifierExpression.of("item")) + .addSelection(IdentifierExpression.of("price")) + .setFilter( + RelationalExpression.of(IdentifierExpression.of("props.colors"), IN, ConstantExpression.ofStrings(List.of("Orange"))) + ) + .build(); + + final Iterator resultDocs = collection.aggregate(query); + assertDocsAndSizeEqualWithoutOrder(dataStoreName, resultDocs, "query/test_json_column_array_lhs_in_filter_aggr_response.json", 1); + } + @ParameterizedTest @ArgumentsSource(AllProvider.class) public void testQueryV1AggregationFilterWithWhereClause(String dataStoreName) throws IOException { diff --git a/document-store/src/integrationTest/resources/query/test_json_column_array_lhs_in_filter_aggr_response.json b/document-store/src/integrationTest/resources/query/test_json_column_array_lhs_in_filter_aggr_response.json new file mode 100644 index 00000000..e951f265 --- /dev/null +++ b/document-store/src/integrationTest/resources/query/test_json_column_array_lhs_in_filter_aggr_response.json @@ -0,0 +1,6 @@ +[ + { + "item": "Soap", + "price": 20 + } +] \ No newline at end of file diff --git a/document-store/src/integrationTest/resources/query/test_primitive_lhs_in_filter_aggr_response.json b/document-store/src/integrationTest/resources/query/test_primitive_lhs_in_filter_aggr_response.json new file mode 100644 index 00000000..f13cab92 --- /dev/null +++ b/document-store/src/integrationTest/resources/query/test_primitive_lhs_in_filter_aggr_response.json @@ -0,0 +1,18 @@ +[ + { + "item":"Comb", + "quantity": 10 + }, + { + "item":"Comb", + "quantity": 5 + }, + { + "item":"Shampoo", + "quantity": 20 + }, + { + "item":"Shampoo", + "quantity": 10 + } +] \ No newline at end of file diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFilterTypeExpressionVisitor.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFilterTypeExpressionVisitor.java index 032597c7..9438fd18 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFilterTypeExpressionVisitor.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFilterTypeExpressionVisitor.java @@ -111,6 +111,8 @@ private boolean isOperatorNeedsFieldAccessor(RelationalOperator operator) { case NOT_CONTAINS: case EXISTS: case NOT_EXISTS: + case IN: + case NOT_IN: return true; default: return false; diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/utils/PostgresUtils.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/utils/PostgresUtils.java index 1c27ce95..a67716bc 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/utils/PostgresUtils.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/utils/PostgresUtils.java @@ -360,8 +360,8 @@ private static StringBuilder prepareFilterStringForInOperator( StreamSupport.stream(values.spliterator(), false) .map( val -> { - paramsBuilder.addObjectParam(val); - return "(jsonb_build_array(" + parsedExpression + ") @> jsonb_build_array(?))"; + paramsBuilder.addObjectParam(val).addObjectParam(val); + return "((jsonb_typeof(to_jsonb("+parsedExpression+")) = 'array' AND to_jsonb("+parsedExpression+") @> jsonb_build_array(?)) OR (jsonb_build_array(" + parsedExpression + ") @> jsonb_build_array(?)))"; }) .collect(Collectors.joining(" OR ")); filterStringBuilder.append(collect);