make inclause behaviour consistent between mongo and postgres (#186)

hypertrace · Dec 15, 2023 · 782fddb · 782fddb
1 parent 35d87df
commit 782fddb
Show file tree

Hide file tree

Showing 7 changed files with 126 additions and 23 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,51 @@ 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/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreTest.java
@@ -666,7 +666,6 @@ public void testSubDocumentUpdate(String dataStoreName) throws IOException {
     // postgres
     // {"foo1":"bar1","subdoc":{"subfoo1":"subbar1","nesteddoc":{"nestedfoo1":"nestedbar1"}},
     // "created_at":"2021-03-15 00:24:50.981147","updated_at":"2021-03-15 00:24:50.981147"}
-    System.out.println(documents.get(0).toJson());
     ObjectNode jsonNode = (ObjectNode) OBJECT_MAPPER.readTree(documents.get(0).toJson());
     String expected =
         "{\"foo1\":\"bar1\",\"subdoc\":{\"subfoo1\":\"subbar1\","

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
@@ -352,6 +352,25 @@ public static String parseNonCompositeFilter(
     return filterString.toString();
   }
 
+  private static StringBuilder prepareFilterStringForInOperator(
+      final String parsedExpression, final Iterable<Object> values, final Builder paramsBuilder) {
+    final StringBuilder filterStringBuilder = new StringBuilder();
+    filterStringBuilder.append("(");
+    final String collect =
+        StreamSupport.stream(values.spliterator(), false)
+            .map(
+                val -> {
+                  paramsBuilder.addObjectParam(val).addObjectParam(val);
+                  return String.format(
+                      "((jsonb_typeof(to_jsonb(%s)) = 'array' AND to_jsonb(%s) @> jsonb_build_array(?)) OR (jsonb_build_array(%s) @> jsonb_build_array(?)))",
+                      parsedExpression, parsedExpression, parsedExpression);
+                })
+            .collect(Collectors.joining(" OR "));
+    filterStringBuilder.append(collect);
+    filterStringBuilder.append(")");
+    return filterStringBuilder;
+  }
+
   private static Object prepareJsonValueForContainsOp(final Object value) {
     if (value instanceof Document) {
       try {
@@ -431,17 +450,31 @@ public static String prepareParsedNonCompositeFilter(
         break;
       case "NOT_IN":
       case "NOT IN":
-        // NOTE: Pl. refer this in non-parsed expression for limitation of this filter
+        // In order to make the behaviour same as for Mongo, the "NOT_IN" operator would match only
+        // if
+        // the LHS and RHS have no intersection at all
+        // NOTE: This doesn't work in case the LHS is a function
         sqlOperator = " NOT IN ";
         isMultiValued = true;
-        value = prepareParameterizedStringForList((Iterable<Object>) value, paramsBuilder);
-        break;
+        filterString
+            .append(" IS NULL OR")
+            .append(" NOT (")
+            .append(
+                prepareFilterStringForInOperator(
+                    preparedExpression, (Iterable<Object>) value, paramsBuilder))
+            .append(")");
+        return filterString.toString();
       case "IN":
+        // In order to make the behaviour same as for Mongo, the "INI" operator would match if the
+        // LHS and RHS have any intersection (i.e. non-empty intersection)
         // NOTE: Pl. refer this in non-parsed expression for limitation of this filter
+        // NOTE: This doesn't work in case the LHS is a function
         sqlOperator = " IN ";
         isMultiValued = true;
-        value = prepareParameterizedStringForList((Iterable<Object>) value, paramsBuilder);
-        break;
+        filterString =
+            prepareFilterStringForInOperator(
+                preparedExpression, (Iterable<Object>) value, paramsBuilder);
+        return filterString.toString();
       case "NOT_EXISTS":
       case "NOT EXISTS":
         sqlOperator = " IS NULL ";

diff --git a/...est/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParserTest.java b/...est/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParserTest.java
@@ -648,19 +648,19 @@ void testFindWithSortingAndPagination() {
             + "document->'quantity' AS \"quantity\", "
             + "document->'date' AS \"date\" "
             + "FROM \"testCollection\" "
-            + "WHERE document->>'item' IN (?, ?, ?, ?) "
+            + "WHERE (((jsonb_typeof(to_jsonb(document->'item')) = 'array' AND to_jsonb(document->'item') @> jsonb_build_array(?)) OR (jsonb_build_array(document->'item') @> jsonb_build_array(?))) OR ((jsonb_typeof(to_jsonb(document->'item')) = 'array' AND to_jsonb(document->'item') @> jsonb_build_array(?)) OR (jsonb_build_array(document->'item') @> jsonb_build_array(?))) OR ((jsonb_typeof(to_jsonb(document->'item')) = 'array' AND to_jsonb(document->'item') @> jsonb_build_array(?)) OR (jsonb_build_array(document->'item') @> jsonb_build_array(?))) OR ((jsonb_typeof(to_jsonb(document->'item')) = 'array' AND to_jsonb(document->'item') @> jsonb_build_array(?)) OR (jsonb_build_array(document->'item') @> jsonb_build_array(?)))) "
             + "ORDER BY document->'quantity' DESC NULLS LAST,document->'item' ASC NULLS FIRST "
             + "OFFSET ? LIMIT ?",
         sql);
 
     Params params = postgresQueryParser.getParamsBuilder().build();
-    assertEquals(6, params.getObjectParams().size());
+    assertEquals(10, params.getObjectParams().size());
     assertEquals("Mirror", params.getObjectParams().get(1));
-    assertEquals("Comb", params.getObjectParams().get(2));
-    assertEquals("Shampoo", params.getObjectParams().get(3));
-    assertEquals("Bottle", params.getObjectParams().get(4));
-    assertEquals(1, params.getObjectParams().get(5));
-    assertEquals(3, params.getObjectParams().get(6));
+    assertEquals("Comb", params.getObjectParams().get(3));
+    assertEquals("Shampoo", params.getObjectParams().get(5));
+    assertEquals("Bottle", params.getObjectParams().get(7));
+    assertEquals(1, params.getObjectParams().get(9));
+    assertEquals(3, params.getObjectParams().get(10));
   }
 
   @Test
@@ -1121,16 +1121,16 @@ void testQueryQ1AggregationFilterWithStringInFilterAlongWithNonAliasFields() {
         "SELECT COUNT(DISTINCT document->>'quantity' ) AS \"qty_count\", "
             + "document->'item' AS \"item\", document->'price' AS \"price\" "
             + "FROM \"testCollection\" GROUP BY document->'item',document->'price' "
-            + "HAVING (COUNT(DISTINCT document->>'quantity' ) <= ?) AND (CAST (document->'item' AS TEXT) IN (?, ?, ?, ?))",
+            + "HAVING (COUNT(DISTINCT document->>'quantity' ) <= ?) AND ((((jsonb_typeof(to_jsonb(CAST (document->'item' AS TEXT))) = 'array' AND to_jsonb(CAST (document->'item' AS TEXT)) @> jsonb_build_array(?)) OR (jsonb_build_array(CAST (document->'item' AS TEXT)) @> jsonb_build_array(?))) OR ((jsonb_typeof(to_jsonb(CAST (document->'item' AS TEXT))) = 'array' AND to_jsonb(CAST (document->'item' AS TEXT)) @> jsonb_build_array(?)) OR (jsonb_build_array(CAST (document->'item' AS TEXT)) @> jsonb_build_array(?))) OR ((jsonb_typeof(to_jsonb(CAST (document->'item' AS TEXT))) = 'array' AND to_jsonb(CAST (document->'item' AS TEXT)) @> jsonb_build_array(?)) OR (jsonb_build_array(CAST (document->'item' AS TEXT)) @> jsonb_build_array(?))) OR ((jsonb_typeof(to_jsonb(CAST (document->'item' AS TEXT))) = 'array' AND to_jsonb(CAST (document->'item' AS TEXT)) @> jsonb_build_array(?)) OR (jsonb_build_array(CAST (document->'item' AS TEXT)) @> jsonb_build_array(?)))))",
         sql);
 
     Params params = postgresQueryParser.getParamsBuilder().build();
-    assertEquals(5, params.getObjectParams().size());
+    assertEquals(9, params.getObjectParams().size());
     assertEquals(10, params.getObjectParams().get(1));
     assertEquals("\"Mirror\"", params.getObjectParams().get(2));
-    assertEquals("\"Comb\"", params.getObjectParams().get(3));
-    assertEquals("\"Shampoo\"", params.getObjectParams().get(4));
-    assertEquals("\"Bottle\"", params.getObjectParams().get(5));
+    assertEquals("\"Comb\"", params.getObjectParams().get(4));
+    assertEquals("\"Shampoo\"", params.getObjectParams().get(6));
+    assertEquals("\"Bottle\"", params.getObjectParams().get(8));
   }
 
   @Test
@@ -1164,16 +1164,16 @@ void testQueryQ1DistinctCountAggregationWithOnlyFilter() {
     assertEquals(
         "SELECT COUNT(DISTINCT document->>'quantity' ) AS \"qty_count\" "
             + "FROM \"testCollection\" "
-            + "WHERE (CAST (document->>'price' AS NUMERIC) <= ?) AND (document->>'item' IN (?, ?, ?, ?))",
+            + "WHERE (CAST (document->>'price' AS NUMERIC) <= ?) AND ((((jsonb_typeof(to_jsonb(document->'item')) = 'array' AND to_jsonb(document->'item') @> jsonb_build_array(?)) OR (jsonb_build_array(document->'item') @> jsonb_build_array(?))) OR ((jsonb_typeof(to_jsonb(document->'item')) = 'array' AND to_jsonb(document->'item') @> jsonb_build_array(?)) OR (jsonb_build_array(document->'item') @> jsonb_build_array(?))) OR ((jsonb_typeof(to_jsonb(document->'item')) = 'array' AND to_jsonb(document->'item') @> jsonb_build_array(?)) OR (jsonb_build_array(document->'item') @> jsonb_build_array(?))) OR ((jsonb_typeof(to_jsonb(document->'item')) = 'array' AND to_jsonb(document->'item') @> jsonb_build_array(?)) OR (jsonb_build_array(document->'item') @> jsonb_build_array(?)))))",
         sql);
 
     Params params = postgresQueryParser.getParamsBuilder().build();
-    assertEquals(5, params.getObjectParams().size());
+    assertEquals(9, params.getObjectParams().size());
     assertEquals(10, params.getObjectParams().get(1));
     assertEquals("Mirror", params.getObjectParams().get(2));
-    assertEquals("Comb", params.getObjectParams().get(3));
-    assertEquals("Shampoo", params.getObjectParams().get(4));
-    assertEquals("Bottle", params.getObjectParams().get(5));
+    assertEquals("Comb", params.getObjectParams().get(4));
+    assertEquals("Shampoo", params.getObjectParams().get(6));
+    assertEquals("Bottle", params.getObjectParams().get(8));
   }
 
   @Test