Merge branch 'main' of github.com:hypertrace/document-store into handle_array_exists_flat_collections

Author: suddendust

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

@@ -1,6 +1,8 @@
 package org.hypertrace.core.documentstore.expression.impl;
 
+import java.util.Optional;
 import lombok.EqualsAndHashCode;
+import org.hypertrace.core.documentstore.parser.SelectTypeExpressionVisitor;
 
 /**
  * Represents an identifier expression for array-typed fields. This allows parsers to apply
@@ -12,11 +14,36 @@
 @EqualsAndHashCode(callSuper = true)
 public class ArrayIdentifierExpression extends IdentifierExpression {
 
+  private final ArrayType arrayType;
+
   public ArrayIdentifierExpression(String name) {
+    this(name, null);
+  }
+
+  public ArrayIdentifierExpression(String name, ArrayType arrayType) {
     super(name);
+    this.arrayType = arrayType;
   }
 
   public static ArrayIdentifierExpression of(String name) {
     return new ArrayIdentifierExpression(name);
   }
+
+  public static ArrayIdentifierExpression of(String name, ArrayType arrayType) {
+    return new ArrayIdentifierExpression(name, arrayType);
+  }
+
+  /** Returns the array type if specified, empty otherwise */
+  public Optional<ArrayType> getArrayType() {
+    return Optional.ofNullable(arrayType);
+  }
+
+  /**
+   * Accepts a SelectTypeExpressionVisitor and dispatches to the ArrayIdentifierExpression-specific
+   * visit method.
+   */
+  @Override
+  public <T> T accept(final SelectTypeExpressionVisitor visitor) {
+    return visitor.visit(this);
+  }
 }

document-store/src/main/java/org/hypertrace/core/documentstore/expression/impl/ArrayIdentifierExpression.java

@@ -0,0 +1,16 @@
+package org.hypertrace.core.documentstore.expression.impl;
+
+import lombok.Getter;
+
+public enum ArrayType {
+  TEXT("text[]"),
+  INTEGER("integer[]"),
+  BOOLEAN("boolean[]"),
+  DOUBLE_PRECISION("double precision[]");
+
+  @Getter private final String postgresType;
+
+  ArrayType(String postgresType) {
+    this.postgresType = postgresType;
+  }
+}

document-store/src/main/java/org/hypertrace/core/documentstore/expression/impl/ArrayType.java

@@ -0,0 +1,13 @@
+package org.hypertrace.core.documentstore.expression.impl;
+
+/** Represents the type of JSON fields in flat collections */
+public enum JsonFieldType {
+  STRING,
+  NUMBER,
+  BOOLEAN,
+  STRING_ARRAY,
+  NUMBER_ARRAY,
+  BOOLEAN_ARRAY,
+  OBJECT_ARRAY,
+  OBJECT
+}

document-store/src/main/java/org/hypertrace/core/documentstore/expression/impl/JsonArrayIdentifierExpression.java

@@ -1,9 +1,11 @@
 package org.hypertrace.core.documentstore.expression.impl;
 
 import java.util.List;
+import java.util.Optional;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
 import org.hypertrace.core.documentstore.parser.FieldTransformationVisitor;
