Array filtering support (Part #4): Postgres array filter impl.

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

@@ -45,7 +45,7 @@
 import org.testcontainers.shaded.com.google.common.collect.Maps;
 import org.testcontainers.utility.DockerImageName;
 
-class ArrayFiltersQueryTest {
+class ArrayFiltersQueryIntegrationTest {
   private static final String COLLECTION_NAME = "galaxy";
 
   private static Map<String, Datastore> datastoreMap;
@@ -123,9 +123,9 @@ private static void createCollectionData() throws IOException {
         Utils.buildDocumentsFromResource("query/array_operators/galaxy.json");
     datastoreMap.forEach(
         (k, v) -> {
-          v.deleteCollection(ArrayFiltersQueryTest.COLLECTION_NAME);
-          v.createCollection(ArrayFiltersQueryTest.COLLECTION_NAME, null);
-          Collection collection = v.getCollection(ArrayFiltersQueryTest.COLLECTION_NAME);
+          v.deleteCollection(ArrayFiltersQueryIntegrationTest.COLLECTION_NAME);
+          v.createCollection(ArrayFiltersQueryIntegrationTest.COLLECTION_NAME, null);
+          Collection collection = v.getCollection(ArrayFiltersQueryIntegrationTest.COLLECTION_NAME);
           collection.bulkUpsert(documents);
         });
   }
@@ -143,13 +143,15 @@ public Stream<Arguments> provideArguments(final ExtensionContext context) {
     }
   }
 
+  @SuppressWarnings("unused")
   private static class MongoProvider implements ArgumentsProvider {
     @Override
     public Stream<Arguments> provideArguments(final ExtensionContext context) {
       return Stream.of(Arguments.of(MONGO_STORE));
     }
   }
 
+  @SuppressWarnings("unused")
   private static class PostgresProvider implements ArgumentsProvider {
     @Override
     public Stream<Arguments> provideArguments(final ExtensionContext context) {
@@ -158,17 +160,16 @@ public Stream<Arguments> provideArguments(final ExtensionContext context) {
   }
 
   @ParameterizedTest
-  @ArgumentsSource(MongoProvider.class)
+  @ArgumentsSource(AllProvider.class)
   void getAllSolarSystemsWithAtLeastOnePlanetHavingBothWaterAndOxygen(final String dataStoreName)
-      throws IOException, JSONException {
+      throws JSONException {
     final Collection collection = getCollection(dataStoreName);
 
     final Query query =
         Query.builder()
             .setFilter(
                 DocumentArrayFilterExpression.builder()
                     .operator(ANY)
-                    // # Can pass in some alias to this?
                     .arraySource(IdentifierExpression.of("additional_info.planets"))
                     .filter(
                         and(
@@ -187,7 +188,11 @@ void getAllSolarSystemsWithAtLeastOnePlanetHavingBothWaterAndOxygen(final String
                                         IdentifierExpression.of("elements"),
                                         RelationalOperator.EQ,
                                         ConstantExpression.of("Water")))
-                                .build()))
+                                .build(),
+                            RelationalExpression.of(
+                                IdentifierExpression.of("meta.num_moons"),
+                                RelationalOperator.GT,
+                                ConstantExpression.of(0))))
                     .build())
             .build();
 
@@ -209,7 +214,7 @@ private String readResource(final String fileName) {
 
   private Collection getCollection(final String dataStoreName) {
     final Datastore datastore = datastoreMap.get(dataStoreName);
-    return datastore.getCollection(ArrayFiltersQueryTest.COLLECTION_NAME);
+    return datastore.getCollection(ArrayFiltersQueryIntegrationTest.COLLECTION_NAME);
   }
 
   private String iteratorToJson(final Iterator<Document> iterator) {

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

@@ -9,15 +9,21 @@
             "Oxygen",
             "Water",
             "Nitrogen"
-          ]
+          ],
+          "meta": {
+            "num_moons": 7
+          }
         },
         {
           "name": "Planet 2",
           "elements": [
             "Oxygen",
             "Helium",
             "Water"
-          ]
+          ],
+          "meta": {
+            "num_moons": 1
+          }
         }
       ]
     }
