feat: support COUNT()

datafusion/functions-aggregate/src/count.rs

@@ -44,7 +44,7 @@ use datafusion_expr::{
     function::AccumulatorArgs, utils::format_state_name, Accumulator, AggregateUDFImpl,
     EmitTo, GroupsAccumulator, Signature, Volatility,
 };
-use datafusion_expr::{Expr, ReversedUDAF};
+use datafusion_expr::{Expr, ReversedUDAF, TypeSignature};
 use datafusion_physical_expr_common::aggregate::groups_accumulator::accumulate::accumulate_indices;
 use datafusion_physical_expr_common::{
     aggregate::count_distinct::{
@@ -95,7 +95,10 @@ impl Default for Count {
 impl Count {
     pub fn new() -> Self {
         Self {
-            signature: Signature::variadic_any(Volatility::Immutable),
+            signature: Signature::one_of(
+                vec![TypeSignature::VariadicAny, TypeSignature::Any(0)],
+                Volatility::Immutable,
+            ),
         }
     }
 }

datafusion/optimizer/src/analyzer/count_empty_rule.rs

@@ -0,0 +1,91 @@
+// 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.
+
+use datafusion_common::tree_node::{Transformed, TreeNode};
+use datafusion_common::Result;
+use datafusion_common::{config::ConfigOptions, tree_node::TransformedResult};
+use datafusion_expr::utils::COUNT_STAR_EXPANSION;
+use datafusion_expr::{
+    expr::{AggregateFunction, AggregateFunctionDefinition, WindowFunction},
+    LogicalPlan, WindowFunctionDefinition,
+};
+use datafusion_expr::{lit, Expr};
+
+use crate::utils::NamePreserver;
+use crate::AnalyzerRule;
+
+/// Rewrite `Count()` to `Count(Expr:Literal(1))`.
+#[derive(Default)]
+pub struct CountEmptyRule {}
+
+impl CountEmptyRule {
+    /// Create a new instance of the rule
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl AnalyzerRule for CountEmptyRule {
+    fn analyze(&self, plan: LogicalPlan, _: &ConfigOptions) -> Result<LogicalPlan> {
+        plan.transform_down_with_subqueries(analyze_internal).data()
+    }
+
+    fn name(&self) -> &str {
+        "count_empty_rule"
+    }
+}
+
+fn is_count_empty_aggregate(aggregate_function: &AggregateFunction) -> bool {
+    matches!(aggregate_function,
+        AggregateFunction {
+            func_def: AggregateFunctionDefinition::UDF(udf),
+            args,
+            ..
+        } if udf.name() == "count" && args.is_empty())
+}
+
+fn is_count_empty_window_aggregate(window_function: &WindowFunction) -> bool {
+    let args = &window_function.args;
+    matches!(window_function.fun,
+        WindowFunctionDefinition::AggregateUDF(ref udaf)
+            if udaf.name() == "count" && args.is_empty())
+}
+
+fn analyze_internal(plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
+    let name_preserver = NamePreserver::new(&plan);
+    plan.map_expressions(|expr| {
+        let original_name = name_preserver.save(&expr)?;
+        let transformed_expr = expr.transform_up(|expr| match expr {
+            Expr::WindowFunction(mut window_function)
+                if is_count_empty_window_aggregate(&window_function) =>
+            {
+                window_function.args = vec![lit(COUNT_STAR_EXPANSION)];
+                Ok(Transformed::yes(Expr::WindowFunction(window_function)))
+            }
+            Expr::AggregateFunction(mut aggregate_function)
+                if is_count_empty_aggregate(&aggregate_function) =>
+            {
+                aggregate_function.args = vec![lit(COUNT_STAR_EXPANSION)];
+                Ok(Transformed::yes(Expr::AggregateFunction(
+                    aggregate_function,
+                )))
+            }
+            _ => Ok(Transformed::no(expr)),
+        })?;
+        transformed_expr.map_data(|data| original_name.restore(data))
+    })
+}

datafusion/optimizer/src/analyzer/count_wildcard_rule.rs

@@ -35,7 +35,7 @@ pub struct CountWildcardRule {}
 
 impl CountWildcardRule {
     pub fn new() -> Self {
-        CountWildcardRule {}
+        Self {}
     }
 }
 

datafusion/optimizer/src/analyzer/mod.rs

@@ -29,6 +29,7 @@ use datafusion_expr::expr::InSubquery;
 use datafusion_expr::expr_rewriter::FunctionRewrite;
 use datafusion_expr::{Expr, LogicalPlan};
 
+use crate::analyzer::count_empty_rule::CountEmptyRule;
 use crate::analyzer::count_wildcard_rule::CountWildcardRule;
 use crate::analyzer::inline_table_scan::InlineTableScan;
 use crate::analyzer::subquery::check_subquery_expr;
@@ -37,6 +38,7 @@ use crate::utils::log_plan;
 
 use self::function_rewrite::ApplyFunctionRewrites;
 
+pub mod count_empty_rule;
 pub mod count_wildcard_rule;
 pub mod function_rewrite;
 pub mod inline_table_scan;
@@ -91,6 +93,7 @@ impl Analyzer {
             Arc::new(InlineTableScan::new()),
             Arc::new(TypeCoercion::new()),
             Arc::new(CountWildcardRule::new()),
+            Arc::new(CountEmptyRule::new()),
         ];
         Self::with_rules(rules)
     }

datafusion/sqllogictest/test_files/count_empty_rule.slt

@@ -0,0 +1,105 @@
+# 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.
+
+statement ok
+CREATE TABLE t1 (a INTEGER, b INTEGER, c INTEGER);
+
+statement ok
+INSERT INTO t1 VALUES
+(1, 2, 3),
+(1, 5, 6),
+(2, 3, 5);
+
+statement ok
+CREATE TABLE t2 (a INTEGER, b INTEGER, c INTEGER);
+
+query TT
+EXPLAIN SELECT COUNT() FROM (SELECT 1 AS a, 2 AS b) AS t;
+----
+logical_plan
+01)Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count()]]
+02)--SubqueryAlias: t
+03)----EmptyRelation
+physical_plan
+01)ProjectionExec: expr=[1 as count()]
+02)--PlaceholderRowExec
+
+query TT
+EXPLAIN SELECT t1.a, COUNT() FROM t1 GROUP BY t1.a;
+----
+logical_plan
+01)Aggregate: groupBy=[[t1.a]], aggr=[[count(Int64(1)) AS count()]]
+02)--TableScan: t1 projection=[a]
+physical_plan
+01)AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count()]
+02)--CoalesceBatchesExec: target_batch_size=8192
+03)----RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4
+04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+05)--------AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count()]
+06)----------MemoryExec: partitions=1, partition_sizes=[1]
+
+query TT
+EXPLAIN SELECT t1.a, COUNT() AS cnt FROM t1 GROUP BY t1.a HAVING COUNT() > 0;
+----
+logical_plan
+01)Projection: t1.a, count() AS cnt
+02)--Filter: count() > Int64(0)
+03)----Aggregate: groupBy=[[t1.a]], aggr=[[count(Int64(1)) AS count()]]
+04)------TableScan: t1 projection=[a]
+physical_plan
+01)ProjectionExec: expr=[a@0 as a, count()@1 as cnt]
+02)--CoalesceBatchesExec: target_batch_size=8192
+03)----FilterExec: count()@1 > 0
+04)------AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count()]
+05)--------CoalesceBatchesExec: target_batch_size=8192
+06)----------RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4
+07)------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+08)--------------AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count()]
+09)----------------MemoryExec: partitions=1, partition_sizes=[1]
+
+query II
+SELECT t1.a, COUNT() AS cnt FROM t1 GROUP BY t1.a HAVING COUNT() > 1;
+----
+1 2
+
+query TT
+EXPLAIN SELECT a, COUNT() OVER (PARTITION BY a) AS count_a FROM t1;
+----
+logical_plan
+01)Projection: t1.a, count() PARTITION BY [t1.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS count_a
+02)--WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [t1.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS count() PARTITION BY [t1.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]]
+03)----TableScan: t1 projection=[a]
+physical_plan
+01)ProjectionExec: expr=[a@0 as a, count() PARTITION BY [t1.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@1 as count_a]
+02)--WindowAggExec: wdw=[count() PARTITION BY [t1.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "count() PARTITION BY [t1.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]
+03)----SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true]
+04)------CoalesceBatchesExec: target_batch_size=8192
+05)--------RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=1
+06)----------MemoryExec: partitions=1, partition_sizes=[1]
+
+query II
+SELECT a, COUNT() OVER (PARTITION BY a) AS count_a FROM t1 ORDER BY a;
+----
+1 2
+1 2
+2 1
+
+statement ok
+DROP TABLE t1;
+
+statement ok
+DROP TABLE t2;

datafusion/sqllogictest/test_files/errors.slt

@@ -103,10 +103,6 @@ SELECT power(1, 2, 3);
 # Wrong window/aggregate function signature
 #
 
-# AggregateFunction with wrong number of arguments
-query error
-select count();
-
 # AggregateFunction with wrong number of arguments
 query error
 select avg(c1, c12) from aggregate_test_100;

datafusion/sqllogictest/test_files/explain.slt

@@ -183,6 +183,7 @@ logical_plan after apply_function_rewrites SAME TEXT AS ABOVE
 logical_plan after inline_table_scan SAME TEXT AS ABOVE
 logical_plan after type_coercion SAME TEXT AS ABOVE
 logical_plan after count_wildcard_rule SAME TEXT AS ABOVE
+logical_plan after count_empty_rule SAME TEXT AS ABOVE
 analyzed_logical_plan SAME TEXT AS ABOVE
 logical_plan after eliminate_nested_union SAME TEXT AS ABOVE
 logical_plan after simplify_expressions SAME TEXT AS ABOVE