diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/query/ColumnMetadata.java b/document-store/src/main/java/org/hypertrace/core/documentstore/query/ColumnMetadata.java
new file mode 100644
index 00000000..b8e8e27a
--- /dev/null
+++ b/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.
+ *
+ *
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.
+ *
+ *
Example usage:
+ * ColumnMetadata metadata = ColumnMetadata.builder() .column("tags", "STRING_ARRAY")
+ * .column("scores", "LONG_ARRAY") .column("attributes", "STRING_MAP") .build();
+ *
+ */
+@Value
+@Builder
+public class ColumnMetadata {
+
+ @Singular
+ Map columns; // column_name -> type_string
+
+ // Known native array type suffixes (from entity-service ValueType)
+ private static final Set 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 NATIVE_SCALAR_TYPES = Set.of(
+ "STRING",
+ "LONG",
+ "DOUBLE",
+ "BYTES",
+ "BOOL",
+ "TIMESTAMP"
+ );
+
+ // Known JSONB types (from entity-service ValueType)
+ private static final Set 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();
+ }
+}
diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/query/Query.java b/document-store/src/main/java/org/hypertrace/core/documentstore/query/Query.java
index 0cb3ee7b..35dcd82d 100644
--- a/document-store/src/main/java/org/hypertrace/core/documentstore/query/Query.java
+++ b/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 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 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() {