[opt](Nereids) remove restrict for count(*) in window (#35220)

support count(*) used for window function CREATE TABLE `t1` ( `id` INT NULL, `dt` TEXT NULL ) DISTRIBUTED BY HASH(`id`) BUCKETS 10 PROPERTIES ( "replication_allocation" = "tag.location.default: 1" ); select *, count(*) over() from t1;
apache · May 23, 2024 · d8c39a9 · d8c39a9
1 parent 67a3ff7
commit d8c39a9
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 90 deletions.
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java
@@ -2068,7 +2068,7 @@ public Expression visitFunctionCallExpression(DorisParser.FunctionCallExpression
 
             List<UnboundStar> unboundStars = ExpressionUtils.collectAll(params, UnboundStar.class::isInstance);
             if (!unboundStars.isEmpty()) {
-                if (functionName.equalsIgnoreCase("count")) {
+                if (ctx.functionIdentifier().dbName == null && functionName.equalsIgnoreCase("count")) {
                     if (unboundStars.size() > 1) {
                         throw new ParseException(
                                 "'*' can only be used once in conjunction with COUNT: " + functionName, ctx);
@@ -2077,9 +2077,10 @@ public Expression visitFunctionCallExpression(DorisParser.FunctionCallExpression
                         throw new ParseException("'*' can not has qualifier: " + unboundStars.size(), ctx);
                     }
                     if (ctx.windowSpec() != null) {
-                        // todo: support count(*) as window function
-                        throw new ParseException(
-                                "COUNT(*) isn't supported as window function; can use COUNT(col)", ctx);
+                        if (isDistinct) {
+                            throw new ParseException("DISTINCT not allowed in analytic function: " + functionName, ctx);
+                        }
+                        return withWindowSpec(ctx.windowSpec(), new Count());
                     }
                     return new Count();
                 }

diff --git a/regression-test/data/nereids_p0/sql_functions/window_functions/test_count_star.out b/regression-test/data/nereids_p0/sql_functions/window_functions/test_count_star.out
@@ -0,0 +1,6 @@
+-- This file is automatically generated. You should know what you did if you want to edit this
+-- !count_star --
+3
+3
+3
+
diff --git a/regression-test/suites/nereids_p0/sql_functions/window_functions/test_count_star.groovy b/regression-test/suites/nereids_p0/sql_functions/window_functions/test_count_star.groovy
@@ -0,0 +1,22 @@
+// 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_count_star") {
+    order_qt_count_star """
+        select count(*) over() from nereids_test_query_db.test;
+    """
+}
diff --git a/regression-test/suites/nereids_p0/sql_functions/window_functions/test_window_function.groovy b/regression-test/suites/nereids_p0/sql_functions/window_functions/test_window_function.groovy
@@ -369,39 +369,27 @@ suite("test_window_function") {
     // test error
     test {
         sql("select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, lag(${k2}) over (partition by ${k1} order by ${k3}) from baseall")
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, lag(${k2}, -1, 1) over (partition by ${k1} order by ${k3}) from baseall"
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, lag(${k2}, 1) over (partition by ${k1} order by ${k3}) from baseall"
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, lead(${k2}) over (partition by ${k1} order by ${k3}) from baseall"
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, lead(${k2}, -1, 1) over (partition by ${k1} order by ${k3}) from baseall"
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, lead(${k2}, 1) over (partition by ${k1} order by ${k3}) from baseall"
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     qt_window_error1"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, first_value(${k2}) over (partition by ${k1}) from baseall order by ${k1}"""
     qt_window_error2"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, first_value(${k2}) over (order by ${k3}) from baseall"""
@@ -410,53 +398,34 @@ suite("test_window_function") {
         sql"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, sum(${k2}) over (partition by ${k1} order by ${k3} rows
             between current row and unbounded preceding) as wj
             from baseall order by ${k1}, wj"""
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, sum(${k2}) over (partition by ${k1} order by ${k3} rows 
             between 0 preceding and 1 following) as wj 
             from baseall order by ${k1}, wj"""
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, sum(${k2}) over (partition by ${k1} order by ${k3} rows 
             between unbounded following and current row) as wj 
             from baseall order by ${k1}, wj"""
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, rank(${k2}) over (partition by ${k1} order by ${k3}) as wj
             from baseall order by ${k1}, wj"""
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, max() over (partition by ${k1} order by ${k3}) as wj
             from baseall order by ${k1}, wj"""
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
-    }
-    test {
-        sql"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, count(*) over (partition by ${k1} order by ${k3}) as wj
-            from baseall order by ${k1}, wj"""
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, count(${k2}) over (order by ${k1} rows partition by ${k3}) as wj
             from baseall order by ${k1}, wj"""
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
 
     // test_query_rank

diff --git a/regression-test/suites/query_p0/sql_functions/window_functions/test_window_function.groovy b/regression-test/suites/query_p0/sql_functions/window_functions/test_window_function.groovy
@@ -367,39 +367,27 @@ suite("test_window_function") {
     // test error
     test {
         sql("select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, lag(${k2}) over (partition by ${k1} order by ${k3}) from baseall")
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, lag(${k2}, -1, 1) over (partition by ${k1} order by ${k3}) from baseall"
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, lag(${k2}, 1) over (partition by ${k1} order by ${k3}) from baseall"
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, lead(${k2}) over (partition by ${k1} order by ${k3}) from baseall"
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, lead(${k2}, -1, 1) over (partition by ${k1} order by ${k3}) from baseall"
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, lead(${k2}, 1) over (partition by ${k1} order by ${k3}) from baseall"
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     qt_window_error1"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, first_value(${k2}) over (partition by ${k1}) from baseall order by ${k1}"""
     qt_window_error2"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, first_value(${k2}) over (order by ${k3}) from baseall"""
@@ -408,53 +396,34 @@ suite("test_window_function") {
         sql"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, sum(${k2}) over (partition by ${k1} order by ${k3} rows
             between current row and unbounded preceding) as wj
             from baseall order by ${k1}, wj"""
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, sum(${k2}) over (partition by ${k1} order by ${k3} rows 
             between 0 preceding and 1 following) as wj 
             from baseall order by ${k1}, wj"""
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, sum(${k2}) over (partition by ${k1} order by ${k3} rows 
             between unbounded following and current row) as wj 
             from baseall order by ${k1}, wj"""
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, rank(${k2}) over (partition by ${k1} order by ${k3}) as wj
             from baseall order by ${k1}, wj"""
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, max() over (partition by ${k1} order by ${k3}) as wj
             from baseall order by ${k1}, wj"""
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
-    }
-    test {
-        sql"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, count(*) over (partition by ${k1} order by ${k3}) as wj
-            from baseall order by ${k1}, wj"""
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
     test {
         sql"""select /*+SET_VAR(parallel_fragment_exec_instance_num=1) */ ${k1}, count(${k2}) over (order by ${k1} rows partition by ${k3}) as wj
             from baseall order by ${k1}, wj"""
-        check { result, exception, startTime, endTime ->
-            assertTrue(exception != null)
-        }
+        exception ""
     }
 
     // test_query_rank