From 1c1e9c66b43654d49e15c8bcbb76e01e2dc0e343 Mon Sep 17 00:00:00 2001 From: morrySnow Date: Wed, 9 Sep 2026 19:11:43 +0800 Subject: [PATCH] [fix](nereids) Preserve AVG accumulator width in distinct rewrite Problem: The multi-distinct rewrite decomposed AVG(DISTINCT BIGINT) into SUM and COUNT, but SUM(BIGINT) can overflow before the result is converted to AVG's return type. Fix: Losslessly widen the shared distinct argument to LARGEINT before creating SUM and COUNT so the rewritten expression preserves AVG's accumulator width and the argument sharing required by multi-distinct aggregation. Tests: - add focused rewrite coverage for the widened shared argument - add a boundary-value regression for AVG(DISTINCT BIGINT) Issue Number: None --- .../analysis/AvgDistinctToSumDivCount.java | 12 +++- .../AvgDistinctToSumDivCountTest.java | 72 +++++++++++++++++++ .../avg_distinct_to_sum_div_count.out | 2 + .../avg_distinct_to_sum_div_count.groovy | 13 ++++ 4 files changed, 97 insertions(+), 2 deletions(-) create mode 100644 fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/AvgDistinctToSumDivCountTest.java diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/AvgDistinctToSumDivCount.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/AvgDistinctToSumDivCount.java index 2740e8269763c2..581d58c585f5a6 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/AvgDistinctToSumDivCount.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/AvgDistinctToSumDivCount.java @@ -29,6 +29,7 @@ import org.apache.doris.nereids.trees.expressions.functions.agg.Sum; import org.apache.doris.nereids.trees.expressions.functions.scalar.NonNullable; import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; +import org.apache.doris.nereids.types.LargeIntType; import org.apache.doris.nereids.util.ExpressionUtils; import org.apache.doris.nereids.util.TypeCoercionUtils; @@ -52,11 +53,18 @@ public Rule build() { .stream() .filter(function -> function instanceof Avg && function.isDistinct()) .collect(ImmutableMap.toImmutableMap(function -> function, function -> { + Expression argument = ((Avg) function).child(); + // AVG(BIGINT) accumulates in LARGEINT, while SUM(BIGINT) accumulates + // in BIGINT. Preserve AVG's wider accumulator when decomposing it. + if (argument.getDataType().isBigIntType()) { + argument = TypeCoercionUtils.castIfNotSameType( + argument, LargeIntType.INSTANCE); + } Sum sum = (Sum) TypeCoercionUtils.processBoundFunction( new Sum(true, ((Avg) function).isAlwaysNullable(), false, - ((Avg) function).child())); + argument)); Count count = (Count) TypeCoercionUtils.processBoundFunction( - new Count(true, ((Avg) function).child())); + new Count(true, argument)); Expression divide = TypeCoercionUtils.castIfNotSameType(TypeCoercionUtils.processDivide( new Divide(sum, count)), function.getDataType()); if (!function.nullable() && divide.nullable()) { diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/AvgDistinctToSumDivCountTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/AvgDistinctToSumDivCountTest.java new file mode 100644 index 00000000000000..6deb4b2aee0a77 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/AvgDistinctToSumDivCountTest.java @@ -0,0 +1,72 @@ +// 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.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator; +import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction; +import org.apache.doris.nereids.trees.expressions.functions.agg.Avg; +import org.apache.doris.nereids.trees.expressions.functions.agg.Count; +import org.apache.doris.nereids.trees.expressions.functions.agg.Sum; +import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; +import org.apache.doris.nereids.trees.plans.logical.LogicalOneRowRelation; +import org.apache.doris.nereids.types.BigIntType; +import org.apache.doris.nereids.types.IntegerType; +import org.apache.doris.nereids.util.MemoPatternMatchSupported; +import org.apache.doris.nereids.util.MemoTestUtils; +import org.apache.doris.nereids.util.PlanChecker; +import org.apache.doris.nereids.util.TypeCoercionUtils; + +import com.google.common.collect.ImmutableList; +import org.junit.jupiter.api.Test; + +import java.util.Set; + +public class AvgDistinctToSumDivCountTest implements MemoPatternMatchSupported { + @Test + public void testBigIntUsesLargeIntAccumulator() { + SlotReference value = new SlotReference("value", BigIntType.INSTANCE, false); + SlotReference other = new SlotReference("other", IntegerType.INSTANCE, false); + LogicalOneRowRelation relation = new LogicalOneRowRelation( + StatementScopeIdGenerator.newRelationId(), ImmutableList.of(value, other)); + + Avg avg = (Avg) TypeCoercionUtils.processBoundFunction(new Avg(true, value)); + Count count = (Count) TypeCoercionUtils.processBoundFunction(new Count(true, other)); + NamedExpression avgOutput = new Alias(avg, "average"); + NamedExpression countOutput = new Alias(count, "count"); + LogicalAggregate aggregate = new LogicalAggregate<>( + ImmutableList.of(), ImmutableList.of(avgOutput, countOutput), relation); + + PlanChecker.from(MemoTestUtils.createConnectContext(), aggregate) + .applyTopDown(new AvgDistinctToSumDivCount()) + .matches(logicalAggregate().when(rewritten -> { + Set functions = rewritten.getAggregateFunctions(); + Sum sum = (Sum) functions.stream() + .filter(Sum.class::isInstance) + .findFirst() + .orElseThrow(() -> new AssertionError("rewritten SUM is missing")); + boolean countSharesWidenedArgument = functions.stream() + .filter(Count.class::isInstance) + .anyMatch(function -> function.child(0).equals(sum.child(0))); + return sum.child(0).getDataType().isLargeIntType() + && countSharesWidenedArgument; + })); + } +} diff --git a/regression-test/data/nereids_rules_p0/avg_distinct_to_sum_div_count/avg_distinct_to_sum_div_count.out b/regression-test/data/nereids_rules_p0/avg_distinct_to_sum_div_count/avg_distinct_to_sum_div_count.out index 38654c9eccaefb..a0f210574beb5f 100644 --- a/regression-test/data/nereids_rules_p0/avg_distinct_to_sum_div_count/avg_distinct_to_sum_div_count.out +++ b/regression-test/data/nereids_rules_p0/avg_distinct_to_sum_div_count/avg_distinct_to_sum_div_count.out @@ -2,3 +2,5 @@ -- !ctas -- 135.55 67.7750 +-- !bigint_avg_distinct_overflow -- +1 2 diff --git a/regression-test/suites/nereids_rules_p0/avg_distinct_to_sum_div_count/avg_distinct_to_sum_div_count.groovy b/regression-test/suites/nereids_rules_p0/avg_distinct_to_sum_div_count/avg_distinct_to_sum_div_count.groovy index 826bda8fd8bae7..a57c4ad8884eab 100644 --- a/regression-test/suites/nereids_rules_p0/avg_distinct_to_sum_div_count/avg_distinct_to_sum_div_count.groovy +++ b/regression-test/suites/nereids_rules_p0/avg_distinct_to_sum_div_count/avg_distinct_to_sum_div_count.groovy @@ -26,4 +26,17 @@ suite("avg_distinct_to_sum_div_count") { sql "drop table if exists testctas;" sql """create table testctas properties("replication_num" = "1") select sum(distinct c1),avg(distinct c1) from t1;""" qt_ctas "select * from testctas;" + + qt_bigint_avg_distinct_overflow """ + SELECT COUNT(*) AS matched_rows, MAX(c) AS distinct_count + FROM ( + SELECT AVG(DISTINCT x) AS a, COUNT(DISTINCT y) AS c + FROM ( + SELECT CAST(9223372036854775807 AS BIGINT) AS x, CAST(1 AS INT) AS y + UNION ALL + SELECT CAST(9223372036854775806 AS BIGINT) AS x, CAST(2 AS INT) AS y + ) t + ) q + WHERE a > 0 + """ }