(improvment)(chat) if exist count() in dsl,set query to NATIVE and only order by field and group by field can add to select (#206)

This commit is contained in:
lexluo09
2023-10-13 14:25:24 +08:00
committed by GitHub
parent 119e5b8c58
commit f605cf0ef9
11 changed files with 99 additions and 49 deletions

View File

@@ -44,16 +44,16 @@ public abstract class BaseSemanticCorrector implements SemanticCorrector {
protected void addFieldsToSelect(SemanticCorrectInfo semanticCorrectInfo, String sql) {
Set<String> selectFields = new HashSet<>(SqlParserSelectHelper.getSelectFields(sql));
Set<String> whereFields = new HashSet<>(SqlParserSelectHelper.getWhereFields(sql));
Set<String> needAddFields = new HashSet<>(SqlParserSelectHelper.getGroupByFields(sql));
needAddFields.addAll(SqlParserSelectHelper.getOrderByFields(sql));
if (CollectionUtils.isEmpty(selectFields) || CollectionUtils.isEmpty(whereFields)) {
if (CollectionUtils.isEmpty(selectFields) || CollectionUtils.isEmpty(needAddFields)) {
return;
}
whereFields.addAll(SqlParserSelectHelper.getOrderByFields(sql));
whereFields.removeAll(selectFields);
whereFields.remove(DateUtils.DATE_FIELD);
String replaceFields = SqlParserAddHelper.addFieldsToSelect(sql, new ArrayList<>(whereFields));
needAddFields.removeAll(selectFields);
needAddFields.remove(DateUtils.DATE_FIELD);
String replaceFields = SqlParserAddHelper.addFieldsToSelect(sql, new ArrayList<>(needAddFields));
semanticCorrectInfo.setSql(replaceFields);
}

View File

@@ -6,6 +6,7 @@ import com.tencent.supersonic.chat.api.pojo.request.ExecuteQueryReq;
import com.tencent.supersonic.chat.api.pojo.response.EntityInfo;
import com.tencent.supersonic.chat.api.pojo.response.QueryResult;
import com.tencent.supersonic.chat.query.QueryManager;
import com.tencent.supersonic.chat.query.llm.dsl.DslQuery;
import com.tencent.supersonic.chat.service.SemanticService;
import com.tencent.supersonic.common.util.ContextUtils;
import java.util.List;
@@ -22,7 +23,8 @@ public class EntityInfoExecuteResponder implements ExecuteResponder {
if (semanticParseInfo == null || semanticParseInfo.getModelId() <= 0L) {
return;
}
if (QueryManager.isPluginQuery(semanticParseInfo.getQueryMode())) {
String queryMode = semanticParseInfo.getQueryMode();
if (QueryManager.isPluginQuery(queryMode) && !DslQuery.QUERY_MODE.equals(queryMode)) {
return;
}
SemanticService semanticService = ContextUtils.getBean(SemanticService.class);

View File

@@ -24,7 +24,7 @@ public class EntityInfoParseResponder implements ParseResponder {
QueryReq queryReq = queryContext.getRequest();
selectedParses.forEach(parseInfo -> {
if (QueryManager.isPluginQuery(parseInfo.getQueryMode())
&& !parseInfo.getQueryMode().equals(DslQuery.QUERY_MODE)) {
&& !DslQuery.QUERY_MODE.equals(parseInfo.getQueryMode())) {
return;
}
//1. set entity info

View File

@@ -46,7 +46,7 @@ public class ChatServiceImpl implements ChatService {
private SolvedQueryManager solvedQueryManager;
public ChatServiceImpl(ChatContextRepository chatContextRepository, ChatRepository chatRepository,
ChatQueryRepository chatQueryRepository, SolvedQueryManager solvedQueryManager) {
ChatQueryRepository chatQueryRepository, SolvedQueryManager solvedQueryManager) {
this.chatContextRepository = chatContextRepository;
this.chatRepository = chatRepository;
this.chatQueryRepository = chatQueryRepository;
@@ -174,9 +174,9 @@ public class ChatServiceImpl implements ChatService {
@Override
public void batchAddParse(ChatContext chatCtx, QueryReq queryReq,
ParseResp parseResult,
List<SemanticParseInfo> candidateParses,
List<SemanticParseInfo> selectedParses) {
ParseResp parseResult,
List<SemanticParseInfo> candidateParses,
List<SemanticParseInfo> selectedParses) {
chatQueryRepository.batchSaveParseInfo(chatCtx, queryReq, parseResult, candidateParses, selectedParses);
}
@@ -205,6 +205,8 @@ public class ChatServiceImpl implements ChatService {
List<SolvedQueryRecallResp> solvedQueryRecallResps = solvedQueryManager.recallSolvedQuery(queryText, agentId);
List<Long> queryIds = solvedQueryRecallResps.stream()
.map(SolvedQueryRecallResp::getQueryId).collect(Collectors.toList());
List<Long> queryIds = solvedQueryRecallResps.stream().map(SolvedQueryRecallResp::getQueryId)
.collect(Collectors.toList());
PageQueryInfoReq pageQueryInfoReq = new PageQueryInfoReq();
pageQueryInfoReq.setIds(queryIds);
pageQueryInfoReq.setPageSize(100);
@@ -219,7 +221,7 @@ public class ChatServiceImpl implements ChatService {
queryResp.getScore() != null && queryResp.getScore() <= lowScoreThreshold)
.map(QueryResp::getQuestionId).collect(Collectors.toSet());
return solvedQueryRecallResps.stream().filter(solvedQueryRecallResp ->
!lowScoreQueryIds.contains(solvedQueryRecallResp.getQueryId()))
!lowScoreQueryIds.contains(solvedQueryRecallResp.getQueryId()))
.collect(Collectors.toList());
}

View File

@@ -1,19 +0,0 @@
package com.tencent.supersonic.common.util.jsqlparser;
import net.sf.jsqlparser.expression.ExpressionVisitorAdapter;
import net.sf.jsqlparser.expression.Function;
public class AggregateFunctionVisitor extends ExpressionVisitorAdapter {
private boolean hasAggregateFunction = false;
public boolean hasAggregateFunction() {
return hasAggregateFunction;
}
@Override
public void visit(Function function) {
super.visit(function);
hasAggregateFunction = true;
}
}

View File

@@ -0,0 +1,21 @@
package com.tencent.supersonic.common.util.jsqlparser;
import java.util.HashSet;
import java.util.Set;
import net.sf.jsqlparser.expression.ExpressionVisitorAdapter;
import net.sf.jsqlparser.expression.Function;
public class FunctionVisitor extends ExpressionVisitorAdapter {
private Set<String> functionNames = new HashSet<>();
public Set<String> getFunctionNames() {
return functionNames;
}
@Override
public void visit(Function function) {
super.visit(function);
functionNames.add(function.getName());
}
}

View File

@@ -1,8 +1,10 @@
package com.tencent.supersonic.common.util.jsqlparser;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.Function;
@@ -13,6 +15,7 @@ import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SelectBody;
import net.sf.jsqlparser.statement.select.SelectItem;
import org.apache.commons.lang3.StringUtils;
import org.springframework.util.CollectionUtils;
/**
* Sql Parser Select function Helper
@@ -21,30 +24,34 @@ import org.apache.commons.lang3.StringUtils;
public class SqlParserSelectFunctionHelper {
public static boolean hasAggregateFunction(String sql) {
if (hasFunction(sql)) {
if (!CollectionUtils.isEmpty(getFunctions(sql))) {
return true;
}
return SqlParserSelectHelper.hasGroupBy(sql);
}
public static boolean hasFunction(String sql) {
public static boolean hasFunction(String sql, String functionName) {
Set<String> functions = getFunctions(sql);
if (!CollectionUtils.isEmpty(functions)) {
return functions.stream().anyMatch(function -> function.equalsIgnoreCase(functionName));
}
return false;
}
public static Set<String> getFunctions(String sql) {
Select selectStatement = SqlParserSelectHelper.getSelect(sql);
SelectBody selectBody = selectStatement.getSelectBody();
if (!(selectBody instanceof PlainSelect)) {
return false;
return new HashSet<>();
}
PlainSelect plainSelect = (PlainSelect) selectBody;
List<SelectItem> selectItems = plainSelect.getSelectItems();
AggregateFunctionVisitor visitor = new AggregateFunctionVisitor();
FunctionVisitor visitor = new FunctionVisitor();
for (SelectItem selectItem : selectItems) {
selectItem.accept(visitor);
}
boolean selectFunction = visitor.hasAggregateFunction();
if (selectFunction) {
return true;
}
return false;
return visitor.getFunctionNames();
}
public static Function getFunction(Expression expression, Map<String, String> fieldNameToAggregate) {

View File

@@ -43,4 +43,34 @@ class SqlParserSelectFunctionHelperTest {
Assert.assertEquals(hasAggregateFunction, true);
}
@Test
void hasFunction() throws JSQLParserException {
String sql = "select 部门,sum (访问次数) from 超音数 where 数据日期 = '2023-08-08' "
+ "and 用户 =alice and 发布日期 ='11' group by 部门 limit 1";
boolean hasFunction = SqlParserSelectFunctionHelper.hasFunction(sql, "sum");
Assert.assertEquals(hasFunction, true);
sql = "select 部门,count (访问次数) from 超音数 where 数据日期 = '2023-08-08' "
+ "and 用户 =alice and 发布日期 ='11' group by 部门 limit 1";
hasFunction = SqlParserSelectFunctionHelper.hasFunction(sql, "count");
Assert.assertEquals(hasFunction, true);
sql = "select 部门,count (*) from 超音数 where 数据日期 = '2023-08-08' "
+ "and 用户 =alice and 发布日期 ='11' group by 部门 limit 1";
hasFunction = SqlParserSelectFunctionHelper.hasFunction(sql, "count");
Assert.assertEquals(hasFunction, true);
sql = "SELECT user_name, pv FROM t_34 WHERE sys_imp_date <= '2023-09-03' "
+ "AND sys_imp_date >= '2023-08-04' GROUP BY user_name ORDER BY sum(pv) DESC LIMIT 10";
hasFunction = SqlParserSelectFunctionHelper.hasFunction(sql, "sum");
Assert.assertEquals(hasFunction, false);
sql = "select 部门,min (访问次数) from 超音数 where 数据日期 = '2023-08-08' "
+ "and 用户 =alice and 发布日期 ='11' group by 部门 limit 1";
hasFunction = SqlParserSelectFunctionHelper.hasFunction(sql, "min");
Assert.assertEquals(hasFunction, true);
}
}

View File

@@ -2,6 +2,7 @@ package com.tencent.supersonic.semantic.query.parser.convert;
import com.tencent.supersonic.common.util.DateUtils;
import com.tencent.supersonic.common.util.jsqlparser.SqlParserReplaceHelper;
import com.tencent.supersonic.common.util.jsqlparser.SqlParserSelectFunctionHelper;
import com.tencent.supersonic.common.util.jsqlparser.SqlParserSelectHelper;
import com.tencent.supersonic.semantic.api.model.enums.TimeDimensionEnum;
import com.tencent.supersonic.semantic.api.model.pojo.SchemaItem;
@@ -81,10 +82,7 @@ public class QueryReqConverter {
queryStructUtils.generateInternalMetricName(databaseReq.getModelId(),
metricTable.getDimensions()))));
}
// if there is no group by in dsl,set MetricTable's aggOption to "NATIVE"
if (!SqlParserSelectHelper.hasGroupBy(databaseReq.getSql())) {
metricTable.setAggOption(AggOption.NATIVE);
}
metricTable.setAggOption(getAggOption(databaseReq));
List<MetricTable> tables = new ArrayList<>();
tables.add(metricTable);
//4.build ParseSqlReq
@@ -104,6 +102,17 @@ public class QueryReqConverter {
return queryStatement;
}
private AggOption getAggOption(QueryDslReq databaseReq) {
// if there is no group by in dsl,set MetricTable's aggOption to "NATIVE"
// if there is count() in dsl,set MetricTable's aggOption to "NATIVE"
String sql = databaseReq.getSql();
if (!SqlParserSelectHelper.hasGroupBy(sql)
|| SqlParserSelectFunctionHelper.hasFunction(sql, "count")) {
return AggOption.NATIVE;
}
return AggOption.DEFAULT;
}
private void convertNameToBizName(QueryDslReq databaseReq, ModelSchemaResp modelSchemaResp) {
Map<String, String> fieldNameToBizNameMap = getFieldNameToBizNameMap(modelSchemaResp);
String sql = databaseReq.getSql();

View File

@@ -61,8 +61,7 @@ public class QueryController {
@PostMapping("/queryStatement")
public Object queryStatement(@RequestBody QueryStatement queryStatement) throws Exception {
Object result = queryService.queryByQueryStatement(queryStatement);
return result;
return queryService.queryByQueryStatement(queryStatement);
}
@PostMapping("/struct/parse")

View File

@@ -89,8 +89,7 @@ public class QueryServiceImpl implements QueryService {
}
public Object queryByQueryStatement(QueryStatement queryStatement) {
QueryResultWithSchemaResp results = semanticQueryEngine.execute(queryStatement);
return results;
return semanticQueryEngine.execute(queryStatement);
}
private QueryStatement convertToQueryStatement(QueryDslReq querySqlCmd, User user) throws Exception {