topdog/supersonic
1
0
Fork 0
mirror of https://github.com/tencentmusic/supersonic.git synced 2025-12-10 11:07:06 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

[improvement][common] Field replacement is performed using a recursive approach, and it supports field replacement with complex expressions (#1859)

Browse Source
This commit is contained in:
lexluo09
2024-10-30 13:17:28 +08:00
committed by GitHub
parent 9644fc4207
commit 5c70607851
5 changed files with 161 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
common/src/main/java/com/tencent/supersonic/common/jsqlparser/FieldReplaceVisitor.java
View File

@@ -2,9 +2,13 @@ package com.tencent.supersonic.common.jsqlparser;
import com.tencent.supersonic.common.util.ContextUtils;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.expression.AnalyticExpression;
import net.sf.jsqlparser.expression.ExpressionVisitorAdapter;
import net.sf.jsqlparser.expression.Function;
import net.sf.jsqlparser.expression.WindowDefinition;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.statement.select.OrderByElement;
import org.springframework.util.CollectionUtils;
import java.util.Map;
@@ -34,4 +38,16 @@ public class FieldReplaceVisitor extends ExpressionVisitorAdapter {
exactReplace.set(originalExactReplace);
}
}
@Override
public void visit(AnalyticExpression expr) {
super.visit(expr);
WindowDefinition windowDefinition = expr.getWindowDefinition();
if (windowDefinition != null
&& !CollectionUtils.isEmpty(windowDefinition.getOrderByElements())) {
for (OrderByElement element : windowDefinition.getOrderByElements()) {
element.getExpression().accept(this);
}
}
}
}

55
common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceHelper.java
View File

