Sync EQ/NEQ on Arrays in Flat Collections with Mongo

document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/builder/PostgresSelectExpressionParserBuilderImpl.java

@@ -1,22 +1,31 @@
 package org.hypertrace.core.documentstore.postgres.query.v1.parser.builder;
 
+import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.EQ;
+import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.NEQ;
 import static org.hypertrace.core.documentstore.postgres.utils.PostgresUtils.getType;
 
-import lombok.AllArgsConstructor;
+import org.hypertrace.core.documentstore.expression.impl.ArrayIdentifierExpression;
+import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
+import org.hypertrace.core.documentstore.expression.impl.JsonFieldType;
+import org.hypertrace.core.documentstore.expression.impl.JsonIdentifierExpression;
 import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
+import org.hypertrace.core.documentstore.expression.type.SelectTypeExpression;
 import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser;
 import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresConstantExpressionVisitor;
 import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresDataAccessorIdentifierExpressionVisitor;
 import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresFieldIdentifierExpressionVisitor;
 import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresFunctionExpressionVisitor;
 import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresSelectTypeExpressionVisitor;
 
-@AllArgsConstructor
 public class PostgresSelectExpressionParserBuilderImpl
     implements PostgresSelectExpressionParserBuilder {
 
   private final PostgresQueryParser postgresQueryParser;
 
+  public PostgresSelectExpressionParserBuilderImpl(PostgresQueryParser postgresQueryParser) {
+    this.postgresQueryParser = postgresQueryParser;
+  }
+
   @Override
   public PostgresSelectTypeExpressionVisitor build(final RelationalExpression expression) {
     switch (expression.getOperator()) {
@@ -29,11 +38,66 @@ public PostgresSelectTypeExpressionVisitor build(final RelationalExpression expr
         return new PostgresFunctionExpressionVisitor(
             new PostgresFieldIdentifierExpressionVisitor(this.postgresQueryParser));
 
+      case EQ:
+      case NEQ:
+        // For EQ/NEQ on array fields, treat like CONTAINS to use -> instead of ->>
+        if (shouldSwitchToContainsFlow(expression)) {
+          // Use field identifier (JSON accessor ->) for array fields
+          return new PostgresFunctionExpressionVisitor(
+              new PostgresFieldIdentifierExpressionVisitor(this.postgresQueryParser));
+        }
+      // Fall through to default for non-array fields
       default:
         return new PostgresFunctionExpressionVisitor(
             new PostgresDataAccessorIdentifierExpressionVisitor(
                 this.postgresQueryParser,
                 getType(expression.getRhs().accept(new PostgresConstantExpressionVisitor()))));
     }
   }
+
+  /**
+   * Checks if this is an EQ/NEQ operator on an array field.
+   *
+   * <p>Only converts to CONTAINS when RHS is a scalar value. If RHS is an array, we want exact
+   * equality match, not containment.
+   *
+   * <p>Handles both:
+   *
+   * <ul>
+   *   <li>{@link JsonIdentifierExpression} with array field type (JSONB arrays)
+   *   <li>{@link ArrayIdentifierExpression} with array type (top-level array columns)
+   * </ul>
+   */
+  private boolean shouldSwitchToContainsFlow(final RelationalExpression expression) {
+    if (expression.getOperator() != EQ && expression.getOperator() != NEQ) {
+      return false;
+    }
+
+    // Check if RHS is an array/iterable - if so, don't convert (since we want an exact match for
+    // such cases)
+    if (expression.getRhs() instanceof ConstantExpression) {
+      ConstantExpression constExpr = (ConstantExpression) expression.getRhs();
+      if (constExpr.getValue() instanceof Iterable) {
+        return false;
+      }
+    }
+
+    return isArrayField(expression.getLhs());
+  }
+
+  private boolean isArrayField(final SelectTypeExpression lhs) {
+    if (lhs instanceof JsonIdentifierExpression) {
+      JsonIdentifierExpression jsonExpr = (JsonIdentifierExpression) lhs;
+      return jsonExpr
+          .getFieldType()
+          .map(
+              fieldType ->
+                  fieldType == JsonFieldType.BOOLEAN_ARRAY
+                      || fieldType == JsonFieldType.STRING_ARRAY
+                      || fieldType == JsonFieldType.NUMBER_ARRAY
+                      || fieldType == JsonFieldType.OBJECT_ARRAY)
+          .orElse(false);
+    }
+    return lhs instanceof ArrayIdentifierExpression;
+  }
 }

document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresArrayEqualityParserSelector.java

@@ -0,0 +1,68 @@
+package org.hypertrace.core.documentstore.postgres.query.v1.parser.filter;
+
+import org.hypertrace.core.documentstore.expression.impl.AggregateExpression;
+import org.hypertrace.core.documentstore.expression.impl.AliasedIdentifierExpression;
+import org.hypertrace.core.documentstore.expression.impl.ArrayIdentifierExpression;
+import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
+import org.hypertrace.core.documentstore.expression.impl.ConstantExpression.DocumentConstantExpression;
+import org.hypertrace.core.documentstore.expression.impl.FunctionExpression;
+import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression;
+import org.hypertrace.core.documentstore.expression.impl.JsonIdentifierExpression;
+import org.hypertrace.core.documentstore.parser.SelectTypeExpressionVisitor;
+
+/**
+ * Selects the appropriate array equality parser based on the LHS expression type.
+ *
+ * <p>For JsonIdentifierExpression: uses JSONB array equality parser
+ *
+ * <p>For ArrayIdentifierExpression: uses top-level array equality parser
+ */
+class PostgresArrayEqualityParserSelector implements SelectTypeExpressionVisitor {
+
+  private static final PostgresRelationalFilterParser jsonArrayEqualityParser =
+      new PostgresJsonArrayEqualityFilterParser();
+  private static final PostgresRelationalFilterParser topLevelArrayEqualityParser =
+      new PostgresTopLevelArrayEqualityFilterParser();
+  private static final PostgresRelationalFilterParser standardParser =
+      new PostgresStandardRelationalFilterParser();
+
+  @Override
+  public PostgresRelationalFilterParser visit(JsonIdentifierExpression expression) {
+    return jsonArrayEqualityParser;
+  }
+
+  @Override
+  public PostgresRelationalFilterParser visit(ArrayIdentifierExpression expression) {
+    return topLevelArrayEqualityParser;
+  }
+
+  @Override
+  public <T> T visit(IdentifierExpression expression) {
+    return (T) standardParser;
+  }
+
+  @Override
+  public <T> T visit(AggregateExpression expression) {
+    return (T) standardParser;
+  }
+
+  @Override
+  public <T> T visit(ConstantExpression expression) {
+    return (T) standardParser;
+  }
+
+  @Override
+  public <T> T visit(DocumentConstantExpression expression) {
+    return (T) standardParser;
+  }
+
+  @Override
+  public <T> T visit(FunctionExpression expression) {
+    return (T) standardParser;
+  }
+
+  @Override
+  public <T> T visit(AliasedIdentifierExpression expression) {
+    return (T) standardParser;
+  }
+}

document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresJsonArrayEqualityFilterParser.java

@@ -0,0 +1,44 @@
+package org.hypertrace.core.documentstore.postgres.query.v1.parser.filter;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
+
+/**
+ * Handles EQ/NEQ operations on JSONB array fields when RHS is also an array, using exact equality
+ * (=) instead of containment (@>).
+ *
+ * <p>Generates: {@code props->'source-loc' = '["hygiene","family-pack"]'::jsonb}
+ */
+class PostgresJsonArrayEqualityFilterParser implements PostgresRelationalFilterParser {
+
+  private static final PostgresStandardRelationalOperatorMapper mapper =
+      new PostgresStandardRelationalOperatorMapper();
+  private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+  @Override
+  public String parse(
+      final RelationalExpression expression, final PostgresRelationalFilterContext context) {
+    final String parsedLhs = expression.getLhs().accept(context.lhsParser());
+    final Object parsedRhs = expression.getRhs().accept(context.rhsParser());
+    final String operator = mapper.getMapping(expression.getOperator(), parsedRhs);
+
+    if (parsedRhs == null) {
+      return String.format("%s %s NULL", parsedLhs, operator);
+    }
+
+    // Convert the array to a JSONB string representation
+    try {
+      String jsonbValue;
+      if (parsedRhs instanceof Iterable) {
+        jsonbValue = OBJECT_MAPPER.writeValueAsString(parsedRhs);
+      } else {
+        jsonbValue = String.valueOf(parsedRhs);
+      }
+      context.getParamsBuilder().addObjectParam(jsonbValue);
+      return String.format("%s %s ?::jsonb", parsedLhs, operator);
+    } catch (JsonProcessingException e) {
+      throw new RuntimeException("Failed to serialize RHS array to JSON", e);
+    }
+  }
+}

document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresNotContainsRelationalFilterParser.java

@@ -35,8 +35,7 @@ private boolean shouldUseJsonParser(
 
     boolean isJsonField = expression.getLhs() instanceof JsonIdentifierExpression;
     boolean isFlatCollection = context.getPgColTransformer().getDocumentType() == DocumentType.FLAT;
-    boolean useJsonParser = !isFlatCollection || isJsonField;
 
-    return useJsonParser;
+    return !isFlatCollection || isJsonField;
   }
 }

