Enable count(Date)

Add IT
opendistro-for-elasticsearch · Dec 11, 2020 · be4e456 · be4e456
1 parent 5635846
commit be4e456
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 0 deletions.
diff --git a/.../com/amazon/opendistroforelasticsearch/sql/expression/aggregation/AggregatorFunction.java b/.../com/amazon/opendistroforelasticsearch/sql/expression/aggregation/AggregatorFunction.java
@@ -36,6 +36,8 @@
 import com.amazon.opendistroforelasticsearch.sql.expression.function.FunctionSignature;
 import com.google.common.collect.ImmutableMap;
 import java.util.Collections;
+import java.util.Date;
+
 import lombok.experimental.UtilityClass;
 
 /**
@@ -92,6 +94,12 @@ private static FunctionResolver count() {
                 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()
     );
   }

diff --git a/...com/amazon/opendistroforelasticsearch/sql/expression/aggregation/CountAggregatorTest.java b/...com/amazon/opendistroforelasticsearch/sql/expression/aggregation/CountAggregatorTest.java
@@ -17,12 +17,16 @@
 
 import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.ARRAY;
 import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.BOOLEAN;
+import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.DATE;
+import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.DATETIME;
 import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.DOUBLE;
 import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.FLOAT;
 import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.INTEGER;
 import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.LONG;
 import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.STRING;
 import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.STRUCT;
+import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.TIMESTAMP;
+
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
@@ -58,6 +62,24 @@ public void count_double_field_expression() {
     assertEquals(4, result.value());
   }
 
+  @Test
+  public void count_date_field_expression() {
+    ExprValue result = aggregation(dsl.count(DSL.ref("date_value", DATE)), tuples);
+    assertEquals(4, result.value());
+  }
+
+  @Test
+  public void count_timestamp_field_expression() {
+    ExprValue result = aggregation(dsl.count(DSL.ref("timestamp_value", TIMESTAMP)), tuples);
+    assertEquals(4, result.value());
+  }
+
+  @Test
+  public void count_datetime_field_expression() {
+    ExprValue result = aggregation(dsl.count(DSL.ref("datetime_value", DATETIME)), tuples);
+    assertEquals(4, result.value());
+  }
+
   @Test
   public void count_arithmetic_expression() {
     ExprValue result = aggregation(dsl.count(