Add ALL and ONE array filter operators for MongoDB and Postgres

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

@@ -2,7 +2,9 @@
 
 import static org.hypertrace.core.documentstore.expression.impl.LogicalExpression.and;
 import static org.hypertrace.core.documentstore.expression.impl.LogicalExpression.not;
+import static org.hypertrace.core.documentstore.expression.operators.ArrayOperator.ALL;
 import static org.hypertrace.core.documentstore.expression.operators.ArrayOperator.ANY;
+import static org.hypertrace.core.documentstore.expression.operators.ArrayOperator.ONE;
 import static org.hypertrace.core.documentstore.model.config.DatabaseType.MONGO;
 import static org.hypertrace.core.documentstore.model.config.DatabaseType.POSTGRES;
 import static org.hypertrace.core.documentstore.utils.Utils.MONGO_STORE;
@@ -11,6 +13,7 @@
 import com.google.common.io.Resources;
 import java.io.IOException;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
 import java.util.Spliterator;
 import java.util.Spliterators;
@@ -344,6 +347,84 @@ void getDocumentsWithEnvironmentIdsSubsetOfGivenList(final String dataStoreName)
     JSONAssert.assertEquals(expected, actual, JSONCompareMode.LENIENT);
   }
 
+  /**
+   * Tests MATCH_ALL semantics: documents whose array attribute contains every value specified in
+   * the filter.
+   */
+  @ParameterizedTest
+  @ArgumentsSource(AllProvider.class)
+  void getDocumentsContainingAllGivenValues(final String dataStoreName)
+      throws JSONException, IOException {
+    final String testCollectionName = "array_match_test";
+    final Datastore datastore = datastoreMap.get(dataStoreName);
+    final Map<Key, Document> testDocuments =
+        Utils.buildDocumentsFromResource("query/array_operators/array_match_test.json");
+    datastore.deleteCollection(testCollectionName);
+    datastore.createCollection(testCollectionName, null);
+    final Collection collection = datastore.getCollection(testCollectionName);
+    collection.bulkUpsert(testDocuments);
+
+    final Query query =
+        Query.builder()
+            .setFilter(
+                ArrayRelationalFilterExpression.builder()
+                    .operator(ALL)
+                    .filter(
+                        RelationalExpression.of(
+                            IdentifierExpression.of("tags"),
+                            RelationalOperator.IN,
+                            ConstantExpression.ofStrings(List.of("red", "blue"))))
+                    .build())
+            .build();
+
+    final Iterator<Document> documents = collection.aggregate(query);
+    final String expected = readResource("array_match_all_result.json");
+    final String actual = iteratorToJson(documents);
+
+    datastore.deleteCollection(testCollectionName);
+
+    JSONAssert.assertEquals(expected, actual, JSONCompareMode.LENIENT);
+  }
+
+  /**
+   * Tests MATCH_ONE semantics: documents whose array attribute has exactly one element, and that
+   * element is one of the values specified in the filter.
+   */
+  @ParameterizedTest
+  @ArgumentsSource(AllProvider.class)
+  void getDocumentsWithExactlyOneElementMatchingGivenValues(final String dataStoreName)
+      throws JSONException, IOException {
+    final String testCollectionName = "array_match_test";
+    final Datastore datastore = datastoreMap.get(dataStoreName);
+    final Map<Key, Document> testDocuments =
+        Utils.buildDocumentsFromResource("query/array_operators/array_match_test.json");
+    datastore.deleteCollection(testCollectionName);
+    datastore.createCollection(testCollectionName, null);
+    final Collection collection = datastore.getCollection(testCollectionName);
+    collection.bulkUpsert(testDocuments);
+
+    final Query query =
+        Query.builder()
+            .setFilter(
+                ArrayRelationalFilterExpression.builder()
+                    .operator(ONE)
+                    .filter(
+                        RelationalExpression.of(
+                            IdentifierExpression.of("tags"),
+                            RelationalOperator.IN,
+                            ConstantExpression.ofStrings(List.of("red", "blue"))))
+                    .build())
+            .build();
+
+    final Iterator<Document> documents = collection.aggregate(query);
+    final String expected = readResource("array_match_one_result.json");
+    final String actual = iteratorToJson(documents);
+
+    datastore.deleteCollection(testCollectionName);
+
+    JSONAssert.assertEquals(expected, actual, JSONCompareMode.LENIENT);
+  }
+
   private String readResource(final String fileName) {
     try {
       return new String(

document-store/src/integrationTest/resources/query/array_operators/array_match_all_result.json

@@ -0,0 +1,10 @@
+[
+  {
+    "name": "Document A",
+    "tags": ["red", "blue"]
+  },
+  {
+    "name": "Document B",
+    "tags": ["red", "blue", "green"]
+  }
+]

document-store/src/integrationTest/resources/query/array_operators/array_match_one_result.json

@@ -0,0 +1,6 @@
+[
+  {
+    "name": "Document C",
+    "tags": ["red"]
+  }
+]

document-store/src/integrationTest/resources/query/array_operators/array_match_test.json

@@ -0,0 +1,31 @@
+[
+  {
+    "_id": 1,
+    "name": "Document A",
+    "tags": ["red", "blue"]
+  },
+  {
+    "_id": 2,
+    "name": "Document B",
+    "tags": ["red", "blue", "green"]
+  },
+  {
+    "_id": 3,
+    "name": "Document C",
+    "tags": ["red"]
+  },
+  {
+    "_id": 4,
+    "name": "Document D",
+    "tags": ["yellow"]
+  },
+  {
+    "_id": 5,
+    "name": "Document E",
+    "tags": []
+  },
+  {
+    "_id": 6,
+    "name": "Document F"
+  }
+]

document-store/src/main/java/org/hypertrace/core/documentstore/expression/operators/ArrayOperator.java

@@ -2,5 +2,9 @@
 
 public enum ArrayOperator {
   ANY,
-  // Can support ALL and NONE later
+  // Array attribute must contain every value specified in the filter
+  ALL,
+  // Array attribute must contain exactly one element, and that element must be one of the values
+  // specified in the filter
+  ONE,
 }

document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoArrayFilterParser.java

@@ -6,10 +6,14 @@
 import static org.hypertrace.core.documentstore.mongo.query.parser.filter.MongoStandardExprRelationalFilterParser.EXPR;
 
 import com.google.common.collect.Maps;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import org.hypertrace.core.documentstore.expression.impl.ArrayFilterExpression;
+import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
+import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
 import org.hypertrace.core.documentstore.expression.operators.ArrayOperator;
+import org.hypertrace.core.documentstore.expression.type.FilterTypeExpression;
 import org.hypertrace.core.documentstore.expression.type.SelectTypeExpression;
 import org.hypertrace.core.documentstore.mongo.MongoUtils;
 import org.hypertrace.core.documentstore.mongo.query.parser.filter.MongoRelationalFilterParserFactory.MongoRelationalFilterContext;
@@ -21,6 +25,12 @@ class MongoArrayFilterParser {
   private static final String IF_NULL = "$ifNull";
   private static final String AS = "as";
   private static final String IN = "in";
+  private static final String SET_IS_SUBSET = "$setIsSubset";
+  private static final String SIZE = "$size";
+  private static final String EQ = "$eq";
+  private static final String IN_OPERATOR = "$in";
+  private static final String ARRAY_ELEM_AT = "$arrayElemAt";
+  private static final String AND = "$and";
 
   private static final Map<ArrayOperator, String> OPERATOR_MAP =
       Maps.immutableEnumMap(Map.ofEntries(entry(ANY, ANY_ELEMENT_TRUE)));
@@ -39,6 +49,17 @@ class MongoArrayFilterParser {
   }
 
   Map<String, Object> parse(final ArrayFilterExpression arrayFilterExpression) {
+    switch (arrayFilterExpression.getOperator()) {
+      case ALL:
+        return parseAllOperator(arrayFilterExpression);
+      case ONE:
+        return parseOneOperator(arrayFilterExpression);
+      default:
+        return parseAnyOperator(arrayFilterExpression);
+    }
+  }
+
+  private Map<String, Object> parseAnyOperator(final ArrayFilterExpression arrayFilterExpression) {
     final String operator =
         Optional.ofNullable(OPERATOR_MAP.get(arrayFilterExpression.getOperator()))
             .orElseThrow(
@@ -103,9 +124,84 @@ Map<String, Object> parse(final ArrayFilterExpression arrayFilterExpression) {
                     entry(INPUT, Map.of(IF_NULL, new Object[] {mapInput, new Object[0]})),
                     entry(AS, alias),
                     entry(IN, filter))));
+    return wrapInExprIfNeeded(arrayFilter);
+  }
+
+  /*
+  {
+    "$expr": {
+      "$setIsSubset": [
+        ["Blue", "Green"],
+        { "$ifNull": ["$colors", []] }
+      ]
+    }
+  }
+   */
+  private Map<String, Object> parseAllOperator(final ArrayFilterExpression arrayFilterExpression) {
+    final Object mapInput = getDollarPrefixedArraySource(arrayFilterExpression);
+    final List<?> values = getFilterValues(arrayFilterExpression);
+
+    final Map<String, Object> setIsSubset =
+        Map.of(
+            SET_IS_SUBSET,
+            List.of(values, Map.of(IF_NULL, new Object[] {mapInput, new Object[0]})));
+    return wrapInExprIfNeeded(setIsSubset);
+  }
+
+  /*
+  {
+    "$expr": {
+      "$and": [
+        { "$eq": [{ "$size": { "$ifNull": ["$colors", []] } }, 1] },
+        { "$in": [{ "$arrayElemAt": [{ "$ifNull": ["$colors", []] }, 0] }, ["Blue", "Green"]] }
+      ]
+    }
+  }
+   */
+  private Map<String, Object> parseOneOperator(final ArrayFilterExpression arrayFilterExpression) {
+    final Object mapInput = getDollarPrefixedArraySource(arrayFilterExpression);
+    final List<?> values = getFilterValues(arrayFilterExpression);
+    final Map<String, Object> arrayWithDefault =
+        Map.of(IF_NULL, new Object[] {mapInput, new Object[0]});
+
+    final Map<String, Object> sizeIsOne = Map.of(EQ, List.of(Map.of(SIZE, arrayWithDefault), 1));
+    final Map<String, Object> firstElementMatches =
+        Map.of(IN_OPERATOR, List.of(Map.of(ARRAY_ELEM_AT, List.of(arrayWithDefault, 0)), values));
+
+    return wrapInExprIfNeeded(Map.of(AND, List.of(sizeIsOne, firstElementMatches)));
+  }
+
+  private String getDollarPrefixedArraySource(final ArrayFilterExpression arrayFilterExpression) {
+    final MongoSelectTypeExpressionParser wrappingParser =
+        new MongoDollarPrefixingIdempotentParser(relationalFilterContext.lhsParser());
+    return arrayFilterExpression.getArraySource().accept(wrappingParser);
+  }
+
+  private List<?> getFilterValues(final ArrayFilterExpression arrayFilterExpression) {
+    final FilterTypeExpression filter = arrayFilterExpression.getFilter();
+    if (!(filter instanceof RelationalExpression)) {
+      throw new UnsupportedOperationException(
+          "Array operator "
+              + arrayFilterExpression.getOperator()
+              + " only supports a relational filter with a constant list of values, got: "
+              + filter);
+    }
+
+    final SelectTypeExpression rhs = ((RelationalExpression) filter).getRhs();
+    if (!(rhs instanceof ConstantExpression)) {
+      throw new UnsupportedOperationException(
+          "Array operator "
+              + arrayFilterExpression.getOperator()
+              + " requires a constant list of values, got: "
+              + rhs);
+    }
+
+    final Object value = ((ConstantExpression) rhs).getValue();
+    return value instanceof List ? (List<?>) value : List.of(value);
+  }
+
+  private Map<String, Object> wrapInExprIfNeeded(final Map<String, Object> filter) {
     // If already wrapped inside `$expr` avoid wrapping again
-    return INSIDE_EXPR.equals(relationalFilterContext.location())
-        ? arrayFilter
-        : Map.of(EXPR, arrayFilter);
+    return INSIDE_EXPR.equals(relationalFilterContext.location()) ? filter : Map.of(EXPR, filter);
   }
 }