document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresRelationalFilterParserFactoryImpl.java

@@ -2,9 +2,11 @@
 
 import static java.util.Map.entry;
 import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.CONTAINS;
+import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.EQ;
 import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.EXISTS;
 import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.IN;
 import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.LIKE;
+import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.NEQ;
 import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.NOT_CONTAINS;
 import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.NOT_EXISTS;
 import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.NOT_IN;
@@ -13,12 +15,18 @@
 import com.google.common.collect.Maps;
 import java.util.Map;
 import org.hypertrace.core.documentstore.DocumentType;
+import org.hypertrace.core.documentstore.expression.impl.ArrayIdentifierExpression;
+import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
+import org.hypertrace.core.documentstore.expression.impl.JsonFieldType;
+import org.hypertrace.core.documentstore.expression.impl.JsonIdentifierExpression;
 import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
 import org.hypertrace.core.documentstore.expression.operators.RelationalOperator;
+import org.hypertrace.core.documentstore.expression.type.SelectTypeExpression;
 import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser;
 
 public class PostgresRelationalFilterParserFactoryImpl
     implements PostgresRelationalFilterParserFactory {
+
   private static final Map<RelationalOperator, PostgresRelationalFilterParser> parserMap =
       Maps.immutableEnumMap(
           Map.ofEntries(
@@ -41,12 +49,139 @@ public PostgresRelationalFilterParser parser(
     boolean isFlatCollection =
         postgresQueryParser.getPgColTransformer().getDocumentType() == DocumentType.FLAT;
 
-    if (expression.getOperator() == CONTAINS) {
+    RelationalOperator operator = expression.getOperator();
+    // Transform EQ/NEQ to CONTAINS/NOT_CONTAINS for array fields with scalar RHS
+    // (but not for unnested fields, which are already scalar)
+    if (shouldConvertEqToContains(expression, postgresQueryParser)) {
+      operator = (expression.getOperator() == EQ) ? CONTAINS : NOT_CONTAINS;
+    }
+
+    if (operator == CONTAINS) {
       return expression.getLhs().accept(new PostgresContainsParserSelector(isFlatCollection));
-    } else if (expression.getOperator() == IN) {
+    } else if (operator == IN) {
       return expression.getLhs().accept(new PostgresInParserSelector(isFlatCollection));
+    } else if (operator == NOT_CONTAINS) {
+      return parserMap.get(NOT_CONTAINS);
+    }
+
+    // For EQ/NEQ on array fields with array RHS, use specialized array equality parser (exact match
+    // instead of containment)
+    if (shouldUseArrayEqualityParser(expression, postgresQueryParser)) {
+      return expression.getLhs().accept(new PostgresArrayEqualityParserSelector());
     }
 
     return parserMap.getOrDefault(expression.getOperator(), postgresStandardRelationalFilterParser);
   }
+
+  /**
+   * Determines if EQ/NEQ should be converted to CONTAINS/NOT_CONTAINS.
+   *
+   * <p>Conversion happens when:
+   *
+   * <ul>
+   *   <li>Operator is EQ or NEQ
+   *   <li>RHS is a SCALAR value (not an array/iterable)
+   *   <li>LHS is a JsonIdentifierExpression with an array field type (STRING_ARRAY, NUMBER_ARRAY,
+   *       etc.) OR
+   *   <li>LHS is an ArrayIdentifierExpression with an array type (TEXT, BIGINT, etc.)
+   *   <li>Field has NOT been unnested (unnested fields are scalar, not arrays)
+   * </ul>
+   *
+   * <p>If RHS is an array, we DO NOT convert - we want exact equality match (= operator), not
+   * containment (@> operator).
+   *
+   * <p>This provides semantic equivalence: checking if an array contains a scalar value is more
+   * intuitive than checking if the array equals the value.
+   */
+  private boolean shouldConvertEqToContains(
+      final RelationalExpression expression, final PostgresQueryParser postgresQueryParser) {
+    if (expression.getOperator() != EQ && expression.getOperator() != NEQ) {
+      return false;
+    }
+
+    // Check if RHS is an array/iterable - if so, don't convert (we want exact match)
+    if (isArrayRhs(expression.getRhs())) {
+      return false;
+    }
+
+    // Check if LHS is an array field
+    if (!isArrayField(expression.getLhs())) {
+      return false;
+    }
+
+    // Check if field has been unnested - unnested fields are scalar, not arrays
+    String fieldName = getFieldName(expression.getLhs());
+    return fieldName == null
+        || !postgresQueryParser
+            .getPgColumnNames()
+            .containsKey(fieldName); // Field is unnested - treat as scalar
+  }
+
+  /**
+   * Determines if we should use the specialized array equality parser.
+   *
+   * <p>Use this parser when:
+   *
+   * <ul>
+   *   <li>Operator is EQ or NEQ
+   *   <li>RHS is an array/iterable (for exact match).
+   *   <li>LHS is either {@link JsonIdentifierExpression} with array type OR {@link
+   *       ArrayIdentifierExpression}
+   *   <li>Field has NOT been unnested (unnested fields are scalar, not arrays)
+   * </ul>
+   */
+  private boolean shouldUseArrayEqualityParser(
+      final RelationalExpression expression, final PostgresQueryParser postgresQueryParser) {
+    if (expression.getOperator() != EQ && expression.getOperator() != NEQ) {
+      return false;
+    }
+
+    // Check if RHS is an array/iterable AND LHS is an array field
+    if (!isArrayRhs(expression.getRhs()) || !isArrayField(expression.getLhs())) {
+      return false;
+    }
+
+    // Check if field has been unnested - unnested fields are scalar, not arrays
+    String fieldName = getFieldName(expression.getLhs());
+    return fieldName == null || !postgresQueryParser.getPgColumnNames().containsKey(fieldName);
+  }
+
+  /**
+   * Checks if the RHS expression contains an array/iterable value. Currently, we don't have a very
+   * clean way to get the RHS data type. //todo: Implement a clean way to get the RHS data type
+   */
+  private boolean isArrayRhs(final SelectTypeExpression rhs) {
+    if (rhs instanceof ConstantExpression) {
+      ConstantExpression constExpr = (ConstantExpression) rhs;
+      return constExpr.getValue() instanceof Iterable;
+    }
+    return false;
+  }
+
+  /** Checks if the LHS expression is an array field. */
+  private boolean isArrayField(final SelectTypeExpression lhs) {
+    if (lhs instanceof JsonIdentifierExpression) {
+      JsonIdentifierExpression jsonExpr = (JsonIdentifierExpression) lhs;
+      return jsonExpr
+          .getFieldType()
+          .map(
+              fieldType ->
+                  fieldType == JsonFieldType.BOOLEAN_ARRAY
+                      || fieldType == JsonFieldType.STRING_ARRAY
+                      || fieldType == JsonFieldType.NUMBER_ARRAY
+                      || fieldType == JsonFieldType.OBJECT_ARRAY)
+          .orElse(false);
+    }
+    return lhs instanceof ArrayIdentifierExpression;
+  }
+
+  /** Extracts the field name from an identifier expression. */
+  private String getFieldName(final SelectTypeExpression lhs) {
+    if (lhs instanceof JsonIdentifierExpression) {
+      return ((JsonIdentifierExpression) lhs).getName();
+    } else if (lhs instanceof ArrayIdentifierExpression) {
+      return ((ArrayIdentifierExpression) lhs).getName();
+    }
+    return null;
+  }
 }