Introduces abstract types for flat collections, address other comments

Author: suddendust

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

@@ -1,6 +1,5 @@
 package org.hypertrace.core.documentstore.expression.impl;
 
-import java.util.Optional;
 import lombok.EqualsAndHashCode;
 import org.hypertrace.core.documentstore.parser.SelectTypeExpressionVisitor;
 
@@ -14,132 +13,66 @@
 @EqualsAndHashCode(callSuper = true)
 public class ArrayIdentifierExpression extends IdentifierExpression {
 
-  private final FlatCollectionDataType arrayElementType;
+  private final DataType arrayElementType;
 
-  public ArrayIdentifierExpression(String name) {
-    this(name, null);
+  ArrayIdentifierExpression(String name) {
+    this(name, DataType.UNSPECIFIED);
   }
 
-  // Package-private: used internally by factory methods
-  ArrayIdentifierExpression(String name, FlatCollectionDataType arrayElementType) {
-    super(name, null);
+  ArrayIdentifierExpression(String name, DataType arrayElementType) {
+    super(name);
     this.arrayElementType = arrayElementType;
   }
 
   public static ArrayIdentifierExpression of(String name) {
     return new ArrayIdentifierExpression(name);
   }
 
-  // Package-private: used internally by factory methods
-  static ArrayIdentifierExpression of(String name, FlatCollectionDataType arrayElementType) {
+  static ArrayIdentifierExpression of(String name, DataType arrayElementType) {
     return new ArrayIdentifierExpression(name, arrayElementType);
   }
 
   public static ArrayIdentifierExpression ofStrings(final String name) {
-    return of(name, PostgresDataType.TEXT);
+    return of(name, DataType.STRING);
   }
 
   public static ArrayIdentifierExpression ofInts(final String name) {
-    return of(name, PostgresDataType.INTEGER);
+    return of(name, DataType.INTEGER);
   }
 
   public static ArrayIdentifierExpression ofLongs(final String name) {
-    return of(name, PostgresDataType.BIGINT);
-  }
-
-  public static ArrayIdentifierExpression ofShorts(final String name) {
-    return of(name, PostgresDataType.SMALLINT);
+    return of(name, DataType.LONG);
   }
 
   public static ArrayIdentifierExpression ofFloats(final String name) {
-    return of(name, PostgresDataType.FLOAT);
+    return of(name, DataType.FLOAT);
   }
 
   public static ArrayIdentifierExpression ofDoubles(final String name) {
-    return of(name, PostgresDataType.DOUBLE);
-  }
-
-  public static ArrayIdentifierExpression ofDecimals(final String name) {
-    return of(name, PostgresDataType.NUMERIC);
+    return of(name, DataType.DOUBLE);
   }
 
   public static ArrayIdentifierExpression ofBooleans(final String name) {
-    return of(name, PostgresDataType.BOOLEAN);
-  }
-
-  public static ArrayIdentifierExpression ofTimestamps(final String name) {
-    return of(name, PostgresDataType.TIMESTAMP);
+    return of(name, DataType.BOOLEAN);
   }
 
   public static ArrayIdentifierExpression ofTimestampsTz(final String name) {
-    return of(name, PostgresDataType.TIMESTAMPTZ);
+    return of(name, DataType.TIMESTAMPTZ);
   }
 
   public static ArrayIdentifierExpression ofDates(final String name) {
-    return of(name, PostgresDataType.DATE);
-  }
-
-  public static ArrayIdentifierExpression ofUuids(final String name) {
-    return of(name, PostgresDataType.UUID);
-  }
-
-  public static ArrayIdentifierExpression ofBytesArray(final String name) {
-    return of(name, PostgresDataType.BYTEA);
-  }
-
-  // Package-private: used internally by getPostgresArrayTypeString()
-  Optional<FlatCollectionDataType> getArrayElementType() {
-    return Optional.ofNullable(arrayElementType);
+    return of(name, DataType.DATE);
   }
 
   /**
-   * Returns the PostgreSQL array type string for this expression, if the element type is specified.
+   * Returns the data type of array elements.
    *
-   * <p>Examples:
+   * <p>This is used by database-specific type extractors to generate appropriate type casts.
    *
-   * <ul>
-   *   <li>TEXT → "text[]"
-   *   <li>INTEGER → "integer[]"
-   *   <li>BOOLEAN → "boolean[]"
-   *   <li>FLOAT → "real[]"
-   *   <li>DOUBLE → "double precision[]"
-   * </ul>
-   *
-   * @return Optional containing the PostgreSQL array type string, or empty if no type is specified
-   * @throws IllegalArgumentException if the element type is not a PostgresDataType
+   * @return The element DataType (UNSPECIFIED if no type was explicitly set)
    */
-  public Optional<String> getPostgresArrayTypeString() {
-    return getArrayElementType().map(this::toPostgresArrayType);
-  }
-
-  private String toPostgresArrayType(FlatCollectionDataType type) {
-    if (!(type instanceof PostgresDataType)) {
-      throw new IllegalArgumentException(
-          "Only PostgresDataType is currently supported, got: " + type.getClass());
-    }
-    PostgresDataType postgresType = (PostgresDataType) type;
-    String baseType = postgresType.getPgTypeName();
-
-    // Map internal type names to SQL array type names
-    switch (baseType) {
-      case "int4":
-        return "integer[]";
-      case "int8":
-        return "bigint[]";
-      case "int2":
-        return "smallint[]";
-      case "float4":
-        return "real[]";
-      case "float8":
-        return "double precision[]";
-      case "bool":
-        return "boolean[]";
-      case "bytea":
-        return "bytea[]";
-      default:
-        // For most types, just append []
-        return baseType + "[]";
-    }
+  public DataType getElementDataType() {
+    return arrayElementType;
   }
 
   /**

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

@@ -0,0 +1,27 @@
+package org.hypertrace.core.documentstore.expression.impl;
+
+/**
+ * Database-agnostic data types for explicit type annotation in queries.
+ *
+ * <p>This enum provides type metadata for {@link IdentifierExpression} and {@link
+ * ArrayIdentifierExpression} fields in flat collections, enabling type-safe query generation
+ * without runtime type inference.
+ *
+ * <p>These types are mapped to database-specific types at query parsing time. For example, when
+ * generating PostgreSQL queries, {@code STRING} maps to {@code text}, {@code INTEGER} maps to
+ * {@code int4}, etc.
+ *
+ * @see ArrayIdentifierExpression
+ * @see IdentifierExpression
+ */
+public enum DataType {
+  UNSPECIFIED,
+  STRING,
+  INTEGER,
+  LONG,
+  FLOAT,
+  DOUBLE,
+  BOOLEAN,
+  TIMESTAMPTZ,
+  DATE
+}

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

@@ -18,10 +18,10 @@
  *
  * <p>Example: IdentifierExpression.of("col1");
  *
- * <p>For flat relational collections, you can optionally provide a {@link FlatCollectionDataType}
- * to enable type-safe query generation without runtime type inference:
+ * <p>For flat relational collections, you can optionally provide a {@link DataType} to enable
+ * type-safe query generation without runtime type inference:
  *
- * <p>Example: IdentifierExpression.of("price", PostgresType.INTEGER);
+ * <p>Example: IdentifierExpression.ofInt("price");
  */
 @Value
 @NonFinal
@@ -30,75 +30,56 @@ public class IdentifierExpression
     implements GroupTypeExpression, SelectTypeExpression, SortTypeExpression {
 
   String name;
-  // Type information of this identifier for flat collections, optional
-  FlatCollectionDataType flatCollectionDataType;
+  // Type information of this identifier for flat collections, this is optional to maintain backward
+  // compatibility
+  DataType dataType;
+
+  IdentifierExpression(String name) {
+    this.name = name;
+    this.dataType = DataType.UNSPECIFIED;
+  }
 
   public static IdentifierExpression of(final String name) {
     Preconditions.checkArgument(name != null && !name.isBlank(), "name is null or blank");
-    return new IdentifierExpression(name, null);
+    return new IdentifierExpression(name);
   }
 
   // Package-private: used internally by factory methods
-  static IdentifierExpression of(
-      final String name, final FlatCollectionDataType flatCollectionDataType) {
+  static IdentifierExpression of(final String name, final DataType dataType) {
     Preconditions.checkArgument(name != null && !name.isBlank(), "name is null or blank");
-    return new IdentifierExpression(name, flatCollectionDataType);
+    return new IdentifierExpression(name, dataType);
   }
 
   public static IdentifierExpression ofString(final String name) {
-    return of(name, PostgresDataType.TEXT);
+    return of(name, DataType.STRING);
   }
 
   public static IdentifierExpression ofInt(final String name) {
-    return of(name, PostgresDataType.INTEGER);
+    return of(name, DataType.INTEGER);
   }
 
   public static IdentifierExpression ofLong(final String name) {
-    return of(name, PostgresDataType.BIGINT);
-  }
-
-  public static IdentifierExpression ofShort(final String name) {
-    return of(name, PostgresDataType.SMALLINT);
+    return of(name, DataType.LONG);
   }
 
   public static IdentifierExpression ofFloat(final String name) {
-    return of(name, PostgresDataType.FLOAT);
+    return of(name, DataType.FLOAT);
   }
 
   public static IdentifierExpression ofDouble(final String name) {
-    return of(name, PostgresDataType.DOUBLE);
-  }
-
-  public static IdentifierExpression ofDecimal(final String name) {
-    return of(name, PostgresDataType.NUMERIC);
+    return of(name, DataType.DOUBLE);
   }
 
   public static IdentifierExpression ofBoolean(final String name) {
-    return of(name, PostgresDataType.BOOLEAN);
-  }
-
-  public static IdentifierExpression ofTimestamp(final String name) {
-    return of(name, PostgresDataType.TIMESTAMP);
+    return of(name, DataType.BOOLEAN);
   }
 
   public static IdentifierExpression ofTimestampTz(final String name) {
-    return of(name, PostgresDataType.TIMESTAMPTZ);
+    return of(name, DataType.TIMESTAMPTZ);
   }
 
   public static IdentifierExpression ofDate(final String name) {
-    return of(name, PostgresDataType.DATE);
-  }
-
-  public static IdentifierExpression ofUuid(final String name) {
-    return of(name, PostgresDataType.UUID);
-  }
-
-  public static IdentifierExpression ofJsonb(final String name) {
-    return of(name, PostgresDataType.JSONB);
-  }
-
-  public static IdentifierExpression ofBytes(final String name) {
-    return of(name, PostgresDataType.BYTEA);
+    return of(name, DataType.DATE);
   }
 
   @Override

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

@@ -9,5 +9,6 @@ public enum JsonFieldType {
   NUMBER_ARRAY,
   BOOLEAN_ARRAY,
   OBJECT_ARRAY,
-  OBJECT
+  OBJECT,
+  UNSPECIFIED
 }

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

@@ -21,7 +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
+  JsonFieldType fieldType;
 
   public static JsonIdentifierExpression of(final String columnName) {
     throw new IllegalArgumentException(