starocean999 commented on code in PR #33630:
URL: https://github.com/apache/doris/pull/33630#discussion_r1579206615
##########
fe/fe-core/src/main/java/org/apache/doris/analysis/StmtRewriter.java:
##########
@@ -1365,4 +1371,166 @@ public static boolean rewriteByPolicy(StatementBase
statementBase, Analyzer anal
}
return reAnalyze;
}
+
+ /**
+ *
+ * @param column the column of SlotRef
+ * @param selectList new selectList for selectStmt
+ * @param groupByExprs group by Exprs for selectStmt
+ * @return true if ref can be rewritten
+ */
+ private static boolean rewriteSelectList(Column column, SelectList
selectList, ArrayList<Expr> groupByExprs) {
+ SlotRef slot = new SlotRef(null, column.getName());
+ if (column.isKey()) {
+ selectList.addItem(new SelectListItem(slot, column.getName()));
+ groupByExprs.add(slot);
+ return true;
+ } else {
+ AggregateType aggregateType = column.getAggregationType();
+ if (aggregateType != AggregateType.SUM && aggregateType !=
AggregateType.MAX
+ && aggregateType != AggregateType.MIN) {
+ return false;
+ } else {
+ FunctionName funcName = new
FunctionName(aggregateType.toString().toLowerCase());
+ List<Expr> arrayList = Lists.newArrayList(slot);
+ FunctionCallExpr func = new FunctionCallExpr(funcName, new
FunctionParams(false, arrayList));
+ selectList.addItem(new SelectListItem(func, column.getName()));
+ return true;
+ }
+ }
+ }
+
+ /**
+ * rewrite stmt for querying random distributed table, construct an
aggregation node for pre-agg
+ * * CREATE TABLE `tbl` (
+ * `k1` BIGINT NULL DEFAULT "10",
+ * `k3` SMALLINT NULL,
+ * `a` BIGINT SUM NULL DEFAULT "0"
+ * ) ENGINE=OLAP
+ * AGGREGATE KEY(`k1`, `k2`)
+ * DISTRIBUTED BY RANDOM BUCKETS 1
+ * PROPERTIES (
+ * "replication_allocation" = "tag.location.default: 1"
+ * )
+ * e.g.,
+ * original: select * from tbl
+ * rewrite: select * from (select k1, k2, sum(pv) from tbl group by k1,
k2) t
+ * do not rewrite if no need two phase agg:
+ * e.g.,
+ * 1. select max(k1) from tbl
+ * 2. select sum(a) from tbl
+ *
+ * @param statementBase stmt to rewrite
+ * @param analyzer the analyzer
+ * @return true if rewritten
+ * @throws UserException
+ */
+ public static boolean rewriteForRandomDistribution(StatementBase
statementBase, Analyzer analyzer)
+ throws UserException {
+ boolean reAnalyze = false;
+ if (!(statementBase instanceof SelectStmt)) {
+ return false;
+ }
+ SelectStmt selectStmt = (SelectStmt) statementBase;
+ for (int i = 0; i < selectStmt.fromClause.size(); i++) {
+ TableRef tableRef = selectStmt.fromClause.get(i);
+ // Recursively rewrite subquery
+ if (tableRef instanceof InlineViewRef) {
+ InlineViewRef viewRef = (InlineViewRef) tableRef;
+ if (rewriteForRandomDistribution(viewRef.getQueryStmt(),
viewRef.getAnalyzer())) {
+ reAnalyze = true;
+ }
+ continue;
+ }
+ TableIf table = tableRef.getTable();
+ if (!(table instanceof OlapTable)) {
+ continue;
+ }
+ // only rewrite random distributed AGG_KEY table
+ OlapTable olapTable = (OlapTable) table;
+ if (olapTable.getKeysType() != KeysType.AGG_KEYS) {
+ continue;
+ }
+ DistributionInfo distributionInfo =
olapTable.getDefaultDistributionInfo();
+ if (distributionInfo.getType() !=
DistributionInfo.DistributionInfoType.RANDOM) {
+ continue;
+ }
+
+ // check agg function and column agg type
+ boolean aggTypeMatch = true;
+ if (selectStmt.getAggInfo() != null) {
+ ArrayList<FunctionCallExpr> aggExprs =
selectStmt.getAggInfo().getAggregateExprs();
+ if (aggExprs.stream().allMatch(expr ->
aggTypeMatch(expr.getFnName().getFunction(), expr))) {
+ continue;
+ }
+ aggTypeMatch = false;
+ }
+ // construct a new InlineViewRef for pre-agg
+ boolean canRewrite = true;
+ SelectList selectList = new SelectList();
+ ArrayList<Expr> groupingExprs = new ArrayList<>();
+ TupleDescriptor desc = tableRef.getDesc();
+ List<Column> columns =
desc.getSlots().stream().map(SlotDescriptor::getColumn).collect(Collectors.toList());
+ columns = columns.isEmpty() || !aggTypeMatch ?
olapTable.getBaseSchema() : columns;
+ for (Column col : columns) {
+ if (!rewriteSelectList(col, selectList, groupingExprs)) {
+ canRewrite = false;
+ break;
+ }
+ }
+ if (!canRewrite) {
+ continue;
+ }
+ Expr whereClause = selectStmt.getWhereClause() == null ? null :
selectStmt.getWhereClause().clone();
+ SelectStmt newSelectSmt = new SelectStmt(selectList,
+ new FromClause(Lists.newArrayList(tableRef)),
+ whereClause,
+ new GroupByClause(groupingExprs,
GroupByClause.GroupingType.GROUP_BY),
+ null,
+ null,
+ LimitElement.NO_LIMIT);
+ InlineViewRef inlineViewRef = new
InlineViewRef(tableRef.getAliasAsName().getTbl(), newSelectSmt);
+ inlineViewRef.setJoinOp(tableRef.getJoinOp());
+ inlineViewRef.setLeftTblRef(tableRef.getLeftTblRef());
+ inlineViewRef.setOnClause(tableRef.getOnClause());
+ tableRef.setOnClause(null);
+ tableRef.setLeftTblRef(null);
+ tableRef.setOnClause(null);
+ if (selectStmt.fromClause.size() > i + 1) {
+ selectStmt.fromClause.get(i + 1).setLeftTblRef(inlineViewRef);
+ }
+ selectStmt.fromClause.set(i, inlineViewRef);
+ selectStmt.analyze(analyzer);
+ reAnalyze = true;
+ }
+ return reAnalyze;
+ }
+
+ /**
+ * check if the agg type of functionCall match the agg type of column
+ * @param functionName the functionName of functionCall
+ * @param expr FunctionCallExpr
+ * @return true if agg type match
+ */
+ private static boolean aggTypeMatch(String functionName, Expr expr) {
+ if (expr.getChildren().isEmpty()) {
+ if (expr instanceof SlotRef) {
Review Comment:
if expr is FunctionCallExpr, why needs if (expr instanceof SlotRef) ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org
For additional commands, e-mail: commits-h...@doris.apache.org
