From e0f7ec0f40e518f30b63192b16709f8d476a16b7 Mon Sep 17 00:00:00 2001
From: jerryjzhang <jerryjzhang@tencent.com>
Date: Sat, 4 Jan 2025 22:24:34 +0800
Subject: [PATCH] [improvement][headless]Add modelId to `QueryColumn`.

---
 .../supersonic/common/pojo/QueryColumn.java   |  1 +
 .../headless/server/utils/QueryUtils.java     | 22 ++++++++++++++-----
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/common/src/main/java/com/tencent/supersonic/common/pojo/QueryColumn.java b/common/src/main/java/com/tencent/supersonic/common/pojo/QueryColumn.java
index c2c5374cf..13a46d08c 100644
--- a/common/src/main/java/com/tencent/supersonic/common/pojo/QueryColumn.java
+++ b/common/src/main/java/com/tencent/supersonic/common/pojo/QueryColumn.java
@@ -17,6 +17,7 @@ public class QueryColumn {
     private String dataFormatType;
     private DataFormat dataFormat;
     private String comment;
+    private Long modelId;
 
     public QueryColumn(String nameEn, String type) {
         this.type = type;
diff --git a/headless/server/src/main/java/com/tencent/supersonic/headless/server/utils/QueryUtils.java b/headless/server/src/main/java/com/tencent/supersonic/headless/server/utils/QueryUtils.java
index e78f5584d..d7b5f2d15 100644
--- a/headless/server/src/main/java/com/tencent/supersonic/headless/server/utils/QueryUtils.java
+++ b/headless/server/src/main/java/com/tencent/supersonic/headless/server/utils/QueryUtils.java
@@ -4,10 +4,7 @@ import com.tencent.supersonic.common.pojo.Constants;
 import com.tencent.supersonic.common.pojo.QueryColumn;
 import com.tencent.supersonic.headless.api.pojo.enums.SemanticType;
 import com.tencent.supersonic.headless.api.pojo.request.QueryMultiStructReq;
-import com.tencent.supersonic.headless.api.pojo.response.MetricResp;
-import com.tencent.supersonic.headless.api.pojo.response.MetricSchemaResp;
-import com.tencent.supersonic.headless.api.pojo.response.SemanticQueryResp;
-import com.tencent.supersonic.headless.api.pojo.response.SemanticSchemaResp;
+import com.tencent.supersonic.headless.api.pojo.response.*;
 import com.tencent.supersonic.headless.core.pojo.QueryStatement;
 import com.tencent.supersonic.headless.core.utils.SqlGenerateUtils;
 import lombok.extern.slf4j.Slf4j;
@@ -39,11 +36,13 @@ public class QueryUtils {
     public void populateQueryColumns(SemanticQueryResp semanticQueryResp,
             SemanticSchemaResp semanticSchemaResp) {
         Map<String, MetricResp> metricRespMap = createMetricRespMap(semanticSchemaResp);
+        Map<String, DimensionResp> dimensionRespMap = createDimRespMap(semanticSchemaResp);
         Map<String, String> namePair = new HashMap<>();
         Map<String, String> nameTypePair = new HashMap<>();
         populateNamePairs(semanticSchemaResp, namePair, nameTypePair);
         List<QueryColumn> columns = semanticQueryResp.getColumns();
-        columns.forEach(column -> processColumn(column, namePair, nameTypePair, metricRespMap));
+        columns.forEach(column -> processColumn(column, namePair, nameTypePair, metricRespMap,
+                dimensionRespMap));
     }
 
     private Map<String, MetricResp> createMetricRespMap(SemanticSchemaResp semanticSchemaResp) {
@@ -52,6 +51,12 @@ public class QueryUtils {
                 .collect(Collectors.toMap(MetricResp::getBizName, a -> a, (k1, k2) -> k1));
     }
 
+    private Map<String, DimensionResp> createDimRespMap(SemanticSchemaResp semanticSchemaResp) {
+        List<DimSchemaResp> dimensions = semanticSchemaResp.getDimensions();
+        return dimensions.stream()
+                .collect(Collectors.toMap(DimensionResp::getBizName, a -> a, (k1, k2) -> k1));
+    }
+
     private void populateNamePairs(SemanticSchemaResp semanticSchemaResp,
             Map<String, String> namePair, Map<String, String> nameTypePair) {
         semanticSchemaResp.getMetrics().forEach(metricDesc -> {
@@ -65,7 +70,8 @@ public class QueryUtils {
     }
 
     private void processColumn(QueryColumn column, Map<String, String> namePair,
-            Map<String, String> nameTypePair, Map<String, MetricResp> metricRespMap) {
+            Map<String, String> nameTypePair, Map<String, MetricResp> metricRespMap,
+            Map<String, DimensionResp> dimensionRespMap) {
         String nameEn = getName(column.getNameEn());
         if (nameEn.contains(JOIN_UNDERLINE)) {
             nameEn = nameEn.split(JOIN_UNDERLINE)[1];
@@ -98,6 +104,10 @@ public class QueryUtils {
         if (metricRespMap.containsKey(nameEn)) {
             column.setDataFormatType(metricRespMap.get(nameEn).getDataFormatType());
             column.setDataFormat(metricRespMap.get(nameEn).getDataFormat());
+            column.setModelId(metricRespMap.get(nameEn).getModelId());
+        }
+        if (dimensionRespMap.containsKey(nameEn)) {
+            column.setModelId(dimensionRespMap.get(nameEn).getModelId());
         }
         // set name by NameEn
         if (StringUtils.isBlank(column.getName()) && StringUtils.isNotBlank(column.getNameEn())) {