From 5870bbfd968444c9c52b02e76c14e3c7b79326a1 Mon Sep 17 00:00:00 2001
From: Zhongnan Su <szhongna@amazon.com>
Date: Wed, 22 Apr 2020 16:23:55 -0700
Subject: [PATCH] [BugFix] Enforce AVG to return double data type (#437)

* Enforce AVG to return double. Update UT, add IT.
---
 .../types/function/AggregateFunction.java     |  3 +-
 ...SemanticAnalyzerAggregateFunctionTest.java |  4 +--
 .../esintgtest/AggregationExpressionIT.java   | 28 +++++++++++++++++++
 .../sql/esintgtest/TestUtils.java             |  5 ++--
 4 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/sql/antlr/semantic/types/function/AggregateFunction.java b/src/main/java/com/amazon/opendistroforelasticsearch/sql/antlr/semantic/types/function/AggregateFunction.java
index 6d99c046a0..8b1f642f0b 100644
--- a/src/main/java/com/amazon/opendistroforelasticsearch/sql/antlr/semantic/types/function/AggregateFunction.java
+++ b/src/main/java/com/amazon/opendistroforelasticsearch/sql/antlr/semantic/types/function/AggregateFunction.java
@@ -18,6 +18,7 @@
 import com.amazon.opendistroforelasticsearch.sql.antlr.semantic.types.Type;
 import com.amazon.opendistroforelasticsearch.sql.antlr.semantic.types.TypeExpression;
 
+import static com.amazon.opendistroforelasticsearch.sql.antlr.semantic.types.base.ESDataType.DOUBLE;
 import static com.amazon.opendistroforelasticsearch.sql.antlr.semantic.types.base.ESDataType.ES_TYPE;
 import static com.amazon.opendistroforelasticsearch.sql.antlr.semantic.types.base.ESDataType.INTEGER;
 import static com.amazon.opendistroforelasticsearch.sql.antlr.semantic.types.base.ESDataType.NUMBER;
@@ -33,7 +34,7 @@ public enum AggregateFunction implements TypeExpression {
     ),
     MAX(func(T(NUMBER)).to(T)),
     MIN(func(T(NUMBER)).to(T)),
-    AVG(func(T(NUMBER)).to(T)),
+    AVG(func(T(NUMBER)).to(DOUBLE)),
     SUM(func(T(NUMBER)).to(T));
 
     private TypeExpressionSpec[] specifications;
diff --git a/src/test/java/com/amazon/opendistroforelasticsearch/sql/antlr/semantic/SemanticAnalyzerAggregateFunctionTest.java b/src/test/java/com/amazon/opendistroforelasticsearch/sql/antlr/semantic/SemanticAnalyzerAggregateFunctionTest.java
index be89e428ec..864324bbf4 100644
--- a/src/test/java/com/amazon/opendistroforelasticsearch/sql/antlr/semantic/SemanticAnalyzerAggregateFunctionTest.java
+++ b/src/test/java/com/amazon/opendistroforelasticsearch/sql/antlr/semantic/SemanticAnalyzerAggregateFunctionTest.java
@@ -91,7 +91,7 @@ public void avgFunctionCallOnBooleanFieldShouldFail() {
         expectValidationFailWithErrorMessages(
             "SELECT AVG(p.active) FROM semantics s, s.projects p",
             "Function [AVG] cannot work with [BOOLEAN].",
-            "Usage: AVG(NUMBER T) -> T"
+            "Usage: AVG(NUMBER T) -> DOUBLE"
         );
     }
 
@@ -152,7 +152,7 @@ public void useAvgFunctionCallWithTextFieldInHavingClauseShouldFail() {
         expectValidationFailWithErrorMessages(
             "SELECT city FROM semantics GROUP BY city HAVING AVG(address) > 10",
             "Function [AVG] cannot work with [TEXT].",
-            "Usage: AVG(NUMBER T) -> T"
+            "Usage: AVG(NUMBER T) -> DOUBLE"
         );
     }
 
diff --git a/src/test/java/com/amazon/opendistroforelasticsearch/sql/esintgtest/AggregationExpressionIT.java b/src/test/java/com/amazon/opendistroforelasticsearch/sql/esintgtest/AggregationExpressionIT.java
index ead9e29a65..7d345cbf56 100644
--- a/src/test/java/com/amazon/opendistroforelasticsearch/sql/esintgtest/AggregationExpressionIT.java
+++ b/src/test/java/com/amazon/opendistroforelasticsearch/sql/esintgtest/AggregationExpressionIT.java
@@ -29,6 +29,7 @@ public class AggregationExpressionIT extends SQLIntegTestCase {
     @Override
     protected void init() throws Exception {
         loadIndex(Index.ACCOUNT);
+        loadIndex(Index.BANK);
     }
 
     @Test
@@ -66,6 +67,33 @@ public void noGroupKeyMaxAddLiteralShouldPass() {
         verifyDataRows(response, rows(41));
     }
 
+    @Test
+    public void noGroupKeyAvgOnIntegerShouldPass() {
+        JSONObject response = executeJdbcRequest(String.format(
+                "SELECT AVG(age) as avg " +
+                        "FROM %s",
+                Index.BANK.getName()));
+
+        verifySchema(response, schema("avg", "avg", "double"));
+        verifyDataRows(response, rows(34));
+    }
+
+    @Test
+    public void hasGroupKeyAvgOnIntegerShouldPass() {
+        JSONObject response = executeJdbcRequest(String.format(
+                "SELECT gender, AVG(age) as avg " +
+                "FROM %s " +
+                "GROUP BY gender",
+                Index.BANK.getName()));
+
+        verifySchema(response,
+                schema("gender", null, "text"),
+                schema("avg", "avg", "double"));
+        verifyDataRows(response,
+                rows("m", 34.25),
+                rows("f", 33.666666666666664d));
+    }
+
     @Test
     public void hasGroupKeyMaxAddMinShouldPass() {
         JSONObject response = executeJdbcRequest(String.format(
diff --git a/src/test/java/com/amazon/opendistroforelasticsearch/sql/esintgtest/TestUtils.java b/src/test/java/com/amazon/opendistroforelasticsearch/sql/esintgtest/TestUtils.java
index 000576020e..83d9cb5342 100644
--- a/src/test/java/com/amazon/opendistroforelasticsearch/sql/esintgtest/TestUtils.java
+++ b/src/test/java/com/amazon/opendistroforelasticsearch/sql/esintgtest/TestUtils.java
@@ -499,8 +499,9 @@ public static String getBankIndexMapping() {
                 "        \"type\": \"text\"\n" +
                 "      },\n" +
                 "      \"gender\": {\n" +
-                "        \"type\": \"text\"\n" +
-                "      },\n" +
+                "        \"type\": \"text\",\n" +
+                "        \"fielddata\": true\n" +
+                "      }," +
                 "      \"lastname\": {\n" +
                 "        \"type\": \"keyword\"\n" +
                 "      },\n" +