[PoC] Add column metadata to Query object

document-store/src/main/java/org/hypertrace/core/documentstore/query/ColumnMetadata.java

@@ -0,0 +1,114 @@
+package org.hypertrace.core.documentstore.query;
+
+import java.util.Map;
+import java.util.Set;
+import lombok.Builder;
+import lombok.Singular;
+import lombok.Value;
+
+/**
+ * Metadata about column types in a collection. Used to inform query parsers about native column
+ * types vs JSONB fields.
+ *
+ * <p>Column types are represented as strings to avoid coupling with entity-service type enums.
+ * Entity-service can pass AttributeKind.name() or ValueType.name() as the type string.
+ *
+ * <p>Example usage: <code>
+ * ColumnMetadata metadata = ColumnMetadata.builder() .column("tags", "STRING_ARRAY")
+ * .column("scores", "LONG_ARRAY") .column("attributes", "STRING_MAP") .build();
+ * </code>
+ */
+@Value
+@Builder
+public class ColumnMetadata {
+
+  @Singular
+  Map<String, String> columns; // column_name -> type_string
+
+  // Known native array type suffixes (from entity-service ValueType)
+  private static final Set<String> NATIVE_ARRAY_TYPES = Set.of(
+      "STRING_ARRAY",
+      "LONG_ARRAY",
+      "DOUBLE_ARRAY",
+      "BOOLEAN_ARRAY"
+  );
+
+  // Known native scalar types (from entity-service ValueType)
+  private static final Set<String> NATIVE_SCALAR_TYPES = Set.of(
+      "STRING",
+      "LONG",
+      "DOUBLE",
+      "BYTES",
+      "BOOL",
+      "TIMESTAMP"
+  );
+
+  // Known JSONB types (from entity-service ValueType)
+  private static final Set<String> JSONB_TYPES = Set.of(
+      "STRING_MAP"
+  );
+
+  /**
+   * Get the column type string for a given column name.
+   *
+   * @param columnName the name of the column
+   * @return the column type string, or null if not specified (defaults to JSONB storage)
+   */
+  public String getColumnType(String columnName) {
+    return columns.get(columnName);
+  }
+
+  /**
+   * Check if a column is a native array type.
+   *
+   * @param columnName the name of the column
+   * @return true if the column is a native array type
+   */
+  public boolean isNativeArrayColumn(String columnName) {
+    String type = getColumnType(columnName);
+    if (type == null) {
+      return false;
+    }
+    // Check exact match first, then check suffix pattern
+    return NATIVE_ARRAY_TYPES.contains(type) || type.endsWith("_ARRAY");
+  }
+
+  /**
+   * Check if a column is a native scalar type.
+   *
+   * @param columnName the name of the column
+   * @return true if the column is a native scalar type
+   */
+  public boolean isNativeScalarColumn(String columnName) {
+    String type = getColumnType(columnName);
+    return type != null && NATIVE_SCALAR_TYPES.contains(type);
+  }
+
+  /**
+   * Check if a column is a native type (array or scalar).
+   *
+   * @param columnName the name of the column
+   * @return true if the column is a native type
+   */
+  public boolean isNativeColumn(String columnName) {
+    return isNativeArrayColumn(columnName) || isNativeScalarColumn(columnName);
+  }
+
+  /**
+   * Check if a column is explicitly a JSONB type (like STRING_MAP).
+   *
+   * @param columnName the name of the column
+   * @return true if the column is explicitly marked as JSONB type
+   */
+  public boolean isJsonbColumn(String columnName) {
+    String type = getColumnType(columnName);
+    return type != null && JSONB_TYPES.contains(type);
+  }
+
+  /**
+   * @return an empty ColumnMetadata instance (all columns default to JSONB storage)
+   */
+  public static ColumnMetadata empty() {
+    return ColumnMetadata.builder().build();
+  }
+}

document-store/src/main/java/org/hypertrace/core/documentstore/query/Query.java

@@ -76,6 +76,7 @@ public class Query {
   Sort sort;
   Pagination pagination; // Missing pagination represents fetching all the records
   FromClause fromClause;
+  ColumnMetadata columnMetadata; // Column type metadata for native vs JSONB columns
 
   @Override
   public String toString() {
@@ -116,6 +117,10 @@ public List<FromTypeExpression> getFromTypeExpressions() {
     return fromClause == null ? emptyList() : unmodifiableList(fromClause.getFromTypeExpressions());
   }
 
+  public ColumnMetadata getColumnMetadata() {
+    return columnMetadata != null ? columnMetadata : ColumnMetadata.empty();
+  }
+
   public static QueryBuilder builder() {
     return new QueryBuilder();
   }
@@ -157,6 +162,7 @@ public static class QueryBuilder {
     private Sort.SortBuilder sortBuilder;
     private Pagination pagination;
     private FromClause.FromClauseBuilder fromClauseBuilder;
+    private ColumnMetadata columnMetadata;
 
     public QueryBuilder setSelection(final Selection selection) {
       this.selectionBuilder = selection.toBuilder();
@@ -306,6 +312,22 @@ public QueryBuilder addFromClauses(final List<FromTypeExpression> expressions) {
       return this;
     }
 
+    public QueryBuilder setColumnMetadata(final ColumnMetadata columnMetadata) {
+      this.columnMetadata = columnMetadata;
+      return this;
+    }
+
+    public QueryBuilder addColumnType(final String columnName, final String type) {
+      if (this.columnMetadata == null) {
+        this.columnMetadata = ColumnMetadata.builder().build();
+      }
+      ColumnMetadata.ColumnMetadataBuilder builder = ColumnMetadata.builder();
+      this.columnMetadata.getColumns().forEach(builder::column);
+      builder.column(columnName, type);
+      this.columnMetadata = builder.build();
+      return this;
+    }
+
     public Query build() {
       return new Query(
           getSelection(),
@@ -314,7 +336,8 @@ public Query build() {
           getAggregationFilter(),
           getSort(),
           pagination,
-          getFrom());
+          getFrom(),
+          columnMetadata);
     }
 
     protected Selection.SelectionBuilder getSelectionBuilder() {