+import org.hypertrace.core.documentstore.parser.SelectTypeExpressionVisitor;
 import org.hypertrace.core.documentstore.postgres.utils.BasicPostgresSecurityValidator;
 
 /**
@@ -19,6 +21,7 @@ public class JsonIdentifierExpression extends IdentifierExpression {
 
   String columnName; // e.g., "customAttr" (the top-level JSONB column)
   List<String> jsonPath; // e.g., ["myAttribute", "nestedField"]
+  JsonFieldType fieldType; // Optional: PRIMITIVE or ARRAY for optimization
 
   public static JsonIdentifierExpression of(final String columnName) {
     throw new IllegalArgumentException(
@@ -33,7 +36,20 @@ public static JsonIdentifierExpression of(final String columnName, final String.
     return of(columnName, List.of(pathElements));
   }
 
+  public static JsonIdentifierExpression of(
+      final String columnName, final JsonFieldType fieldType, final String... pathElements) {
+    if (pathElements == null || pathElements.length == 0) {
+      throw new IllegalArgumentException("JSON path cannot be null or empty");
+    }
+    return of(columnName, fieldType, List.of(pathElements));
+  }
+
   public static JsonIdentifierExpression of(final String columnName, final List<String> jsonPath) {
+    return of(columnName, null, jsonPath);
+  }
+
+  public static JsonIdentifierExpression of(
+      final String columnName, final JsonFieldType fieldType, final List<String> jsonPath) {
     BasicPostgresSecurityValidator.getDefault().validateIdentifier(columnName);
 
     if (jsonPath == null || jsonPath.isEmpty()) {
@@ -46,13 +62,20 @@ public static JsonIdentifierExpression of(final String columnName, final List<St
 
     // Construct full name for compatibility: "customAttr.myAttribute"
     String fullName = columnName + "." + String.join(".", unmodifiablePath);
-    return new JsonIdentifierExpression(fullName, columnName, unmodifiablePath);
+    return new JsonIdentifierExpression(fullName, columnName, unmodifiablePath, fieldType);
   }
 
-  protected JsonIdentifierExpression(String name, String columnName, List<String> jsonPath) {
+  protected JsonIdentifierExpression(
+      String name, String columnName, List<String> jsonPath, JsonFieldType fieldType) {
     super(name);
     this.columnName = columnName;
     this.jsonPath = jsonPath;
+    this.fieldType = fieldType;
+  }
+
+  /** Returns the JSON field type if specified, empty otherwise */
+  public Optional<JsonFieldType> getFieldType() {
+    return Optional.ofNullable(fieldType);
   }
 
   /**
@@ -68,6 +91,15 @@ public <T> T accept(final FieldTransformationVisitor<T> visitor) {
     return visitor.visit(this);
   }
 
+  /**
+   * Accepts a SelectTypeExpressionVisitor and dispatches to the JsonIdentifierExpression-specific
+   * visit method.
+   */
+  @Override
+  public <T> T accept(final SelectTypeExpressionVisitor visitor) {
+    return visitor.visit(this);
+  }
+
   @Override
   public String toString() {
     return String.format(

document-store/src/main/java/org/hypertrace/core/documentstore/expression/impl/JsonFieldType.java

@@ -2,10 +2,12 @@
 
 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;
 
 public interface SelectTypeExpressionVisitor {
   <T> T visit(final AggregateExpression expression);
@@ -19,4 +21,20 @@ public interface SelectTypeExpressionVisitor {
   <T> T visit(final IdentifierExpression expression);
 
   <T> T visit(final AliasedIdentifierExpression expression);
+
+  /**
+   * Visit an ArrayIdentifierExpression. Default implementation delegates to
+   * visit(IdentifierExpression) since ArrayIdentifierExpression extends IdentifierExpression.
+   */
+  default <T> T visit(final ArrayIdentifierExpression expression) {
+    return visit((IdentifierExpression) expression);
+  }
+
+  /**
+   * Visit a JsonIdentifierExpression. Default implementation delegates to
+   * visit(IdentifierExpression) since JsonIdentifierExpression extends IdentifierExpression.
+   */
+  default <T> T visit(final JsonIdentifierExpression expression) {
+    return visit((IdentifierExpression) expression);
+  }
 }

document-store/src/main/java/org/hypertrace/core/documentstore/expression/impl/JsonIdentifierExpression.java

@@ -0,0 +1,66 @@
+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;
+import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresContainsRelationalFilterParserNonJsonField;
+
+class PostgresContainsParserSelector implements SelectTypeExpressionVisitor {
+
+  private static final PostgresContainsRelationalFilterParserInterface jsonFieldContainsParser =
+      new PostgresContainsRelationalFilterParser();
+  private static final PostgresContainsRelationalFilterParserInterface nonJsonFieldContainsParser =
+      new PostgresContainsRelationalFilterParserNonJsonField();
+
+  private final boolean isFlatCollection;
+
+  PostgresContainsParserSelector(boolean isFlatCollection) {
+    this.isFlatCollection = isFlatCollection;
+  }
+
+  @Override
+  public PostgresRelationalFilterParser visit(JsonIdentifierExpression expression) {
+    return jsonFieldContainsParser;
+  }
+
+  @Override
+  public PostgresRelationalFilterParser visit(ArrayIdentifierExpression expression) {
+    return isFlatCollection ? nonJsonFieldContainsParser : jsonFieldContainsParser;
+  }
+
+  @Override
+  public PostgresRelationalFilterParser visit(IdentifierExpression expression) {
+    return isFlatCollection ? nonJsonFieldContainsParser : jsonFieldContainsParser;
+  }
+
+  @Override
+  public PostgresRelationalFilterParser visit(AggregateExpression expression) {
+    return isFlatCollection ? nonJsonFieldContainsParser : jsonFieldContainsParser;
+  }
+
+  @Override
+  public PostgresRelationalFilterParser visit(ConstantExpression expression) {
+    return isFlatCollection ? nonJsonFieldContainsParser : jsonFieldContainsParser;
+  }
+
+  @Override
+  public PostgresRelationalFilterParser visit(DocumentConstantExpression expression) {
+    return isFlatCollection ? nonJsonFieldContainsParser : jsonFieldContainsParser;
+  }
+
+  @Override
+  public PostgresRelationalFilterParser visit(FunctionExpression expression) {
+    return isFlatCollection ? nonJsonFieldContainsParser : jsonFieldContainsParser;
+  }
+
+  @Override
+  public PostgresRelationalFilterParser visit(AliasedIdentifierExpression expression) {
+    return isFlatCollection ? nonJsonFieldContainsParser : jsonFieldContainsParser;
+  }
+}