@@ -154,37 +154,20 @@ public class SqlReplaceHelper {
public static String replaceFields(String sql, Map<String, String> fieldNameMap,
boolean exactReplace) {
Select selectStatement = SqlSelectHelper.getSelect(sql);
List<PlainSelect> plainSelectList = SqlSelectHelper.getWithItem(selectStatement);
if (selectStatement instanceof PlainSelect) {
PlainSelect plainSelect = (PlainSelect) selectStatement;
plainSelectList.add(plainSelect);
getFromSelect(plainSelect.getFromItem(), plainSelectList);
} else if (selectStatement instanceof SetOperationList) {
SetOperationList setOperationList = (SetOperationList) selectStatement;
if (!CollectionUtils.isEmpty(setOperationList.getSelects())) {
setOperationList.getSelects().forEach(subSelectBody -> {
PlainSelect subPlainSelect = (PlainSelect) subSelectBody;
plainSelectList.add(subPlainSelect);
getFromSelect(subPlainSelect.getFromItem(), plainSelectList);
});
Set<Select> plainSelectList = SqlSelectHelper.getAllSelect(selectStatement);
for (Select plainSelect : plainSelectList) {
if (plainSelect instanceof PlainSelect) {
replaceFieldsInPlainOneSelect(fieldNameMap, exactReplace,
(PlainSelect) plainSelect);
}
List<OrderByElement> orderByElements = setOperationList.getOrderByElements();
if (!CollectionUtils.isEmpty(orderByElements)) {
for (OrderByElement orderByElement : orderByElements) {
orderByElement.accept(new OrderByReplaceVisitor(fieldNameMap, exactReplace));
if (plainSelect instanceof SetOperationList) {
replaceFieldsInSetOperationList(fieldNameMap, exactReplace,
(SetOperationList) plainSelect);
}
}
} else {
return sql;
}
List<PlainSelect> plainSelects = SqlSelectHelper.getPlainSelects(plainSelectList);
for (PlainSelect plainSelect : plainSelects) {
replaceFieldsInPlainOneSelect(fieldNameMap, exactReplace, plainSelect);
}
return selectStatement.toString();
}
private static void replaceFieldsInPlainOneSelect(Map<String, String> fieldNameMap,
boolean exactReplace, PlainSelect plainSelect) {
// 1. replace where fields
@@ -236,22 +219,24 @@ public class SqlReplaceHelper {
List<Join> joins = plainSelect.getJoins();
if (!CollectionUtils.isEmpty(joins)) {
for (Join join : joins) {
if (!CollectionUtils.isEmpty(join.getOnExpressions())) {
if (CollectionUtils.isEmpty(join.getOnExpressions())) {
continue;
}
join.getOnExpressions().stream().forEach(onExpression -> {
onExpression.accept(visitor);
});
}
if (!(join.getRightItem() instanceof ParenthesedSelect)) {
continue;
}
ParenthesedSelect parenthesedSelect = (ParenthesedSelect) join.getRightItem();
List<PlainSelect> plainSelectList = new ArrayList<>();
plainSelectList.add(parenthesedSelect.getPlainSelect());
List<PlainSelect> subPlainSelects =
SqlSelectHelper.getPlainSelects(plainSelectList);
for (PlainSelect subPlainSelect : subPlainSelects) {
replaceFieldsInPlainOneSelect(fieldNameMap, exactReplace, subPlainSelect);
}
private static void replaceFieldsInSetOperationList(Map<String, String> fieldNameMap,
boolean exactReplace, SetOperationList operationList) {
List<OrderByElement> orderByElements = operationList.getOrderByElements();
if (!CollectionUtils.isEmpty(orderByElements)) {
for (OrderByElement orderByElement : orderByElements) {
orderByElement.accept(new OrderByReplaceVisitor(fieldNameMap, exactReplace));
}
}
}

104
common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlSelectHelper.java
View File

@@ -26,6 +26,7 @@ import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.select.FromItem;
import net.sf.jsqlparser.statement.select.GroupByElement;
import net.sf.jsqlparser.statement.select.Join;
import net.sf.jsqlparser.statement.select.LateralView;
@@ -50,7 +51,9 @@ import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;
/** Sql Parser Select Helper */
/**
* Sql Parser Select Helper
*/
@Slf4j
public class SqlSelectHelper {
@@ -808,4 +811,103 @@ public class SqlSelectHelper {
}
}
}
public static Set<Select> getAllSelect(Select selectStatement) {
Set<Select> selects = new HashSet<>();
collectSelects(selectStatement, selects);
return selects;
}
private static void collectSelects(Select select, Set<Select> selects) {
if (select == null) {
return;
}
if (select instanceof PlainSelect) {
PlainSelect plainSelect = (PlainSelect) select;
selects.add(plainSelect);
collectFromItemPlainSelects(plainSelect.getFromItem(), selects);
collectWithItemPlainSelects(plainSelect.getWithItemsList(), selects);
collectJoinsPlainSelects(plainSelect.getJoins(), selects);
collectNestedPlainSelects(plainSelect, selects);
} else if (select instanceof SetOperationList) {
SetOperationList setOperationList = (SetOperationList) select;
selects.add(setOperationList);
if (!CollectionUtils.isEmpty(setOperationList.getSelects())) {
for (Select subSelectBody : setOperationList.getSelects()) {
collectSelects(subSelectBody, selects);
}
}
} else if (select instanceof WithItem) {
WithItem withItem = (WithItem) select;
collectSelects(withItem.getSelect(), selects);
} else if (select instanceof ParenthesedSelect) {
ParenthesedSelect parenthesedSelect = (ParenthesedSelect) select;
collectSelects(parenthesedSelect.getPlainSelect(), selects);
}
}
private static void collectJoinsPlainSelects(List<Join> joins, Set<Select> selects) {
if (CollectionUtils.isEmpty(joins)) {
return;
}
for (Join join : joins) {
FromItem rightItem = join.getRightItem();
if (!(rightItem instanceof ParenthesedSelect)) {
continue;
}
ParenthesedSelect parenthesedSelect = (ParenthesedSelect) rightItem;
selects.add(parenthesedSelect.getPlainSelect());
}
}
private static void collectFromItemPlainSelects(FromItem fromItem, Set<Select> selects) {
if (fromItem instanceof ParenthesedSelect) {
ParenthesedSelect parenthesedSelect = (ParenthesedSelect) fromItem;
collectSelects(parenthesedSelect.getSelect(), selects);
}
}
public static void collectWithItemPlainSelects(List<WithItem> withItemList,
Set<Select> selects) {
if (CollectionUtils.isEmpty(withItemList)) {
return;
}
for (WithItem withItem : withItemList) {
collectSelects(withItem.getSelect(), selects);
}
}
private static void collectNestedPlainSelects(PlainSelect plainSelect, Set<Select> selects) {
ExpressionVisitorAdapter expressionVisitor = new ExpressionVisitorAdapter() {
@Override
public void visit(Select subSelect) {
if (subSelect instanceof ParenthesedSelect) {
ParenthesedSelect parenthesedSelect = (ParenthesedSelect) subSelect;
if (parenthesedSelect.getSelect() instanceof PlainSelect) {
selects.add(parenthesedSelect.getPlainSelect());
}
}
}
};
plainSelect.accept(new SelectVisitorAdapter() {
@Override
public void visit(PlainSelect plainSelect) {
Expression whereExpression = plainSelect.getWhere();
if (whereExpression != null) {
whereExpression.accept(expressionVisitor);
}
Expression having = plainSelect.getHaving();
if (Objects.nonNull(having)) {
having.accept(expressionVisitor);
}
List<SelectItem<?>> selectItems = plainSelect.getSelectItems();
if (!CollectionUtils.isEmpty(selectItems)) {
for (SelectItem selectItem : selectItems) {
selectItem.accept(expressionVisitor);
}
}
}
});
}
}

17
common/src/test/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceFieldsTest.java
View File

@@ -263,4 +263,21 @@ class SqlReplaceFieldsTest extends SqlReplaceHelperTest {
+ "SELECT * FROM daily_visits", replaceSql);
}
@Test
void testReplaceFields18() {
String replaceSql = "WITH\n" + " latest_data AS (\n" + " SELECT\n" + " 粉丝数,\n"
+ " ROW_NUMBER() OVER (\n" + " ORDER BY\n" + " 数据日期 DESC\n"
+ " ) AS __row_num__\n" + " FROM\n" + " 问答艺人数据集\n" + " WHERE\n"
+ " (TME歌手ID = '1')\n" + " AND (\n" + " 数据日期 >= '2024-10-22'\n"
+ " AND 数据日期 <= '2024-10-29'\n" + " )\n" + " )\n" + "SELECT\n"
+ " AVG(__粉丝数__)\n" + "FROM\n" + " latest_data\n" + "WHERE\n"
+ " __row_num__ = 1";
replaceSql = SqlReplaceHelper.replaceFields(replaceSql, fieldToBizName);
Assert.assertEquals("WITH latest_data AS (SELECT fans_cnt, ROW_NUMBER() OVER "
+ "(ORDER BY sys_imp_date DESC) AS __row_num__ FROM 问答艺人数据集 WHERE (TME歌手ID = '1') "
+ "AND (sys_imp_date >= '2024-10-22' AND sys_imp_date <= '2024-10-29')) SELECT AVG(__粉丝数__) "
+ "FROM latest_data WHERE __row_num__ = 1", replaceSql);
}
}

1
common/src/test/java/com/tencent/supersonic/common/jsqlparser/SqlReplaceHelperTest.java
View File

@@ -348,6 +348,7 @@ class SqlReplaceHelperTest {
fieldToBizName.put("歌曲发布时间", "song_publis_date");
fieldToBizName.put("歌曲发布年份", "song_publis_year");
fieldToBizName.put("访问次数", "pv");
fieldToBizName.put("粉丝数", "fans_cnt");
return fieldToBizName;
}