Consolidate count function to accept all field type

opendistro-for-elasticsearch · Dec 15, 2020 · f7c3f51 · f7c3f51
1 parent c826138
commit f7c3f51
Showing 1 changed file with 11 additions and 27 deletions.
diff --git a/.../com/amazon/opendistroforelasticsearch/sql/expression/aggregation/AggregatorFunction.java b/.../com/amazon/opendistroforelasticsearch/sql/expression/aggregation/AggregatorFunction.java
@@ -28,15 +28,21 @@
 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;
 import com.amazon.opendistroforelasticsearch.sql.expression.function.FunctionName;
 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() {