Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Enhancement] Support Generated Column rewrite in complex Query (backport #50398) #50829

Closed
wants to merge 1 commit into from
Closed

[Enhancement] Support Generated Column rewrite in complex Query (backport #50398) #50829

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 0 additions & 11 deletions fe/fe-core/src/main/java/com/starrocks/sql/analyzer/AnalyzeState.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@
import com.starrocks.analysis.FunctionCallExpr;
import com.starrocks.analysis.LimitElement;
import com.starrocks.analysis.OrderByElement;
import com.starrocks.analysis.SlotRef;
import com.starrocks.common.IdGenerator;
import com.starrocks.sql.ast.Relation;
import com.starrocks.sql.ast.SelectRelation;
Expand Down Expand Up @@ -56,8 +55,6 @@ public class AnalyzeState {
private Scope orderScope;
private List<Expr> orderSourceExpressions;

private Map<Expr, SlotRef> generatedExprToColumnRef = new HashMap<>();

/**
* outputExprInOrderByScope is used to record which expressions in outputExpression are to be
* recorded in the first level of OrderByScope (order by expressions can refer to columns in output)
Expand Down Expand Up @@ -257,12 +254,4 @@ public ExprId getNextNondeterministicId() {
public List<Expr> getColumnNotInGroupBy() {
return columnNotInGroupBy;
}

public void setGeneratedExprToColumnRef(Map<Expr, SlotRef> generatedExprToColumnRef) {
this.generatedExprToColumnRef = generatedExprToColumnRef;
}

public Map<Expr, SlotRef> getGeneratedExprToColumnRef() {
return generatedExprToColumnRef;
}
}
183 changes: 172 additions & 11 deletions fe/fe-core/src/main/java/com/starrocks/sql/analyzer/QueryAnalyzer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,154 @@ public void analyze(StatementBase node, Scope parent) {
new Visitor().process(node, parent);
}

private class GeneratedColumnExprMappingCollector implements AstVisitor<Void, Scope> {
public GeneratedColumnExprMappingCollector() {
}

public Void process(ParseNode node, Scope scope) {
return node.accept(this, scope);
}

private void reAnalyzeExpressionBasedOnCurrentScope(SelectRelation childSelectRelation, Scope scope,
Map<Expr, SlotRef> resultGeneratedExprToColumnRef) {
if (childSelectRelation.getGeneratedExprToColumnRef() == null ||
childSelectRelation.getGeneratedExprToColumnRef().isEmpty()) {
return;
}
// 1. get all available generated column from child selectRelation
// available means that:
// a. generated column output from child selectRelation directly.
// b. all reference column of generated column output from child selectRelation directly.
List<SlotRef> outputSlotRef = childSelectRelation.getOutputExpression()
.stream().filter(e -> e instanceof SlotRef)
.map(e -> (SlotRef) e).collect(Collectors.toList());
boolean hasStar = childSelectRelation.getSelectList()
.getItems().stream().anyMatch(SelectListItem::isStar);
Map<Expr, SlotRef> generatedExprToColumnRef = new HashMap<>();
for (Map.Entry<Expr, SlotRef> entry : childSelectRelation.getGeneratedExprToColumnRef().entrySet()) {
List<SlotRef> allRefColumns = Lists.newArrayList();
entry.getKey().collect(SlotRef.class, allRefColumns);
allRefColumns.add(entry.getValue());
if (hasStar || outputSlotRef.containsAll(allRefColumns)) {
generatedExprToColumnRef.put(entry.getKey().clone(), (SlotRef) entry.getValue().clone());
}
}

// 2. rewrite(rename slotRef) generated column expression(unAnalyzed) using alias in current scope
Map<String, String> slotRefToAlias = new HashMap<>();
for (SelectListItem item : childSelectRelation.getSelectList().getItems()) {
if (item.isStar()) {
slotRefToAlias.clear();
break;
}

if (!(item.getExpr() instanceof SlotRef) || (item.getAlias() == null || item.getAlias().isEmpty())) {
continue;
}

slotRefToAlias.put(((SlotRef) item.getExpr()).toSql(), item.getAlias());
}
List<SlotRef> allRefSlotRefs = new ArrayList<>();
for (Map.Entry<Expr, SlotRef> entry : generatedExprToColumnRef.entrySet()) {
List<SlotRef> refColumns = Lists.newArrayList();
entry.getKey().collect(SlotRef.class, refColumns);

allRefSlotRefs.addAll(refColumns);
allRefSlotRefs.add(entry.getValue());
}
for (SlotRef slotRef : allRefSlotRefs) {
if (!slotRefToAlias.isEmpty()) {
String alias = slotRefToAlias.get(slotRef.toSql());
if (alias != null) {
slotRef.setColumnName(alias);
}
}
slotRef.setTblName(null);
}

// 3. analyze generated column expression based on current scope
for (Map.Entry<Expr, SlotRef> entry : generatedExprToColumnRef.entrySet()) {
entry.getKey().reset();
entry.getValue().reset();

try {
ExpressionAnalyzer.analyzeExpression(entry.getKey(), new AnalyzeState(), scope, session);
ExpressionAnalyzer.analyzeExpression(entry.getValue(), new AnalyzeState(), scope, session);
} catch (Exception ignore) {
// ignore generated column rewrite if hit any exception
generatedExprToColumnRef.clear();
}
}
resultGeneratedExprToColumnRef.putAll(generatedExprToColumnRef);
}

@Override
public Void visitTable(TableRelation tableRelation, Scope scope) {
Table table = tableRelation.getTable();
Map<Expr, SlotRef> generatedExprToColumnRef = new HashMap<>();
for (Column column : table.getBaseSchema()) {
Expr generatedColumnExpression = column.getGeneratedColumnExpr(table.getIdToColumn());
if (generatedColumnExpression != null) {
SlotRef slotRef = new SlotRef(null, column.getName());
ExpressionAnalyzer.analyzeExpression(generatedColumnExpression, new AnalyzeState(), scope, session);
ExpressionAnalyzer.analyzeExpression(slotRef, new AnalyzeState(), scope, session);
generatedExprToColumnRef.put(generatedColumnExpression, slotRef);
}
}
tableRelation.setGeneratedExprToColumnRef(generatedExprToColumnRef);
return null;
}

@Override
public Void visitSelect(SelectRelation selectRelation, Scope scope) {
selectRelation.setGeneratedExprToColumnRef(selectRelation.getRelation().getGeneratedExprToColumnRef());
return null;
}

@Override
public Void visitSubquery(SubqueryRelation subquery, Scope scope) {
QueryRelation queryRelation = subquery.getQueryStatement().getQueryRelation();
if (queryRelation instanceof SelectRelation) {
SelectRelation childSelectRelation = (SelectRelation) queryRelation;
reAnalyzeExpressionBasedOnCurrentScope(childSelectRelation, scope, subquery.getGeneratedExprToColumnRef());
}
return null;
}

@Override
public Void visitJoin(JoinRelation joinRelation, Scope scope) {
Relation leftRelation = joinRelation.getLeft();
Relation rightRelation = joinRelation.getRight();
joinRelation.getGeneratedExprToColumnRef().putAll(leftRelation.getGeneratedExprToColumnRef());
joinRelation.getGeneratedExprToColumnRef().putAll(rightRelation.getGeneratedExprToColumnRef());
return null;
}

@Override
public Void visitView(ViewRelation node, Scope scope) {
QueryRelation queryRelation = node.getQueryStatement().getQueryRelation();
if (queryRelation instanceof SubqueryRelation) {
node.setGeneratedExprToColumnRef(queryRelation.getGeneratedExprToColumnRef());
} else if (queryRelation instanceof SelectRelation) {
SelectRelation childSelectRelation = (SelectRelation) queryRelation;
reAnalyzeExpressionBasedOnCurrentScope(childSelectRelation, scope, node.getGeneratedExprToColumnRef());
}
return null;
}

@Override
public Void visitCTE(CTERelation cteRelation, Scope scope) {
QueryRelation queryRelation = cteRelation.getCteQueryStatement().getQueryRelation();
if (queryRelation instanceof SubqueryRelation) {
cteRelation.setGeneratedExprToColumnRef(queryRelation.getGeneratedExprToColumnRef());
} else if (queryRelation instanceof SelectRelation) {
SelectRelation childSelectRelation = (SelectRelation) queryRelation;
reAnalyzeExpressionBasedOnCurrentScope(childSelectRelation, scope, cteRelation.getGeneratedExprToColumnRef());
}
return null;
}
}

private class Visitor implements AstVisitor<Scope, Scope> {
public Visitor() {
}
Expand Down Expand Up @@ -217,6 +365,7 @@ public Scope visitSelect(SelectRelation selectRelation, Scope scope) {
Scope sourceScope = process(resolvedRelation, scope);
sourceScope.setParent(scope);

<<<<<<< HEAD
Map<Expr, SlotRef> generatedExprToColumnRef = new HashMap<>();
new AstVisitor<Void, Void>() {
@Override
Expand Down Expand Up @@ -319,6 +468,9 @@ public Void visitJoin(JoinRelation joinRelation, Void context) {
}.visit(resolvedRelation);
analyzeState.setGeneratedExprToColumnRef(generatedExprToColumnRef);

=======
selectRelation.accept(new RewriteAliasVisitor(sourceScope, session), null);
>>>>>>> 96b97a70e4 ([Enhancement] Support Generated Column rewrite in complex Query (#50398))
SelectAnalyzer selectAnalyzer = new SelectAnalyzer(session);
selectAnalyzer.analyze(
analyzeState,
Expand All @@ -332,6 +484,8 @@ public Void visitJoin(JoinRelation joinRelation, Void context) {
selectRelation.getLimit());

selectRelation.fillResolvedAST(analyzeState);
GeneratedColumnExprMappingCollector collector = new GeneratedColumnExprMappingCollector();
collector.process(selectRelation, sourceScope);
return analyzeState.getOutputScope();
}

Expand Down Expand Up @@ -596,17 +750,8 @@ public Scope visitTable(TableRelation node, Scope outerScope) {
Scope scope = new Scope(RelationId.of(node), new RelationFields(fields.build()));
node.setScope(scope);

Map<Expr, SlotRef> getGeneratedExprToColumnRef = new HashMap<>();
for (Column column : table.getBaseSchema()) {
Expr generatedColumnExpression = column.getGeneratedColumnExpr(table.getIdToColumn());
if (generatedColumnExpression != null) {
SlotRef slotRef = new SlotRef(null, column.getName());
ExpressionAnalyzer.analyzeExpression(generatedColumnExpression, new AnalyzeState(), scope, session);
ExpressionAnalyzer.analyzeExpression(slotRef, new AnalyzeState(), scope, session);
getGeneratedExprToColumnRef.put(generatedColumnExpression, slotRef);
}
}
node.setGeneratedExprToColumnRef(getGeneratedExprToColumnRef);
GeneratedColumnExprMappingCollector collector = new GeneratedColumnExprMappingCollector();
collector.process(node, scope);

return scope;
}
Expand Down Expand Up @@ -657,6 +802,10 @@ public Scope visitCTE(CTERelation cteRelation, Scope context) {
}
Scope scope = new Scope(RelationId.of(cteRelation), new RelationFields(outputFields.build()));
cteRelation.setScope(scope);

GeneratedColumnExprMappingCollector collector = new GeneratedColumnExprMappingCollector();
collector.process(cteRelation, scope);

return scope;
}

Expand Down Expand Up @@ -761,6 +910,10 @@ public Scope visitJoin(JoinRelation join, Scope parentScope) {
leftScope.getRelationFields().joinWith(rightScope.getRelationFields()));
}
join.setScope(scope);

GeneratedColumnExprMappingCollector collector = new GeneratedColumnExprMappingCollector();
collector.process(join, scope);

return scope;
}

Expand Down Expand Up @@ -882,6 +1035,10 @@ public Scope visitSubquery(SubqueryRelation subquery, Scope context) {

analyzeOrderByClause(subquery, scope);
subquery.setScope(scope);

GeneratedColumnExprMappingCollector collector = new GeneratedColumnExprMappingCollector();
collector.process(subquery, scope);

return scope;
}

Expand Down Expand Up @@ -959,6 +1116,10 @@ public Scope visitView(ViewRelation node, Scope scope) {

Scope viewScope = new Scope(RelationId.of(node), new RelationFields(fields));
node.setScope(viewScope);

GeneratedColumnExprMappingCollector collector = new GeneratedColumnExprMappingCollector();
collector.process(node, viewScope);

return viewScope;
}

Expand Down
18 changes: 18 additions & 0 deletions fe/fe-core/src/main/java/com/starrocks/sql/ast/Relation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,15 +14,19 @@

package com.starrocks.sql.ast;

import com.starrocks.analysis.Expr;
import com.starrocks.analysis.ParseNode;
import com.starrocks.analysis.SlotRef;
import com.starrocks.analysis.TableName;
import com.starrocks.sql.analyzer.RelationFields;
import com.starrocks.sql.analyzer.Scope;
import com.starrocks.sql.common.ErrorType;
import com.starrocks.sql.common.StarRocksPlannerException;
import com.starrocks.sql.parser.NodePosition;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

public abstract class Relation implements ParseNode {
private Scope scope;
Expand All @@ -39,6 +43,12 @@ public abstract class Relation implements ParseNode {
// generated by Security Policy rewriting does not perform permission verification.
private boolean createByPolicyRewritten = false;

/**
* generatedExprToColumnRef stores the mapping relationship
* between generated expressions and generated columns
*/
private Map<Expr, SlotRef> generatedExprToColumnRef = new HashMap<>();

protected final NodePosition pos;

protected Relation(NodePosition pos) {
Expand Down Expand Up @@ -109,6 +119,14 @@ public List<String> getExplicitColumnNames() {
return explicitColumnNames;
}

public void setGeneratedExprToColumnRef(Map<Expr, SlotRef> generatedExprToColumnRef) {
this.generatedExprToColumnRef = generatedExprToColumnRef;
}

public Map<Expr, SlotRef> getGeneratedExprToColumnRef() {
return generatedExprToColumnRef;
}

@Override
public <R, C> R accept(AstVisitor<R, C> visitor, C context) {
return visitor.visitRelation(this, context);
Expand Down
14 changes: 0 additions & 14 deletions fe/fe-core/src/main/java/com/starrocks/sql/ast/SelectRelation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,13 +20,11 @@
import com.starrocks.analysis.GroupByClause;
import com.starrocks.analysis.LimitElement;
import com.starrocks.analysis.OrderByElement;
import com.starrocks.analysis.SlotRef;
import com.starrocks.sql.analyzer.AnalyzeState;
import com.starrocks.sql.analyzer.FieldId;
import com.starrocks.sql.analyzer.Scope;
import com.starrocks.sql.parser.NodePosition;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

Expand Down Expand Up @@ -79,12 +77,6 @@ public class SelectRelation extends QueryRelation {

private Map<Expr, FieldId> columnReferences;

/**
* materializeExpressionToColumnRef stores the mapping relationship
* between generated expressions and generated columns
*/
private Map<Expr, SlotRef> generatedExprToColumnRef = new HashMap<>();

public SelectRelation(
SelectList selectList,
Relation fromRelation,
Expand Down Expand Up @@ -160,8 +152,6 @@ public void fillResolvedAST(AnalyzeState analyzeState) {

this.columnReferences = analyzeState.getColumnReferences();

this.generatedExprToColumnRef = analyzeState.getGeneratedExprToColumnRef();

this.setScope(analyzeState.getOutputScope());
}

Expand Down Expand Up @@ -307,8 +297,4 @@ public boolean hasAnalyticInfo() {
public List<Expr> getOutputExpression() {
return outputExpr;
}

public Map<Expr, SlotRef> getGeneratedExprToColumnRef() {
return generatedExprToColumnRef;
}
}
Loading
Loading