From d953e5c8f443b2bc58e8eb1174306b0e5ef353aa Mon Sep 17 00:00:00 2001 From: jakevin Date: Fri, 27 Oct 2023 18:45:21 +0800 Subject: [PATCH] [feature](Nereids): Convert topn(x, 1) to max(x) (#26004) --- .../expression/ExpressionOptimization.java | 4 +- .../rules/expression/rules/TopnToMax.java | 55 +++++++++++++++++++ .../rules/expression/rules/TopnToMaxTest.java | 40 ++++++++++++++ .../nereids_p0/expression/topn_to_max.out | 8 +++ .../nereids_p0/expression/topn_to_max.groovy | 49 +++++++++++++++++ 5 files changed, 155 insertions(+), 1 deletion(-) create mode 100644 fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/TopnToMax.java create mode 100644 fe/fe-core/src/test/java/org/apache/doris/nereids/rules/expression/rules/TopnToMaxTest.java create mode 100644 regression-test/data/nereids_p0/expression/topn_to_max.out create mode 100644 regression-test/suites/nereids_p0/expression/topn_to_max.groovy diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/ExpressionOptimization.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/ExpressionOptimization.java index 869cd6b99e..a37e6f4754 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/ExpressionOptimization.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/ExpressionOptimization.java @@ -26,6 +26,7 @@ import org.apache.doris.nereids.rules.expression.rules.SimplifyComparisonPredica import org.apache.doris.nereids.rules.expression.rules.SimplifyDecimalV3Comparison; import org.apache.doris.nereids.rules.expression.rules.SimplifyInPredicate; import org.apache.doris.nereids.rules.expression.rules.SimplifyRange; +import org.apache.doris.nereids.rules.expression.rules.TopnToMax; import com.google.common.collect.ImmutableList; @@ -44,7 +45,8 @@ public class ExpressionOptimization extends ExpressionRewrite { SimplifyRange.INSTANCE, OrToIn.INSTANCE, ArrayContainToArrayOverlap.INSTANCE, - CaseWhenToIf.INSTANCE + CaseWhenToIf.INSTANCE, + TopnToMax.INSTANCE ); private static final ExpressionRuleExecutor EXECUTOR = new ExpressionRuleExecutor(OPTIMIZE_REWRITE_RULES); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/TopnToMax.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/TopnToMax.java new file mode 100644 index 0000000000..30e76cfe22 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/TopnToMax.java @@ -0,0 +1,55 @@ +// 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.expression.rules; + +import org.apache.doris.nereids.rules.expression.ExpressionRewriteContext; +import org.apache.doris.nereids.rules.expression.ExpressionRewriteRule; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction; +import org.apache.doris.nereids.trees.expressions.functions.agg.Max; +import org.apache.doris.nereids.trees.expressions.functions.agg.TopN; +import org.apache.doris.nereids.trees.expressions.literal.IntegerLikeLiteral; +import org.apache.doris.nereids.trees.expressions.visitor.DefaultExpressionRewriter; + +/** + * Convert topn(x, 1) to max(x) + */ +public class TopnToMax extends DefaultExpressionRewriter implements + ExpressionRewriteRule { + + public static final TopnToMax INSTANCE = new TopnToMax(); + + @Override + public Expression rewrite(Expression expr, ExpressionRewriteContext ctx) { + return expr.accept(this, null); + } + + @Override + public Expression visitAggregateFunction(AggregateFunction aggregateFunction, ExpressionRewriteContext context) { + if (!(aggregateFunction instanceof TopN)) { + return aggregateFunction; + } + TopN topN = (TopN) aggregateFunction; + if (topN.arity() == 2 && topN.child(1) instanceof IntegerLikeLiteral + && ((IntegerLikeLiteral) topN.child(1)).getIntValue() == 1) { + return new Max(topN.child(0)); + } else { + return aggregateFunction; + } + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/expression/rules/TopnToMaxTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/expression/rules/TopnToMaxTest.java new file mode 100644 index 0000000000..8f3c682de5 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/expression/rules/TopnToMaxTest.java @@ -0,0 +1,40 @@ +// 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.expression.rules; + +import org.apache.doris.nereids.rules.expression.ExpressionRewriteTestHelper; +import org.apache.doris.nereids.rules.expression.ExpressionRuleExecutor; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.functions.agg.Max; +import org.apache.doris.nereids.trees.expressions.functions.agg.TopN; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.types.StringType; + +import com.google.common.collect.ImmutableList; +import org.junit.jupiter.api.Test; + +class TopnToMaxTest extends ExpressionRewriteTestHelper { + @Test + void testSimplifyComparisonPredicateRule() { + executor = new ExpressionRuleExecutor(ImmutableList.of(TopnToMax.INSTANCE)); + + Slot slot = new SlotReference("a", StringType.INSTANCE); + assertRewrite(new TopN(slot, Literal.of(1)), new Max(slot)); + } +} diff --git a/regression-test/data/nereids_p0/expression/topn_to_max.out b/regression-test/data/nereids_p0/expression/topn_to_max.out new file mode 100644 index 0000000000..6c8d190500 --- /dev/null +++ b/regression-test/data/nereids_p0/expression/topn_to_max.out @@ -0,0 +1,8 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !sql -- +1 1 +2 2 + +-- !sql -- +2 + diff --git a/regression-test/suites/nereids_p0/expression/topn_to_max.groovy b/regression-test/suites/nereids_p0/expression/topn_to_max.groovy new file mode 100644 index 0000000000..ae848b5a24 --- /dev/null +++ b/regression-test/suites/nereids_p0/expression/topn_to_max.groovy @@ -0,0 +1,49 @@ +// 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. + +suite("test_topn_to_max") { + sql 'set enable_nereids_planner=true' + sql 'set enable_fallback_to_original_planner=false' + + sql 'drop table if exists test_topn_to_max;' + + sql '''create table test_topn_to_max (k1 int, k2 string) distributed by hash(k1) buckets 3 properties('replication_num' = '1');''' + sql '''insert into test_topn_to_max values (1, "1"), (2, "2");''' + + + order_qt_sql ''' + select k1, topn(k2, 1) + from test_topn_to_max + group by k1; + ''' + res = sql ''' + explain rewritten plan select k1, max(k2) + from test_topn_to_max + group by k1; + ''' + assertTrue(res.toString().contains("max")) + + order_qt_sql ''' + select topn(k2, 1) + from test_topn_to_max; + ''' + res = sql ''' + explain rewritten plan select max(k2) + from test_topn_to_max; + ''' + assertTrue(res.toString().contains("max")) +}