From 83f899b32baef67b859cab902e7da965108b8c97 Mon Sep 17 00:00:00 2001 From: starocean999 <40539150+starocean999@users.noreply.github.com> Date: Thu, 19 Sep 2024 09:38:48 +0800 Subject: [PATCH] [enhancement](nereids)support subquery in LogicalGenerator (#40663) 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 --- .../doris/nereids/jobs/executor/Analyzer.java | 2 + .../apache/doris/nereids/rules/RuleType.java | 1 + .../rules/analysis/NormalizeGenerate.java | 64 +++++++++++++++++++ .../nereids_p0/subquery/test_subquery.groovy | 2 + 4 files changed, 69 insertions(+) create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/NormalizeGenerate.java diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Analyzer.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Analyzer.java index 1ffbac97d741a4..6f6c022117c337 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Analyzer.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Analyzer.java @@ -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; @@ -170,6 +171,7 @@ private static List buildAnalyzerJobs(Optional new CollectJoinConstraint() ), topDown(new LeadingJoin()), + bottomUp(new NormalizeGenerate()), bottomUp(new SubqueryToApply()), topDown(new MergeProjects()) ); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java index d345d9057e9b43..67bdef4ef85d65 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java @@ -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), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/NormalizeGenerate.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/NormalizeGenerate.java new file mode 100644 index 00000000000000..200dc04630cec0 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/NormalizeGenerate.java @@ -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 subqueries = ExpressionUtils.collectToList( + generate.getExpressions(), SubqueryExpr.class::isInstance); + Map replaceMap = new HashMap<>(); + ImmutableList.Builder 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 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); + } +} diff --git a/regression-test/suites/nereids_p0/subquery/test_subquery.groovy b/regression-test/suites/nereids_p0/subquery/test_subquery.groovy index c8121d03b312b5..82b858cf10be78 100644 --- a/regression-test/suites/nereids_p0/subquery/test_subquery.groovy +++ b/regression-test/suites/nereids_p0/subquery/test_subquery.groovy @@ -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;""" }