From f7c3f5190c092117488a9be1c2b7a3f708cf9163 Mon Sep 17 00:00:00 2001
From: Harold Wang <harowang@amazon.com>
Date: Mon, 14 Dec 2020 23:15:46 -0800
Subject: [PATCH] Consolidate count function to accept all field type

---
 .../aggregation/AggregatorFunction.java       | 38 ++++++-------------
 1 file changed, 11 insertions(+), 27 deletions(-)

diff --git a/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/expression/aggregation/AggregatorFunction.java b/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/expression/aggregation/AggregatorFunction.java
index 68f1305da7..e467c38585 100644
--- a/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/expression/aggregation/AggregatorFunction.java
+++ b/core/src/main/java/com/amazon/opendistroforelasticsearch/sql/expression/aggregation/AggregatorFunction.java
@@ -28,6 +28,8 @@
 import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.TIME;
 import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.TIMESTAMP;
 
+import com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType;
+import com.amazon.opendistroforelasticsearch.sql.data.type.ExprType;
 import com.amazon.opendistroforelasticsearch.sql.expression.function.BuiltinFunctionName;
 import com.amazon.opendistroforelasticsearch.sql.expression.function.BuiltinFunctionRepository;
 import com.amazon.opendistroforelasticsearch.sql.expression.function.FunctionBuilder;
@@ -35,8 +37,12 @@
 import com.amazon.opendistroforelasticsearch.sql.expression.function.FunctionResolver;
 import com.amazon.opendistroforelasticsearch.sql.expression.function.FunctionSignature;
 import com.google.common.collect.ImmutableMap;
+
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Date;
+import java.util.List;
+import java.util.stream.Collectors;
 
 import lombok.experimental.UtilityClass;
 
@@ -75,33 +81,11 @@ private static FunctionResolver avg() {
 
   private static FunctionResolver count() {
     FunctionName functionName = BuiltinFunctionName.COUNT.getName();
-    return new FunctionResolver(
-        functionName,
-        new ImmutableMap.Builder<FunctionSignature, FunctionBuilder>()
-            .put(new FunctionSignature(functionName, Collections.singletonList(INTEGER)),
-                arguments -> new CountAggregator(arguments, INTEGER))
-            .put(new FunctionSignature(functionName, Collections.singletonList(LONG)),
-                arguments -> new CountAggregator(arguments, INTEGER))
-            .put(new FunctionSignature(functionName, Collections.singletonList(FLOAT)),
-                arguments -> new CountAggregator(arguments, INTEGER))
-            .put(new FunctionSignature(functionName, Collections.singletonList(DOUBLE)),
-                arguments -> new CountAggregator(arguments, INTEGER))
-            .put(new FunctionSignature(functionName, Collections.singletonList(STRING)),
-                arguments -> new CountAggregator(arguments, INTEGER))
-            .put(new FunctionSignature(functionName, Collections.singletonList(STRUCT)),
-                arguments -> new CountAggregator(arguments, INTEGER))
-            .put(new FunctionSignature(functionName, Collections.singletonList(ARRAY)),
-                arguments -> new CountAggregator(arguments, INTEGER))
-            .put(new FunctionSignature(functionName, Collections.singletonList(BOOLEAN)),
-                arguments -> new CountAggregator(arguments, INTEGER))
-            .put(new FunctionSignature(functionName, Collections.singletonList(DATE)),
-                arguments -> new CountAggregator(arguments, INTEGER))
-            .put(new FunctionSignature(functionName, Collections.singletonList(DATETIME)),
-                arguments -> new CountAggregator(arguments, INTEGER))
-            .put(new FunctionSignature(functionName, Collections.singletonList(TIMESTAMP)),
-                arguments -> new CountAggregator(arguments, INTEGER))
-            .build()
-    );
+    FunctionResolver functionResolver = new FunctionResolver(functionName,
+        ExprCoreType.coreTypes().stream().collect(Collectors.toMap(
+          type -> new FunctionSignature(functionName, Collections.singletonList(type)),
+          type -> arguments -> new CountAggregator(arguments, INTEGER))));
+    return functionResolver;
   }
 
   private static FunctionResolver sum() {