diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/EliminateGroupBy.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/EliminateGroupBy.java index 628bd3cd8234b7..24692b58ee89e6 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/EliminateGroupBy.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/EliminateGroupBy.java @@ -27,7 +27,6 @@ import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction; import org.apache.doris.nereids.trees.expressions.functions.agg.AnyValue; import org.apache.doris.nereids.trees.expressions.functions.agg.Avg; -import org.apache.doris.nereids.trees.expressions.functions.agg.AvgWeighted; import org.apache.doris.nereids.trees.expressions.functions.agg.Count; import org.apache.doris.nereids.trees.expressions.functions.agg.Max; import org.apache.doris.nereids.trees.expressions.functions.agg.MaxBy; @@ -66,8 +65,9 @@ public class EliminateGroupBy extends OneRewriteRuleFactory { private static final ImmutableSet> supportedBasicFunctions = ImmutableSet.of(Sum.class, Avg.class, Min.class, Max.class, Median.class, AnyValue.class); + // AvgWeighted must retain its aggregate: even one row can have zero weight or overflow v * w. private static final ImmutableSet> supportedTwoArgsFunctions - = ImmutableSet.of(MinBy.class, MaxBy.class, AvgWeighted.class, Percentile.class); + = ImmutableSet.of(MinBy.class, MaxBy.class, Percentile.class); private static final ImmutableSet> supportedDevLikeFunctions = ImmutableSet.of(Stddev.class, StddevSamp.class, Variance.class, VarianceSamp.class); private static final ImmutableSet> supportedFunctionSum0 diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/EliminateGroupByAvgWeightedTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/EliminateGroupByAvgWeightedTest.java new file mode 100644 index 00000000000000..88e015a758af65 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/EliminateGroupByAvgWeightedTest.java @@ -0,0 +1,57 @@ +// 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.rewrite; + +import org.apache.doris.nereids.trees.expressions.functions.agg.AvgWeighted; +import org.apache.doris.nereids.util.MemoPatternMatchSupported; +import org.apache.doris.nereids.util.PlanChecker; +import org.apache.doris.utframe.TestWithFeService; + +import org.junit.jupiter.api.Test; + +class EliminateGroupByAvgWeightedTest extends TestWithFeService implements MemoPatternMatchSupported { + @Override + protected void runBeforeAll() throws Exception { + createDatabase("eliminate_group_by_avg_weighted"); + connectContext.setDatabase("eliminate_group_by_avg_weighted"); + connectContext.getSessionVariable().setDisableNereidsRules("PRUNE_EMPTY_PARTITION"); + createTable("CREATE TABLE t (k INT NOT NULL, v DOUBLE NOT NULL, w DOUBLE NOT NULL)" + + " UNIQUE KEY(k) DISTRIBUTED BY HASH(k) BUCKETS 1" + + " PROPERTIES ('replication_num' = '1')"); + } + + @Test + void retainAvgWeighted() { + for (String weight : new String[] {"w", "0", "2", "-2"}) { + PlanChecker.from(connectContext) + .analyze("SELECT k, avg_weighted(v, " + weight + ") FROM t GROUP BY k") + .rewrite() + .matches(logicalAggregate().when(agg -> agg.getAggregateFunctions().stream() + .anyMatch(function -> function instanceof AvgWeighted))); + } + } + + @Test + void retainAvgWeightedWithOtherAggregates() { + PlanChecker.from(connectContext) + .analyze("SELECT k, avg_weighted(v, w), sum(v), max(v) FROM t GROUP BY k") + .rewrite() + .matches(logicalAggregate().when(agg -> agg.getAggregateFunctions().stream() + .anyMatch(function -> function instanceof AvgWeighted))); + } +} diff --git a/regression-test/data/nereids_rules_p0/eliminate_gby_key/eliminate_group_by.out b/regression-test/data/nereids_rules_p0/eliminate_gby_key/eliminate_group_by.out index 5c264971256ff5..037e9672190ef8 100644 --- a/regression-test/data/nereids_rules_p0/eliminate_gby_key/eliminate_group_by.out +++ b/regression-test/data/nereids_rules_p0/eliminate_gby_key/eliminate_group_by.out @@ -151,8 +151,10 @@ PhysicalResultSink -- !two_args_func_avg_weighted_shape -- PhysicalResultSink --PhysicalProject -----filter((test_unique2.__DORIS_DELETE_SIGN__ = 0)) -------PhysicalOlapScan[test_unique2] +----hashAgg[GLOBAL] +------PhysicalProject +--------filter((test_unique2.__DORIS_DELETE_SIGN__ = 0)) +----------PhysicalOlapScan[test_unique2] -- !two_args_func_percentile_shape -- PhysicalResultSink