Skip to content

Commit

Permalink
[enhancement](nereids)support subquery in LogicalGenerator (#40663)
Browse files Browse the repository at this point in the history
select e1 from (select 1) t lateral view explode(**(select
sequence(CURRENT_DATE(), date_add(CURRENT_DATE(), interval 2 day)))**)
t2 as e1;

The **subquery** in explode is supported by this pr
  • Loading branch information
starocean999 authored Sep 19, 2024
1 parent cee07d6 commit 83f899b
Show file tree
Hide file tree
Showing 4 changed files with 69 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
import org.apache.doris.nereids.rules.analysis.HavingToFilter;
import org.apache.doris.nereids.rules.analysis.LeadingJoin;
import org.apache.doris.nereids.rules.analysis.NormalizeAggregate;
import org.apache.doris.nereids.rules.analysis.NormalizeGenerate;
import org.apache.doris.nereids.rules.analysis.NormalizeRepeat;
import org.apache.doris.nereids.rules.analysis.OneRowRelationExtractAggregate;
import org.apache.doris.nereids.rules.analysis.ProjectToGlobalAggregate;
Expand Down Expand Up @@ -170,6 +171,7 @@ private static List<RewriteJob> buildAnalyzerJobs(Optional<CustomTableResolver>
new CollectJoinConstraint()
),
topDown(new LeadingJoin()),
bottomUp(new NormalizeGenerate()),
bottomUp(new SubqueryToApply()),
topDown(new MergeProjects())
);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,7 @@ public enum RuleType {
NORMALIZE_AGGREGATE(RuleTypeClass.REWRITE),
NORMALIZE_SORT(RuleTypeClass.REWRITE),
NORMALIZE_REPEAT(RuleTypeClass.REWRITE),
NORMALIZE_GENERATE(RuleTypeClass.REWRITE),
EXTRACT_AND_NORMALIZE_WINDOW_EXPRESSIONS(RuleTypeClass.REWRITE),
SIMPLIFY_WINDOW_EXPRESSION(RuleTypeClass.REWRITE),
CHECK_AND_STANDARDIZE_WINDOW_FUNCTION_AND_FRAME(RuleTypeClass.REWRITE),
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.

package org.apache.doris.nereids.rules.analysis;

import org.apache.doris.nereids.rules.Rule;
import org.apache.doris.nereids.rules.RuleType;
import org.apache.doris.nereids.trees.expressions.Alias;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.SubqueryExpr;
import org.apache.doris.nereids.trees.expressions.functions.Function;
import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
import org.apache.doris.nereids.util.ExpressionUtils;

import com.google.common.collect.ImmutableList;

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

/**
* NormalizeGenerate
*/
public class NormalizeGenerate extends OneAnalysisRuleFactory {
@Override
public Rule build() {
return logicalGenerate()
.when(generate -> generate.getGenerators().stream()
.anyMatch(expr -> expr.containsType(SubqueryExpr.class)))
.then(generate -> {
List<Expression> subqueries = ExpressionUtils.collectToList(
generate.getExpressions(), SubqueryExpr.class::isInstance);
Map<Expression, Expression> replaceMap = new HashMap<>();
ImmutableList.Builder<Alias> builder = ImmutableList.builder();
for (Expression expr : subqueries) {
Alias alias = new Alias(expr);
builder.add(alias);
replaceMap.put(expr, alias.toSlot());
}
LogicalProject logicalProject = new LogicalProject(builder.build(), generate.child());
List<Function> newGenerators = new ArrayList<>(generate.getGenerators().size());
for (Function function : generate.getGenerators()) {
newGenerators.add((Function) ExpressionUtils.replace(function, replaceMap));
}
return generate.withGenerators(newGenerators).withChildren(ImmutableList.of(logicalProject));
})
.toRule(RuleType.NORMALIZE_GENERATE);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -295,4 +295,6 @@ suite("test_subquery") {
contains("partitions=3/")
}
sql """drop table if exists scalar_subquery_t"""

sql """select e1 from (select 1) t lateral view explode((select sequence(CURRENT_DATE(), date_add(CURRENT_DATE(), interval 2 day)))) t2 as e1;"""
}

0 comments on commit 83f899b

Please sign in to comment.