(improvement)(Headless) distinct select fields in S2CorrectSQL (#912)

chat/server/src/main/java/com/tencent/supersonic/chat/server/parser/NL2SQLParser.java

@@ -8,7 +8,9 @@ import com.tencent.supersonic.headless.api.pojo.request.QueryReq;
 import com.tencent.supersonic.headless.api.pojo.response.ParseResp;
 import com.tencent.supersonic.headless.server.service.ChatQueryService;
 import java.util.List;
+import lombok.extern.slf4j.Slf4j;
 
+@Slf4j
 public class NL2SQLParser implements ChatParser {
 
     @Override

chat/server/src/main/java/com/tencent/supersonic/chat/server/persistence/repository/ChatQueryRepository.java

@@ -35,4 +35,5 @@ public interface ChatQueryRepository {
     List<ChatParseDO> getParseInfoList(List<Long> questionIds);
 
     Boolean deleteChatQuery(Long questionId);
+
 }

chat/server/src/main/java/com/tencent/supersonic/chat/server/persistence/repository/impl/ChatQueryRepositoryImpl.java

@@ -192,4 +192,5 @@ public class ChatQueryRepositoryImpl implements ChatQueryRepository {
     public Boolean deleteChatQuery(Long questionId) {
         return chatQueryDOMapper.deleteByPrimaryKey(questionId);
     }
+
 }

common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/SqlAddHelper.java

@@ -5,6 +5,7 @@ import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 import java.util.ArrayList;
+
 import lombok.extern.slf4j.Slf4j;
 import net.sf.jsqlparser.expression.Expression;
 import net.sf.jsqlparser.expression.Function;
@@ -12,6 +13,7 @@ import net.sf.jsqlparser.expression.LongValue;
 import net.sf.jsqlparser.expression.Parenthesis;
 import net.sf.jsqlparser.expression.StringValue;
 import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
+import net.sf.jsqlparser.expression.operators.conditional.OrExpression;
 import net.sf.jsqlparser.expression.operators.relational.ComparisonOperator;
 import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
 import net.sf.jsqlparser.schema.Column;
@@ -225,7 +227,7 @@ public class SqlAddHelper {
     }
 
     private static void addAggregateToSelectItems(List<SelectItem> selectItems,
-            Map<String, String> fieldNameToAggregate) {
+                                                  Map<String, String> fieldNameToAggregate) {
         for (SelectItem selectItem : selectItems) {
             if (selectItem instanceof SelectExpressionItem) {
                 SelectExpressionItem selectExpressionItem = (SelectExpressionItem) selectItem;
@@ -240,7 +242,7 @@ public class SqlAddHelper {
     }
 
     private static void addAggregateToOrderByItems(List<OrderByElement> orderByElements,
-            Map<String, String> fieldNameToAggregate) {
+                                                   Map<String, String> fieldNameToAggregate) {
         if (orderByElements == null) {
             return;
         }
@@ -255,7 +257,7 @@ public class SqlAddHelper {
     }
 
     private static void addAggregateToGroupByItems(GroupByElement groupByElement,
-            Map<String, String> fieldNameToAggregate) {
+                                                   Map<String, String> fieldNameToAggregate) {
         if (groupByElement == null) {
             return;
         }
@@ -276,13 +278,22 @@ public class SqlAddHelper {
     }
 
     private static void modifyWhereExpression(Expression whereExpression,
-            Map<String, String> fieldNameToAggregate) {
+                                              Map<String, String> fieldNameToAggregate) {
         if (SqlSelectHelper.isLogicExpression(whereExpression)) {
-            AndExpression andExpression = (AndExpression) whereExpression;
+            if (whereExpression instanceof AndExpression) {
-            Expression leftExpression = andExpression.getLeftExpression();
+                AndExpression andExpression = (AndExpression) whereExpression;
-            Expression rightExpression = andExpression.getRightExpression();
+                Expression leftExpression = andExpression.getLeftExpression();
-            modifyWhereExpression(leftExpression, fieldNameToAggregate);
+                Expression rightExpression = andExpression.getRightExpression();
-            modifyWhereExpression(rightExpression, fieldNameToAggregate);
+                modifyWhereExpression(leftExpression, fieldNameToAggregate);
+                modifyWhereExpression(rightExpression, fieldNameToAggregate);
+            }
+            if (whereExpression instanceof OrExpression) {
+                OrExpression orExpression = (OrExpression) whereExpression;
+                Expression leftExpression = orExpression.getLeftExpression();
+                Expression rightExpression = orExpression.getRightExpression();
+                modifyWhereExpression(leftExpression, fieldNameToAggregate);
+                modifyWhereExpression(rightExpression, fieldNameToAggregate);
+            }
         } else if (whereExpression instanceof Parenthesis) {
             modifyWhereExpression(((Parenthesis) whereExpression).getExpression(), fieldNameToAggregate);
         } else {

common/src/main/java/com/tencent/supersonic/common/util/jsqlparser/SqlRemoveHelper.java

@@ -29,6 +29,7 @@ import net.sf.jsqlparser.statement.select.SelectItem;
 import net.sf.jsqlparser.statement.select.SelectVisitorAdapter;
 import org.springframework.util.CollectionUtils;
 
+import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
 import java.util.Objects;
@@ -60,6 +61,32 @@ public class SqlRemoveHelper {
         return selectStatement.toString();
     }
 
+    public static String removeSameFieldFromSelect(String sql) {
+        Select selectStatement = SqlSelectHelper.getSelect(sql);
+        if (selectStatement == null) {
+            return sql;
+        }
+        SelectBody selectBody = selectStatement.getSelectBody();
+        if (!(selectBody instanceof PlainSelect)) {
+            return sql;
+        }
+        List<SelectItem> selectItems = ((PlainSelect) selectBody).getSelectItems();
+        Set<String> fields = new HashSet<>();
+        selectItems.removeIf(selectItem -> {
+            if (selectItem instanceof SelectExpressionItem) {
+                SelectExpressionItem selectExpressionItem = (SelectExpressionItem) selectItem;
+                String field = selectExpressionItem.getExpression().toString();
+                if (fields.contains(field)) {
+                    return true;
+                }
+                fields.add(field);
+            }
+            return false;
+        });
+        ((PlainSelect) selectBody).setSelectItems(selectItems);
+        return selectStatement.toString();
+    }
+
     public static String removeWhereCondition(String sql, Set<String> removeFieldNames) {
         Select selectStatement = SqlSelectHelper.getSelect(sql);
         SelectBody selectBody = selectStatement.getSelectBody();

common/src/test/java/com/tencent/supersonic/common/util/jsqlparser/SqlRemoveHelperTest.java

@@ -1,6 +1,7 @@
 package com.tencent.supersonic.common.util.jsqlparser;
 
 import java.util.HashSet;
+import java.util.List;
 import java.util.Set;
 import org.junit.Assert;
 import org.junit.jupiter.api.Test;
@@ -10,6 +11,21 @@ import org.junit.jupiter.api.Test;
  */
 class SqlRemoveHelperTest {
 
+    @Test
+    void testRemoveSameFieldFromSelect() {
+        String sql = "select 歌曲名,歌手名,粉丝数,粉丝数,sum(粉丝数),sum(粉丝数),avg(播放量),avg(播放量)"
+                + " from 歌曲库 where sum(粉丝数) > 20000 and  2>1 and "
+                + "sum(播放量) > 20000 and 1=1  HAVING sum(播放量) > 20000 and 3>1";
+        sql = SqlRemoveHelper.removeSameFieldFromSelect(sql);
+        System.out.println(sql);
+        sql = "SELECT 结算播放量 FROM 艺人 WHERE (歌手名 IN ('林俊杰', '陈奕迅')) AND (数据日期 >= '2024-04-04' AND 数据日期 <= '2024-04-04')";
+        List<FieldExpression> fieldExpressionList = SqlSelectHelper.getWhereExpressions(sql);
+        fieldExpressionList.stream().forEach(fieldExpression -> {
+            System.out.println(fieldExpression.toString());
+        });
+
+    }
+
     @Test
     void testRemoveWhereHavingCondition() {
         String sql = "select 歌曲名 from 歌曲库 where sum(粉丝数) > 20000 and  2>1 and "

headless/core/src/main/java/com/tencent/supersonic/headless/core/chat/corrector/GrammarCorrector.java

@@ -1,5 +1,6 @@
 package com.tencent.supersonic.headless.core.chat.corrector;
 
+import com.tencent.supersonic.common.util.jsqlparser.SqlRemoveHelper;
 import com.tencent.supersonic.headless.api.pojo.SemanticParseInfo;
 import com.tencent.supersonic.headless.core.pojo.QueryContext;
 import lombok.extern.slf4j.Slf4j;
@@ -28,5 +29,12 @@ public class GrammarCorrector extends BaseSemanticCorrector {
         for (BaseSemanticCorrector corrector : correctors) {
             corrector.correct(queryContext, semanticParseInfo);
         }
+        removeSameFieldFromSelect(semanticParseInfo);
+    }
+
+    public void removeSameFieldFromSelect(SemanticParseInfo semanticParseInfo) {
+        String correctS2SQL = semanticParseInfo.getSqlInfo().getCorrectS2SQL();
+        correctS2SQL = SqlRemoveHelper.removeSameFieldFromSelect(correctS2SQL);
+        semanticParseInfo.getSqlInfo().setCorrectS2SQL(correctS2SQL);
     }
 }

headless/core/src/main/java/com/tencent/supersonic/headless/core/chat/parser/llm/LLMRequestService.java

@@ -61,11 +61,12 @@ public class LLMRequestService {
     }
 
     public LLMReq getLlmReq(QueryContext queryCtx, Long dataSetId,
-            SemanticSchema semanticSchema, List<LLMReq.ElementValue> linkingValues) {
+                            SemanticSchema semanticSchema, List<LLMReq.ElementValue> linkingValues) {
         Map<Long, String> dataSetIdToName = semanticSchema.getDataSetIdToName();
         String queryText = queryCtx.getQueryText();
 
         LLMReq llmReq = new LLMReq();
+
         llmReq.setQueryText(queryText);
         LLMReq.FilterCondition filterCondition = new LLMReq.FilterCondition();
         llmReq.setFilterCondition(filterCondition);
@@ -103,7 +104,7 @@ public class LLMRequestService {
     }
 
     protected List<String> getFieldNameList(QueryContext queryCtx, Long dataSetId,
-            LLMParserConfig llmParserConfig) {
+                                            LLMParserConfig llmParserConfig) {
 
         Set<String> results = getTopNFieldNames(queryCtx, dataSetId, llmParserConfig);
 

headless/core/src/main/java/com/tencent/supersonic/headless/core/chat/parser/llm/LLMSqlParser.java

@@ -11,6 +11,7 @@ import com.tencent.supersonic.headless.core.chat.query.llm.s2sql.LLMResp;
 import com.tencent.supersonic.headless.core.chat.query.llm.s2sql.LLMSqlResp;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.collections4.MapUtils;
+
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -28,6 +29,7 @@ public class LLMSqlParser implements SemanticParser {
         try {
             //2.get dataSetId from queryCtx and chatCtx.
             Long dataSetId = requestService.getDataSetId(queryCtx);
+            log.info("dataSetId:{}", dataSetId);
             if (dataSetId == null) {
                 return;
             }

headless/core/src/main/java/com/tencent/supersonic/headless/core/chat/parser/llm/SqlPromptGenerator.java

@@ -6,11 +6,11 @@ import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang3.tuple.Pair;
 import org.springframework.stereotype.Component;
 
+import java.util.List;
 import java.util.ArrayList;
+import java.util.Map;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.List;
-import java.util.Map;
 
 @Component
 @Slf4j
@@ -85,7 +85,7 @@ public class SqlPromptGenerator {
     }
 
     public List<List<Map<String, String>>> getExampleCombos(List<Map<String, String>> exampleList, int numFewShots,
-            int numSelfConsistency) {
+                                                            int numSelfConsistency) {
         List<List<Map<String, String>>> results = new ArrayList<>();
         for (int i = 0; i < numSelfConsistency; i++) {
             List<Map<String, String>> shuffledList = new ArrayList<>(exampleList);
@@ -118,7 +118,7 @@ public class SqlPromptGenerator {
     }
 
     public List<String> generateSqlPromptPool(LLMReq llmReq, List<String> schemaLinkStrPool,
-            List<List<Map<String, String>>> fewshotExampleListPool) {
+                                              List<List<Map<String, String>>> fewshotExampleListPool) {
         List<String> sqlPromptPool = new ArrayList<>();
         for (int i = 0; i < schemaLinkStrPool.size(); i++) {
             String schemaLinkStr = schemaLinkStrPool.get(i);

headless/core/src/main/java/com/tencent/supersonic/headless/core/pojo/QueryContext.java

@@ -15,6 +15,7 @@ import lombok.Builder;
 import lombok.Data;
 import lombok.NoArgsConstructor;
 
+
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.List;

headless/server/src/main/java/com/tencent/supersonic/headless/server/persistence/repository/impl/ChatContextRepositoryImpl.java

@@ -8,7 +8,6 @@ import com.tencent.supersonic.headless.server.persistence.dataobject.ChatContext
 import com.tencent.supersonic.headless.server.persistence.mapper.ChatContextMapper;
 import com.tencent.supersonic.headless.server.persistence.repository.ChatContextRepository;
 import lombok.extern.slf4j.Slf4j;
-import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.annotation.Primary;
 import org.springframework.stereotype.Repository;
 
@@ -17,7 +16,7 @@ import org.springframework.stereotype.Repository;
 @Slf4j
 public class ChatContextRepositoryImpl implements ChatContextRepository {
 
-    @Autowired(required = false)
+
     private final ChatContextMapper chatContextMapper;
 
     public ChatContextRepositoryImpl(ChatContextMapper chatContextMapper) {