Skip to content

Fix the scope analysis of the CTEs #729

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

Merged
merged 1 commit into from
Aug 6, 2024
Merged

Fix the scope analysis of the CTEs #729

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
15 changes: 8 additions & 7 deletions wren-base/src/main/java/io/wren/base/sqlrewrite/analyzer/StatementAnalyzer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -160,7 +160,8 @@ protected Scope visitTable(Table node, Optional<Scope> scope)
if (withQuery.isPresent()) {
// currently we only care about the table that is actually a model instead of a alias table that use cte table
// return empty scope here.
Scope outputScope = createScopeForCommonTableExpression(node, withQuery.get(), scope);
Optional<Scope> withScope = analysis.tryGetScope(withQuery.get().getQuery());
Scope outputScope = createScopeForCommonTableExpression(node, withQuery.get(), withScope);
analysis.setScope(node, outputScope);
return outputScope;
}
Expand Down Expand Up @@ -202,8 +203,7 @@ protected Scope visitTable(Table node, Optional<Scope> scope)
private Scope createScopeForCommonTableExpression(Table table, WithQuery withQuery, Optional<Scope> scope)
{
Query query = withQuery.getQuery();
Analysis analyzed = new Analysis(query);
Optional<Scope> queryScope = Optional.ofNullable(analyze(analyzed, query, sessionContext, wrenMDL));
Optional<Scope> queryScope = analysis.tryGetScope(query);
List<Field> fields;
Optional<List<Identifier>> columnNames = withQuery.getColumnNames();
if (columnNames.isPresent()) {
Expand Down Expand Up @@ -243,7 +243,7 @@ private List<Field> createScopeForQuery(Query query, QualifiedName scopeName, Op
else {
SingleColumn singleColumn = (SingleColumn) selectItem;
String name = singleColumn.getAlias().map(Identifier::getValue)
.or(() -> Optional.ofNullable(QueryUtil.getQualifiedName(singleColumn.getExpression())).map(QualifiedName::toString))
.or(() -> Optional.ofNullable(QueryUtil.getQualifiedName(singleColumn.getExpression()).getSuffix()))
.orElse(singleColumn.getExpression().toString());
if (scope.isPresent()) {
Optional<Field> fieldOptional = scope.get().getRelationType().resolveAnyField(QueryUtil.getQualifiedName(singleColumn.getExpression()));
Expand Down Expand Up @@ -402,7 +402,7 @@ private Scope analyzeFrom(QuerySpecification node, Optional<Scope> scope)

private void analyzeWhere(Expression node, Scope scope)
{
ExpressionAnalysis expressionAnalysis = analyzeExpression(node, scope);
analyzeExpression(node, scope);
}

private void analyzeWindowSpecification(WindowSpecification windowSpecification, Scope scope)
Expand Down Expand Up @@ -561,7 +561,7 @@ protected Scope visitTableSubquery(TableSubquery node, Optional<Scope> scope)
private Optional<Scope> analyzeWith(Query node, Optional<Scope> scope)
{
if (node.getWith().isEmpty()) {
return Optional.empty();
return scope.map(s -> Scope.builder().parent(Optional.of(s)).build());
}

With with = node.getWith().get();
Expand All @@ -572,7 +572,8 @@ private Optional<Scope> analyzeWith(Query node, Optional<Scope> scope)
if (withScopeBuilder.containsNamedQuery(name)) {
throw new IllegalArgumentException(format("WITH query name '%s' specified more than once", name));
}
process(withQuery.getQuery(), withScopeBuilder.build());
Scope withQueryScope = process(withQuery.getQuery(), withScopeBuilder.build());
analysis.setScope(withQuery.getQuery(), withQueryScope);
withScopeBuilder.namedQuery(name, withQuery);
}

Expand Down
52 changes: 52 additions & 0 deletions wren-base/src/test/java/io/wren/base/sqlrewrite/analyzer/TestDecisionPointAnalyzer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -676,6 +676,58 @@ WITH t1 as (SELECT * FROM customer), t2 as (SELECT * FROM orders)
new ExprSource("t1.custkey", "customer", "custkey", new NodeLocation(2, 33)),
new ExprSource("t2.custkey", "orders", "custkey", new NodeLocation(2, 50))));
}

statement = parseSql("""
WITH t1 as (SELECT customer.custkey FROM customer),
t2 as (SELECT t1.custkey FROM t1)
SELECT t2.custkey FROM t2
""");
result = DecisionPointAnalyzer.analyze(statement, DEFAULT_SESSION_CONTEXT, mdl);
assertThat(result.size()).isEqualTo(3);
QueryAnalysis queryAnalysis = result.get(2);
assertThat(queryAnalysis.getSelectItems().get(0).getExprSources().size()).isEqualTo(1);

statement = parseSql("""
WITH t1 as (SELECT customer.custkey FROM customer),
t2 as (SELECT t1.custkey as custkey_alias FROM t1)
SELECT custkey_alias FROM t2
""");
result = DecisionPointAnalyzer.analyze(statement, DEFAULT_SESSION_CONTEXT, mdl);
assertThat(result.size()).isEqualTo(3);
queryAnalysis = result.get(2);
assertThat(queryAnalysis.getSelectItems().get(0).getExprSources().size()).isEqualTo(1);

statement = parseSql("""
WITH t1 as (SELECT customer.custkey FROM customer),
t2 as (SELECT * FROM t1)
SELECT custkey FROM t2
""");
result = DecisionPointAnalyzer.analyze(statement, DEFAULT_SESSION_CONTEXT, mdl);
assertThat(result.size()).isEqualTo(3);
queryAnalysis = result.get(2);
assertThat(queryAnalysis.getSelectItems().get(0).getExprSources().size()).isEqualTo(1);

statement = parseSql("""
WITH t1 as (SELECT customer.custkey FROM customer),
t2 as (SELECT * FROM t1)
SELECT custkey FROM t2
""");
result = DecisionPointAnalyzer.analyze(statement, DEFAULT_SESSION_CONTEXT, mdl);
assertThat(result.size()).isEqualTo(3);
queryAnalysis = result.get(2);
assertThat(queryAnalysis.getSelectItems().get(0).getExprSources().size()).isEqualTo(1);


// we only analyze the top-level expression
statement = parseSql("""
WITH t1 as (SELECT customer.custkey FROM customer),
t2 as (SELECT (t1.custkey + 1) as custkey_plus FROM t1)
SELECT custkey_plus FROM t2
Comment on lines +723 to +725
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We only analyze the top-level expression. We won't dig into the child of this expression. So, the number of expression sources of custkey_plus is zero.

""");
result = DecisionPointAnalyzer.analyze(statement, DEFAULT_SESSION_CONTEXT, mdl);
assertThat(result.size()).isEqualTo(3);
queryAnalysis = result.get(2);
assertThat(queryAnalysis.getSelectItems().get(0).getExprSources().size()).isEqualTo(0);
}

@Test
Expand Down
Loading