[fix](aggregate) Preserve AVG accumulator width in distinct rewrite

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()) {

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<LogicalOneRowRelation> aggregate = new LogicalAggregate<>(
+                ImmutableList.of(), ImmutableList.of(avgOutput, countOutput), relation);
+
+        PlanChecker.from(MemoTestUtils.createConnectContext(), aggregate)
+                .applyTopDown(new AvgDistinctToSumDivCount())
+                .matches(logicalAggregate().when(rewritten -> {
+                    Set<AggregateFunction> 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;
+                }));
+    }
+}

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

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
+    """
 }