branch-4.1: [fix](fd) Guard uniform aggregate inference by participation #67881

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalAggregate.java

@@ -29,8 +29,6 @@
 import org.apache.doris.nereids.trees.expressions.functions.ExpressionTrait;
 import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
 import org.apache.doris.nereids.trees.expressions.functions.agg.AggregatePhase;
-import org.apache.doris.nereids.trees.expressions.functions.agg.Count;
-import org.apache.doris.nereids.trees.expressions.functions.agg.Ndv;
 import org.apache.doris.nereids.trees.expressions.literal.Literal;
 import org.apache.doris.nereids.trees.plans.Plan;
 import org.apache.doris.nereids.trees.plans.PlanType;
@@ -446,7 +444,7 @@ private boolean isUniformGroupByUnique(NamedExpression namedExpression) {
             return false;
         }
         Expression agg = namedExpression.child(0);
-        return agg instanceof Count || agg instanceof Ndv;
+        return ExpressionUtils.isUniformAgg(agg);
     }
 
     @Override

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalHashAggregate.java

@@ -29,8 +29,6 @@
 import org.apache.doris.nereids.trees.expressions.Slot;
 import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
 import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateParam;
-import org.apache.doris.nereids.trees.expressions.functions.agg.Count;
-import org.apache.doris.nereids.trees.expressions.functions.agg.Ndv;
 import org.apache.doris.nereids.trees.expressions.functions.agg.NullableAggregateFunction;
 import org.apache.doris.nereids.trees.plans.AggMode;
 import org.apache.doris.nereids.trees.plans.AggPhase;
@@ -407,7 +405,7 @@ private boolean isUniformGroupByUnique(NamedExpression namedExpression) {
             return false;
         }
         Expression agg = namedExpression.child(0);
-        return agg instanceof Count || agg instanceof Ndv;
+        return ExpressionUtils.isUniformAgg(agg);
     }
 
     @Override

fe/fe-core/src/main/java/org/apache/doris/nereids/util/ExpressionUtils.java

@@ -55,8 +55,10 @@
 import org.apache.doris.nereids.trees.expressions.functions.BoundFunction;
 import org.apache.doris.nereids.trees.expressions.functions.NoneMovableFunction;
 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.Max;
 import org.apache.doris.nereids.trees.expressions.functions.agg.Min;
+import org.apache.doris.nereids.trees.expressions.functions.agg.Ndv;
 import org.apache.doris.nereids.trees.expressions.functions.agg.Sum;
 import org.apache.doris.nereids.trees.expressions.functions.generator.Explode;
 import org.apache.doris.nereids.trees.expressions.functions.generator.ExplodeBitmap;
@@ -1186,6 +1188,24 @@ public static boolean isInjectiveAgg(Expression agg) {
         return agg instanceof Sum || agg instanceof Avg || agg instanceof Max || agg instanceof Min;
     }
 
+    /**
+     * Whether a single-row group always produces the same aggregate result.
+     *
+     * <p>COUNT(*) always consumes its only row. Argument-based COUNT and NDV consume the row only
+     * when every argument is non-null, so nullable arguments may produce either zero or one across
+     * otherwise single-row groups. Keep the proof conservative and inspect the complete argument
+     * expressions rather than only their input slots.</p>
+     */
+    public static boolean isUniformAgg(Expression agg) {
+        if (agg instanceof Count && ((Count) agg).isCountStar()) {
+            return true;
+        }
+        if (!(agg instanceof Count || agg instanceof Ndv)) {
+            return false;
+        }
+        return agg.getArguments().stream().allMatch(Expression::notNullable);
+    }
+
     public static <E> Set<E> mutableCollect(List<? extends Expression> expressions,
             Predicate<TreeNode<Expression>> predicate) {
         Set<E> set = new HashSet<>();

fe/fe-core/src/test/java/org/apache/doris/nereids/properties/UniformTest.java

@@ -45,6 +45,13 @@ protected void runBeforeAll() throws Exception {
                 + "UNIQUE KEY(id)\n"
                 + "distributed by hash(id) buckets 10\n"
                 + "properties('replication_num' = '1');");
+        createTable("create table test.uniform_agg_witness (\n"
+                + "pk int not null,\n"
+                + "b int not null,\n"
+                + "v int null)\n"
+                + "UNIQUE KEY(pk)\n"
+                + "distributed by hash(pk) buckets 10\n"
+                + "properties('replication_num' = '1');");
         connectContext.setDatabase("test");
         connectContext.getSessionVariable().setDisableNereidsRules("PRUNE_EMPTY_PARTITION");
     }
@@ -79,6 +86,33 @@ void testAgg() {
 
     }
 