@@ -46,7 +52,10 @@
             "Oxygen",
             "Nitrogen",
             "Water"
-          ]
+          ],
+          "meta": {
+            "num_moons": 1
+          }
         },
         {
           "name": "Mars",

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

@@ -10,15 +10,21 @@
             "Oxygen",
             "Water",
             "Nitrogen"
-          ]
+          ],
+          "meta": {
+            "num_moons": 7
+          }
         },
         {
           "name": "Planet 2",
           "elements": [
             "Oxygen",
             "Helium",
             "Water"
-          ]
+          ],
+          "meta": {
+            "num_moons": 1
+          }
         }
       ]
     }
@@ -61,7 +67,10 @@
             "Hydrogen",
             "Helium",
             "Methane"
-          ]
+          ],
+          "meta": {
+            "num_moons": 10
+          }
         },
         {
           "name": "Planet 8",
@@ -100,7 +109,10 @@
             "Oxygen",
             "Nitrogen",
             "Water"
-          ]
+          ],
+          "meta": {
+            "num_moons": 1
+          }
         },
         {
           "name": "Mars",

document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresArrayRelationalWrappingFilterVisitorProvider.java

@@ -0,0 +1,33 @@
+package org.hypertrace.core.documentstore.postgres.query.v1.vistors;
+
+import lombok.AllArgsConstructor;
+import org.hypertrace.core.documentstore.expression.type.SelectTypeExpression;
+import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser;
+
+@AllArgsConstructor
+public class PostgresArrayRelationalWrappingFilterVisitorProvider
+    implements PostgresWrappingFilterVisitorProvider {
+
+  private PostgresQueryParser postgresQueryParser;
+  private String arraySource;
+  private String alias;
+
+  @Override
+  public PostgresSelectTypeExpressionVisitor getForRelational(
+      final PostgresSelectTypeExpressionVisitor baseVisitor, final SelectTypeExpression rhs) {
+    // Override the base visitor,
+    // pick the LHS field name (elements.inner),
+    // replace it with alias (inner), and
+    // optionally trim double quotes (TRIM('"' FROM inner::text))
+    return new PostgresIdentifierTrimmingExpressionVisitor(
+        new PostgresIdentifierReplacingExpressionVisitor(
+            new PostgresIdentifierExpressionVisitor(postgresQueryParser), arraySource, alias),
+        rhs);
+  }
+
+  @Override
+  public PostgresSelectTypeExpressionVisitor getForNonRelational(
+      PostgresSelectTypeExpressionVisitor baseVisitor) {
+    return getForRelational(baseVisitor, null);
+  }
+}

document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresDocumentArrayWrappingFilterVisitorProvider.java

@@ -0,0 +1,31 @@
+package org.hypertrace.core.documentstore.postgres.query.v1.vistors;
+
+import lombok.AllArgsConstructor;
+import org.hypertrace.core.documentstore.expression.type.SelectTypeExpression;
+import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser;
+
+@AllArgsConstructor
+public class PostgresDocumentArrayWrappingFilterVisitorProvider
+    implements PostgresWrappingFilterVisitorProvider {
+
+  private PostgresQueryParser postgresQueryParser;
+  private String alias;
+
+  @Override
+  public PostgresSelectTypeExpressionVisitor getForRelational(
+      final PostgresSelectTypeExpressionVisitor baseVisitor, final SelectTypeExpression rhs) {
+    // Override the existing parser,
+    // parse the field name (meta.num_moons),
+    // get data accessor expression prefixed with alias (planets->'meta'->>'num_moons'), and
+    // cast according to the value type (CAST (planets->'meta'->>'num_moons' AS NUMERIC) > ?)
+    return new PostgresRelationalIdentifierAccessingExpressionVisitor(
+        new PostgresIdentifierExpressionVisitor(postgresQueryParser), alias, rhs);
+  }
+
+  @Override
+  public PostgresSelectTypeExpressionVisitor getForNonRelational(
+      final PostgresSelectTypeExpressionVisitor baseVisitor) {
+    // Any LHS field name (elements) is to be prefixed with current alias (inner)
+    return new PostgresIdentifierAccessingExpressionVisitor(baseVisitor, alias);
+  }
+}

document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFilterTypeExpressionVisitor.java

@@ -4,11 +4,13 @@
 import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.EQ;
 import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.IN;
 import static org.hypertrace.core.documentstore.postgres.PostgresCollection.ID;
+import static org.hypertrace.core.documentstore.postgres.utils.PostgresUtils.encodeAliasForNestedField;
 import static org.hypertrace.core.documentstore.postgres.utils.PostgresUtils.prepareParsedNonCompositeFilter;
 
 import java.util.Optional;
 import java.util.stream.Collector;
 import java.util.stream.Collectors;
+import javax.annotation.Nullable;
 import org.apache.commons.lang3.StringUtils;
 import org.hypertrace.core.documentstore.Key;
 import org.hypertrace.core.documentstore.expression.impl.ArrayRelationalFilterExpression;
@@ -27,9 +29,17 @@
 
 public class PostgresFilterTypeExpressionVisitor implements FilterTypeExpressionVisitor {
   protected PostgresQueryParser postgresQueryParser;
+  @Nullable private final PostgresWrappingFilterVisitorProvider wrappingVisitorProvider;
 
   public PostgresFilterTypeExpressionVisitor(PostgresQueryParser postgresQueryParser) {
+    this(postgresQueryParser, null);
+  }
+
+  public PostgresFilterTypeExpressionVisitor(
+      PostgresQueryParser postgresQueryParser,
+      final PostgresWrappingFilterVisitorProvider wrappingVisitorProvider) {
     this.postgresQueryParser = postgresQueryParser;
+    this.wrappingVisitorProvider = wrappingVisitorProvider;
   }
 
   @SuppressWarnings("unchecked")
@@ -51,7 +61,11 @@
   public String visit(final RelationalExpression expression) {
     final PostgresSelectExpressionParserBuilder parserBuilder =
         new PostgresSelectExpressionParserBuilderImpl(postgresQueryParser);
-    final PostgresSelectTypeExpressionVisitor lhsVisitor = parserBuilder.build(expression);
+    final PostgresSelectTypeExpressionVisitor baseVisitor = parserBuilder.build(expression);
+    final PostgresSelectTypeExpressionVisitor lhsVisitor =
+        Optional.ofNullable(wrappingVisitorProvider)
+            .map(visitor -> visitor.getForRelational(baseVisitor, expression.getRhs()))
+            .orElse(baseVisitor);
 
     final PostgresRelationalFilterContext context =
         PostgresRelationalFilterContext.builder()
@@ -80,16 +94,44 @@
             postgresQueryParser.getParamsBuilder());
   }
 
-  @SuppressWarnings("unchecked")
+  @SuppressWarnings({"unchecked", "SwitchStatementWithTooFewBranches"})
   @Override
   public String visit(final ArrayRelationalFilterExpression expression) {
-    throw new UnsupportedOperationException();
+    /*
+    EXISTS
+     (SELECT 1
+      FROM jsonb_array_elements(COALESCE(planets->'elements', '[]'::jsonb)) AS elements
+      WHERE TRIM('"' FROM elements::text) = 'Oxygen'
+     )
+     */
+    switch (expression.getOperator()) {
+      case ANY:
+        return getFilterStringForAnyOperator(expression);
+
+      default:
+        throw new UnsupportedOperationException(
+            "Unsupported array operator: " + expression.getOperator());
+    }
   }
 
-  @SuppressWarnings("unchecked")
+  @SuppressWarnings({"unchecked", "SwitchStatementWithTooFewBranches"})
   @Override
   public String visit(final DocumentArrayFilterExpression expression) {
-    throw new UnsupportedOperationException();
+    /*
+    EXISTS
+    (SELECT 1
+     FROM  jsonb_array_elements(COALESCE(document->'planets', '[]'::jsonb)) AS planets
+     WHERE <parsed_containing_filter_with_aliased_field_names>
+     )
+     */
+    switch (expression.getOperator()) {
+      case ANY:
+        return getFilterStringForAnyOperator(expression);
+
+      default:
+        throw new UnsupportedOperationException(
+            "Unsupported array operator: " + expression.getOperator());
+    }
   }
 
   public static Optional<String> getFilterClause(PostgresQueryParser postgresQueryParser) {
@@ -113,4 +155,75 @@
     throw new UnsupportedOperationException(
         String.format("Query operation:%s not supported", operator));
   }
+
+  private String getFilterStringForAnyOperator(final ArrayRelationalFilterExpression expression) {
+    // Convert 'elements' to planets->'elements' where planets could be an alias for an upper
+    // level array filter
+    // For the first time (if 'elements' was not under any nested array, say a top-level field),
+    // use the field identifier visitor to make it document->'elements'
+    final PostgresIdentifierExpressionVisitor identifierVisitor =
+        new PostgresIdentifierExpressionVisitor(postgresQueryParser);
+    final PostgresSelectTypeExpressionVisitor arrayPathVisitor =
+        wrappingVisitorProvider == null
+            ? new PostgresFieldIdentifierExpressionVisitor(identifierVisitor)
+            : wrappingVisitorProvider.getForNonRelational(identifierVisitor);
+    final String parsedLhs = expression.getArraySource().accept(arrayPathVisitor);
+
+    // Extract the field name
+    final String identifierName =
+        expression
+            .getArraySource()
+            .accept(new PostgresIdentifierExpressionVisitor(postgresQueryParser));
+
+    // If the field name is 'elements.inner', alias becomes 'elements_dot_inner'
+    final String alias = encodeAliasForNestedField(identifierName);
+
+    // Any LHS field name (elements) is to be prefixed with current alias (elements_dot_inner)
+    final PostgresWrappingFilterVisitorProvider visitorProvider =
+        new PostgresArrayRelationalWrappingFilterVisitorProvider(
+            postgresQueryParser, identifierName, alias);
+    final String parsedFilter =
+        expression
+            .getFilter()
+            .accept(new PostgresFilterTypeExpressionVisitor(postgresQueryParser, visitorProvider));
+
+    return String.format(
+        "EXISTS (SELECT 1 FROM jsonb_array_elements(COALESCE(%s, '[]'::jsonb)) AS \"%s\" WHERE %s)",
+        parsedLhs, alias, parsedFilter);
+  }
+
+  private String getFilterStringForAnyOperator(final DocumentArrayFilterExpression expression) {
+    // Convert 'elements' to planets->'elements' where planets could be an alias for an upper
+    // level array filter
+    // For the first time (if 'elements' was not under any nested array, say a top-level field),
+    // use the field identifier visitor to make it document->'elements'
+    final PostgresIdentifierExpressionVisitor identifierVisitor =
+        new PostgresIdentifierExpressionVisitor(postgresQueryParser);
+    final PostgresSelectTypeExpressionVisitor arrayPathVisitor =
+        wrappingVisitorProvider == null
+            ? new PostgresFieldIdentifierExpressionVisitor(identifierVisitor)
+            : wrappingVisitorProvider.getForNonRelational(identifierVisitor);
+    final String parsedLhs = expression.getArraySource().accept(arrayPathVisitor);
+
+    // Extract the field name
+    final String identifierName =
+        expression
+            .getArraySource()
+            .accept(new PostgresIdentifierExpressionVisitor(postgresQueryParser));
+
+    // If the field name is 'elements.inner', alias becomes 'elements_dot_inner'
+    final String alias = encodeAliasForNestedField(identifierName);
+
+    // Any LHS field name (elements) is to be prefixed with current alias (elements_dot_inner)
+    final PostgresWrappingFilterVisitorProvider wrapper =
+        new PostgresDocumentArrayWrappingFilterVisitorProvider(postgresQueryParser, alias);
+    final String parsedFilter =
+        expression
+            .getFilter()
+            .accept(new PostgresFilterTypeExpressionVisitor(postgresQueryParser, wrapper));
+
+    return String.format(
+        "EXISTS (SELECT 1 FROM jsonb_array_elements(COALESCE(%s, '[]'::jsonb)) AS \"%s\" WHERE %s)",
+        parsedLhs, alias, parsedFilter);
+  }
 }