address review comments

hypertrace · Dec 15, 2023 · d1d4ac7 · d1d4ac7
1 parent a039f9d
commit d1d4ac7
Show file tree

Hide file tree

Showing 5 changed files with 60 additions and 2 deletions.
diff --git a/...store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java b/...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<Document> 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<Document> 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/...c/integrationTest/resources/query/test_json_column_array_lhs_in_filter_aggr_response.json b/...c/integrationTest/resources/query/test_json_column_array_lhs_in_filter_aggr_response.json
@@ -0,0 +1,6 @@
+[
+  {
+    "item": "Soap",
+    "price": 20
+  }
+]
diff --git a/...store/src/integrationTest/resources/query/test_primitive_lhs_in_filter_aggr_response.json b/...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
+  }
+]
diff --git a/...ace/core/documentstore/postgres/query/v1/vistors/PostgresFilterTypeExpressionVisitor.java b/...ace/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/...t-store/src/main/java/org/hypertrace/core/documentstore/postgres/utils/PostgresUtils.java b/...t-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);