+    @Test
+    void testSingleRowAggregateUniformityRequiresStableParticipation() {
+        assertAggregateUniform("select count(*) from uniform_agg_witness group by pk", true);
+        assertAggregateUniform("select count(b) from uniform_agg_witness group by pk", true);
+        assertAggregateUniform("select ndv(b) from uniform_agg_witness group by pk", true);
+        assertAggregateUniform("select count(distinct pk, b) from uniform_agg_witness group by pk", true);
+        assertAggregateUniform("select count(if(v is null, 1, 0)) from uniform_agg_witness group by pk", true);
+
+        assertAggregateUniform("select count(v) from uniform_agg_witness group by pk", false);
+        assertAggregateUniform("select ndv(v) from uniform_agg_witness group by pk", false);
+        assertAggregateUniform("select count(distinct b, v) from uniform_agg_witness group by pk", false);
+        assertAggregateUniform(
+                "select count(if(v is null, 1, null)) from uniform_agg_witness group by pk", false);
+        assertAggregateUniform("select count(cast(pk as tinyint)) from uniform_agg_witness group by pk", false);
+        assertAggregateUniform("select count(try_cast(pk as tinyint)) from uniform_agg_witness group by pk", false);
+        assertAggregateUniform("select count(u.name) from uniform_agg_witness w "
+                + "left join uni u on w.pk = u.id group by w.pk", false);
+    }
+
+    private void assertAggregateUniform(String sql, boolean expected) {
+        Plan plan = PlanChecker.from(connectContext)
+                .analyze(sql)
+                .getPlan();
+        Assertions.assertEquals(expected, plan.getLogicalProperties().getTrait()
+                .isUniform(plan.getOutput().get(0)), sql);
+    }
+
     @Test
     void testTopNLimit() {
         Plan plan = PlanChecker.from(connectContext)

regression-test/data/nereids_rules_p0/eliminate_gby_key/eliminate_group_by_key_by_uniform.out

@@ -8,6 +8,21 @@
 
 -- !tranform_to_scalar_agg_not_null_column --
 
+-- !nullable_count_not_uniform --
+7	0	1	0
+7	1	1	1
+
+-- !nullable_ndv_not_uniform --
+7	0	1	0
+7	1	1	1
+
+-- !multi_argument_count_not_uniform --
+7	0	1
+7	1	1
+
+-- !non_nullable_count_uniform --
+7	1	2
+
 -- !project_const --
 \N	1
 \N	1

regression-test/suites/nereids_rules_p0/eliminate_gby_key/eliminate_group_by_key_by_uniform.groovy

@@ -42,6 +42,97 @@ suite("eliminate_group_by_key_by_uniform") {
     qt_empty_tranform_multi_column "select a, min(a), sum(a), count(a) from eli_gbk_by_uniform_t where a = 1 group by a, b,'abc' order by 1,2,3,4"
     qt_tranform_to_scalar_agg_not_null_column "select b, min(a), sum(a), count(a) from eli_gbk_by_uniform_t where b = 1 group by a, b order by 1,2,3,4"
 
+    sql "drop table if exists uniform_agg_witness"
+    sql """
+        create table uniform_agg_witness (
+            pk int not null,
+            b int not null,
+            v int null
+        ) unique key(pk)
+        distributed by hash(pk) buckets 1
+        properties("replication_num"="1")
+    """
+    sql "insert into uniform_agg_witness values (1, 7, null), (2, 7, 9)"
+
+    def nullableCountPlan = sql("""
+        explain select b, c, count(*) as n, sum(h) as sh
+        from (
+            select pk, b, count(v) as c, ndv(v) as h
+            from uniform_agg_witness
+            group by pk, b
+        ) s
+        group by b, c
+    """).collect { it[0] }.join("\n")
+    assertTrue((nullableCountPlan =~ /group by: b\[#\d+\], c\[#\d+\]/).find(),
+            "nullable COUNT must remain in the outer group keys:\n${nullableCountPlan}")
+
+    def nullableNdvPlan = sql("""
+        explain select b, h, count(*) as n, sum(c) as sc
+        from (
+            select pk, b, count(v) as c, ndv(v) as h
+            from uniform_agg_witness
+            group by pk, b
+        ) s
+        group by b, h
+    """).collect { it[0] }.join("\n")
+    assertTrue((nullableNdvPlan =~ /group by: b\[#\d+\], h\[#\d+\]/).find(),
+            "nullable NDV must remain in the outer group keys:\n${nullableNdvPlan}")
+
+    def nonNullableCountPlan = sql("""
+        explain select b, c, count(*) as n
+        from (
+            select pk, b, count(b) as c
+            from uniform_agg_witness
+            group by pk, b
+        ) s
+        group by b, c
+    """).collect { it[0] }.join("\n")
+    assertTrue((nonNullableCountPlan =~ /group by: b\[#\d+\]/).find(),
+            "non-null COUNT should keep the safe group-key elimination:\n${nonNullableCountPlan}")
+    assertFalse((nonNullableCountPlan =~ /group by: b\[#\d+\], c\[#\d+\]/).find(),
+            "non-null COUNT should not remain in the outer group keys:\n${nonNullableCountPlan}")
+
+    order_qt_nullable_count_not_uniform """
+        select b, c, count(*) as n, sum(h) as sh
+        from (
+            select pk, b, count(v) as c, ndv(v) as h
+            from uniform_agg_witness
+            group by pk, b
+        ) s
+        group by b, c
+        order by b, c
+    """
+    order_qt_nullable_ndv_not_uniform """
+        select b, h, count(*) as n, sum(c) as sc
+        from (
+            select pk, b, count(v) as c, ndv(v) as h
+            from uniform_agg_witness
+            group by pk, b
+        ) s
+        group by b, h
+        order by b, h
+    """
+    order_qt_multi_argument_count_not_uniform """
+        select b, c, count(*) as n
+        from (
+            select pk, b, count(distinct b, v) as c
+            from uniform_agg_witness
+            group by pk, b
+        ) s
+        group by b, c
+        order by b, c
+    """
+    order_qt_non_nullable_count_uniform """
+        select b, c, count(*) as n
+        from (
+            select pk, b, count(b) as c
+            from uniform_agg_witness
+            group by pk, b
+        ) s
+        group by b, c
+        order by b, c
+    """
+
     qt_project_const "select sum(c1), c2 from (select a c1,1 c2, d c3 from eli_gbk_by_uniform_t) t group by c2,c3 order by 1,2;"
     qt_project_slot_uniform "select max(c3), c1,c2,c3 from (select a c1,1 c2, d c3 from eli_gbk_by_uniform_t where a=1) t group by c1,c2,c3 order by 1,2,3,4;"
 
@@ -332,4 +423,4 @@ GROUP BY
   );
 
     """
-}
+}