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() {