ArrayIdentifierExpression and JsonArrayIdentifierExpression to pass type info to flat collections (#249)

* ArrayIdentifierExpression and JsonArrayIdentifierExpression to pass type info to flat collections

* Added test cases

* Fixed failing test cases

* Fixed failing test cases

Author: suddendust

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

@@ -0,0 +1,22 @@
+package org.hypertrace.core.documentstore.expression.impl;
+
+import lombok.EqualsAndHashCode;
+
+/**
+ * Represents an identifier expression for array-typed fields. This allows parsers to apply
+ * array-specific logic (e.g., cardinality checks for EXISTS operators to exclude empty arrays).
+ *
+ * <p>Similar to {@link JsonIdentifierExpression}, this provides type information to parsers so they
+ * can generate appropriate database-specific queries for array operations.
+ */
+@EqualsAndHashCode(callSuper = true)
+public class ArrayIdentifierExpression extends IdentifierExpression {
+
+  public ArrayIdentifierExpression(String name) {
+    super(name);
+  }
+
+  public static ArrayIdentifierExpression of(String name) {
+    return new ArrayIdentifierExpression(name);
+  }
+}

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

@@ -0,0 +1,50 @@
+package org.hypertrace.core.documentstore.expression.impl;
+
+import java.util.List;
+import lombok.EqualsAndHashCode;
+import org.hypertrace.core.documentstore.postgres.utils.BasicPostgresSecurityValidator;
+
+/**
+ * Represents an identifier expression for array-typed fields inside JSONB columns. This allows
+ * parsers to apply array-specific logic (e.g., jsonb_array_length checks for EXISTS operators to
+ * exclude empty arrays).
+ *
+ * <p>Example: For a JSONB column "attributes" with a nested array field "tags":
+ *
+ * <pre>{"attributes": {"tags": ["value1", "value2"]}}</pre>
+ *
+ * Use: {@code JsonArrayIdentifierExpression.of("attributes", "tags")}
+ */
+@EqualsAndHashCode(callSuper = true)
+public class JsonArrayIdentifierExpression extends JsonIdentifierExpression {
+
+  public static JsonArrayIdentifierExpression of(
+      final String columnName, final String... pathElements) {
+    if (pathElements == null || pathElements.length == 0) {
+      throw new IllegalArgumentException("JSON path cannot be null or empty for array field");
+    }
+    return of(columnName, List.of(pathElements));
+  }
+
+  public static JsonArrayIdentifierExpression of(
+      final String columnName, final List<String> jsonPath) {
+    // Validate inputs
+    BasicPostgresSecurityValidator.getDefault().validateIdentifier(columnName);
+
+    if (jsonPath == null || jsonPath.isEmpty()) {
+      throw new IllegalArgumentException("JSON path cannot be null or empty for array field");
+    }
+
+    BasicPostgresSecurityValidator.getDefault().validateJsonPath(jsonPath);
+
+    List<String> unmodifiablePath = List.copyOf(jsonPath);
+
+    // Construct full name for compatibility: "customAttr.myAttribute"
+    String fullName = columnName + "." + String.join(".", unmodifiablePath);
+    return new JsonArrayIdentifierExpression(fullName, columnName, unmodifiablePath);
+  }
+
+  private JsonArrayIdentifierExpression(String name, String columnName, List<String> jsonPath) {
+    super(name, columnName, jsonPath);
+  }
+}

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

@@ -2,7 +2,7 @@
 
 import java.util.List;
 import lombok.EqualsAndHashCode;
-import lombok.Value;
+import lombok.Getter;
 import org.hypertrace.core.documentstore.parser.FieldTransformationVisitor;
 import org.hypertrace.core.documentstore.postgres.utils.BasicPostgresSecurityValidator;
 
@@ -13,13 +13,18 @@
  *
  * <p>This generates SQL like: customAttr -> 'myAttribute' -> 'nestedField' (returns JSON)
  */
-@Value
+@Getter
 @EqualsAndHashCode(callSuper = true)
 public class JsonIdentifierExpression extends IdentifierExpression {
 
   String columnName; // e.g., "customAttr" (the top-level JSONB column)
   List<String> jsonPath; // e.g., ["myAttribute", "nestedField"]
 
+  public static JsonIdentifierExpression of(final String columnName) {
+    throw new IllegalArgumentException(
+        "JSON path cannot be null or empty. Use of(columnName, path...) instead.");
+  }
+
   public static JsonIdentifierExpression of(final String columnName, final String... pathElements) {
     if (pathElements == null || pathElements.length == 0) {
       // In this case, use IdentifierExpression
@@ -44,7 +49,7 @@ public static JsonIdentifierExpression of(final String columnName, final List<St
     return new JsonIdentifierExpression(fullName, columnName, unmodifiablePath);
   }
 
-  private JsonIdentifierExpression(String name, String columnName, List<String> jsonPath) {
+  protected JsonIdentifierExpression(String name, String columnName, List<String> jsonPath) {
     super(name);
     this.columnName = columnName;
     this.jsonPath = jsonPath;

document-store/src/test/java/org/hypertrace/core/documentstore/expression/impl/ArrayIdentifierExpressionTest.java

@@ -0,0 +1,72 @@
+package org.hypertrace.core.documentstore.expression.impl;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+
+import org.junit.jupiter.api.Test;
+
+class ArrayIdentifierExpressionTest {
+
+  @Test
+  void testOfCreatesInstance() {
+    ArrayIdentifierExpression expression = ArrayIdentifierExpression.of("tags");
+
+    assertEquals("tags", expression.getName());
+  }
+
+  @Test
+  void testEqualsAndHashCode() {
+    ArrayIdentifierExpression expr1 = ArrayIdentifierExpression.of("tags");
+    ArrayIdentifierExpression expr2 = ArrayIdentifierExpression.of("tags");
+
+    // Test equality - should be equal
+    assertEquals(expr1, expr2, "Expressions with same name should be equal");
+
+    // Test hashCode
+    assertEquals(
+        expr1.hashCode(), expr2.hashCode(), "Expressions with same name should have same hashCode");
+  }
+
+  @Test
+  void testNotEqualsWithDifferentName() {
+    ArrayIdentifierExpression expr1 = ArrayIdentifierExpression.of("tags");
+    ArrayIdentifierExpression expr2 = ArrayIdentifierExpression.of("categories");
+
+    // Test inequality
+    assertNotEquals(expr1, expr2, "Expressions with different names should not be equal");
+  }
+
+  @Test
+  void testNotEqualsWithIdentifierExpression() {
+    ArrayIdentifierExpression arrayExpr = ArrayIdentifierExpression.of("tags");
+    IdentifierExpression identExpr = IdentifierExpression.of("tags");
+
+    // Even though they have the same name, they are different types
+    assertNotEquals(
+        arrayExpr, identExpr, "ArrayIdentifierExpression should not equal IdentifierExpression");
+  }
+
+  @Test
+  void testInheritsFromIdentifierExpression() {
+    ArrayIdentifierExpression expression = ArrayIdentifierExpression.of("tags");
+
+    // Verify it's an instance of parent class
+    assertEquals(
+        IdentifierExpression.class,
+        expression.getClass().getSuperclass(),
+        "ArrayIdentifierExpression should extend IdentifierExpression");
+  }
+
+  @Test
+  void testMultipleInstancesWithSameNameAreEqual() {
+    ArrayIdentifierExpression expr1 = ArrayIdentifierExpression.of("categoryTags");
+    ArrayIdentifierExpression expr2 = ArrayIdentifierExpression.of("categoryTags");
+    ArrayIdentifierExpression expr3 = ArrayIdentifierExpression.of("categoryTags");
+
+    assertEquals(expr1, expr2);
+    assertEquals(expr2, expr3);
+    assertEquals(expr1, expr3);
+    assertEquals(expr1.hashCode(), expr2.hashCode());
+    assertEquals(expr2.hashCode(), expr3.hashCode());
+  }
+}