topdog/supersonic
1
0
Fork 0
mirror of https://github.com/tencentmusic/supersonic.git synced 2025-12-14 05:43:51 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

[improvement][headless] Support merging SQL queries with WITH using the LIMIT clause (#1795)

Browse Source
This commit is contained in:
lexluo09
2024-10-13 21:37:34 +08:00
committed by GitHub
parent 6d993b4785
commit 28b7847389
2 changed files with 135 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
common/src/main/java/com/tencent/supersonic/common/calcite/SqlMergeWithUtils.java
View File

@@ -6,6 +6,8 @@ import org.apache.calcite.sql.SqlIdentifier;
import org.apache.calcite.sql.SqlLiteral;
import org.apache.calcite.sql.SqlNode;
import org.apache.calcite.sql.SqlNodeList;
import org.apache.calcite.sql.SqlOrderBy;
import org.apache.calcite.sql.SqlSelect;
import org.apache.calcite.sql.SqlWith;
import org.apache.calcite.sql.SqlWithItem;
import org.apache.calcite.sql.SqlWriterConfig;
@@ -42,30 +44,46 @@ public class SqlMergeWithUtils {
// Create a new WITH item for parentWithName without quotes
SqlWithItem withItem = new SqlWithItem(SqlParserPos.ZERO,
new SqlIdentifier(parentWithName, SqlParserPos.ZERO), // false
// to
// avoid
// quotes
null, sqlNode2, SqlLiteral.createBoolean(false, SqlParserPos.ZERO));
new SqlIdentifier(parentWithName, SqlParserPos.ZERO), null, sqlNode2,
SqlLiteral.createBoolean(false, SqlParserPos.ZERO));
// Add the new WITH item to the list
withItemList.add(withItem);
}
// Check if the main SQL node contains a LIMIT clause
SqlNode limitNode = null;
if (sqlNode1 instanceof SqlOrderBy) {
SqlOrderBy sqlOrderBy = (SqlOrderBy) sqlNode1;
limitNode = sqlOrderBy.fetch;
sqlNode1 = sqlOrderBy.query;
} else if (sqlNode1 instanceof SqlSelect) {
SqlSelect sqlSelect = (SqlSelect) sqlNode1;
limitNode = sqlSelect.getFetch();
sqlSelect.setFetch(null);
sqlNode1 = sqlSelect;
}
// Extract existing WITH items from sqlNode1 if it is a SqlWith
if (sqlNode1 instanceof SqlWith) {
SqlWith sqlWith = (SqlWith) sqlNode1;
withItemList.addAll(sqlWith.withList.getList());
sqlNode1 = sqlWith.body;
}
// Create a new SqlWith node
SqlWith finalSqlNode = new SqlWith(SqlParserPos.ZERO,
new SqlNodeList(withItemList, SqlParserPos.ZERO), sqlNode1);
// If there was a LIMIT clause, wrap the finalSqlNode in a SqlOrderBy with the LIMIT
SqlNode resultNode = finalSqlNode;
if (limitNode != null) {
resultNode = new SqlOrderBy(SqlParserPos.ZERO, finalSqlNode, SqlNodeList.EMPTY, null,
limitNode);
}
// Custom SqlPrettyWriter configuration to avoid quoting identifiers
SqlWriterConfig config = Configuration.getSqlWriterConfig(engineType);
// Pretty print the final SQL
SqlPrettyWriter writer = new SqlPrettyWriter(config);
return writer.format(finalSqlNode);
return writer.format(resultNode);
}
}