From 6b2a14e5892d212903ee4dcc065fd870b934430c Mon Sep 17 00:00:00 2001
From: lexluo09 <39718951+lexluo09@users.noreply.github.com>
Date: Wed, 11 Oct 2023 13:14:42 +0800
Subject: [PATCH] (improvement)(chat) The group by field can be added even if
 the group by field already exists (#192)

---
 .../chat/responder/execute/EntityInfoExecuteResponder.java  | 6 ++++++
 .../common/util/jsqlparser/SqlParserUpdateHelper.java       | 6 +++++-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/chat/core/src/main/java/com/tencent/supersonic/chat/responder/execute/EntityInfoExecuteResponder.java b/chat/core/src/main/java/com/tencent/supersonic/chat/responder/execute/EntityInfoExecuteResponder.java
index 4659e0bcc..2b5d67230 100644
--- a/chat/core/src/main/java/com/tencent/supersonic/chat/responder/execute/EntityInfoExecuteResponder.java
+++ b/chat/core/src/main/java/com/tencent/supersonic/chat/responder/execute/EntityInfoExecuteResponder.java
@@ -17,6 +17,9 @@ public class EntityInfoExecuteResponder implements ExecuteResponder {
 
     @Override
     public void fillResponse(QueryResult queryResult, SemanticParseInfo semanticParseInfo, ExecuteQueryReq queryReq) {
+        if (semanticParseInfo == null || semanticParseInfo.getModelId() <= 0L) {
+            return;
+        }
         SemanticService semanticService = ContextUtils.getBean(SemanticService.class);
         User user = queryReq.getUser();
         queryResult.setEntityInfo(semanticService.getEntityInfo(semanticParseInfo, user));
@@ -27,6 +30,9 @@ public class EntityInfoExecuteResponder implements ExecuteResponder {
             return;
         }
         String primaryEntityBizName = entityInfo.getModelInfo().getPrimaryEntityBizName();
+        if (CollectionUtils.isEmpty(queryResult.getQueryColumns())) {
+            return;
+        }
         boolean existPrimaryEntityName = queryResult.getQueryColumns().stream()
                 .anyMatch(queryColumn -> primaryEntityBizName.equals(queryColumn.getNameEn()));
 
diff --git a/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/SqlParserUpdateHelper.java b/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/SqlParserUpdateHelper.java
index 0191d9b6c..a1f9ffa98 100644
--- a/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/SqlParserUpdateHelper.java
+++ b/common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/SqlParserUpdateHelper.java
@@ -312,7 +312,7 @@ public class SqlParserUpdateHelper {
     }
 
     public static String addGroupBy(String sql, Set<String> groupByFields) {
-        if (SqlParserSelectHelper.hasGroupBy(sql) || CollectionUtils.isEmpty(groupByFields)) {
+        if (CollectionUtils.isEmpty(groupByFields)) {
             return sql;
         }
         Select selectStatement = SqlParserSelectHelper.getSelect(sql);
@@ -324,6 +324,10 @@ public class SqlParserUpdateHelper {
 
         PlainSelect plainSelect = (PlainSelect) selectBody;
         GroupByElement groupByElement = new GroupByElement();
+        List<String> originalGroupByFields = SqlParserSelectHelper.getGroupByFields(sql);
+        if (!CollectionUtils.isEmpty(originalGroupByFields)) {
+            groupByFields.addAll(originalGroupByFields);
+        }
         for (String groupByField : groupByFields) {
             groupByElement.addGroupByExpression(new Column(groupByField));
         }