Increase the precision of sum return type (#3974)

Signed-off-by: Heng Qian <qianheng@amazon.com>
(cherry picked from commit db2a8bf)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: github-actions[bot]

core/src/main/java/org/opensearch/sql/executor/OpenSearchTypeSystem.java

@@ -5,10 +5,14 @@
 
 package org.opensearch.sql.executor;
 
+import static org.apache.calcite.sql.type.SqlTypeName.APPROX_TYPES;
+import static org.apache.calcite.sql.type.SqlTypeName.INT_TYPES;
+
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.calcite.rel.type.RelDataTypeSystem;
 import org.apache.calcite.rel.type.RelDataTypeSystemImpl;
+import org.apache.calcite.sql.type.BasicSqlType;
 import org.apache.calcite.sql.type.SqlTypeName;
 
 public class OpenSearchTypeSystem extends RelDataTypeSystemImpl {
@@ -20,14 +24,30 @@ private OpenSearchTypeSystem() {}
   public RelDataType deriveAvgAggType(RelDataTypeFactory typeFactory, RelDataType argumentType) {
     if (SqlTypeName.DECIMAL == argumentType.getSqlTypeName()) {
       return typeFactory.createTypeWithNullability(highPrecision(typeFactory, argumentType), false);
-    } else if (SqlTypeName.INT_TYPES.contains(argumentType.getSqlTypeName())) {
+    } else if (INT_TYPES.contains(argumentType.getSqlTypeName())) {
       return typeFactory.createTypeWithNullability(
           typeFactory.createSqlType(SqlTypeName.DOUBLE), false);
     } else {
       return argumentType;
     }
   }
 
+  @Override
+  public RelDataType deriveSumType(RelDataTypeFactory typeFactory, RelDataType argumentType) {
+    argumentType = super.deriveSumType(typeFactory, argumentType);
+    if (argumentType instanceof BasicSqlType) {
+      SqlTypeName typeName = argumentType.getSqlTypeName();
+      if (INT_TYPES.contains(typeName)) {
+        return typeFactory.createTypeWithNullability(
+            typeFactory.createSqlType(SqlTypeName.BIGINT), argumentType.isNullable());
+      } else if (APPROX_TYPES.contains(typeName)) {
+        return typeFactory.createTypeWithNullability(
+            typeFactory.createSqlType(SqlTypeName.DOUBLE), argumentType.isNullable());
+      }
+    }
+    return argumentType;
+  }
+
   /**
    * Compute a higher precision version of a type.
    *

integ-test/src/test/java/org/opensearch/sql/calcite/tpch/CalcitePPLTpchIT.java

@@ -249,8 +249,8 @@ public void testQ12() throws IOException {
     verifySchemaInOrder(
         actual,
         schema("l_shipmode", "string"),
-        schema("high_line_count", "int"),
-        schema("low_line_count", "int"));
+        schema("high_line_count", "bigint"),
+        schema("low_line_count", "bigint"));
     verifyDataRows(actual, rows("MAIL", 5, 5), rows("SHIP", 5, 10));
   }
 

integ-test/src/yamlRestTest/resources/rest-api-spec/test/issues/3958.yml

@@ -0,0 +1,57 @@
+setup:
+  - do:
+      query.settings:
+        body:
+          transient:
+            plugins.calcite.enabled : true
+            plugins.calcite.fallback.allowed : false
+
+  - do:
+      indices.create:
+        index: tmp
+        body:
+          settings:
+            number_of_shards: 1
+            number_of_replicas: 0
+          mappings:
+            properties:
+              demo:
+                type: short
+
+
+---
+teardown:
+  - do:
+      query.settings:
+        body:
+          transient:
+            plugins.calcite.enabled : false
+            plugins.calcite.fallback.allowed : true
+
+---
+"Handle sum results exceeds the max value of short":
+  - skip:
+      features:
+        - headers
+        - allowed_warnings
+  - do:
+      bulk:
+        index: tmp
+        refresh: true
+        body:
+          - '{"index": {}}'
+          - '{"demo": 20000}'
+          - '{"index": {}}'
+          - '{"demo": 20000}'
+
+  - do:
+      allowed_warnings:
+        - 'Loading the fielddata on the _id field is deprecated and will be removed in future versions. If you require sorting or aggregating on this field you should also include the id in the body of your documents, and map this field as a keyword field that has [doc_values] enabled'
+      headers:
+        Content-Type: 'application/json'
+      ppl:
+        body:
+          query: 'source=tmp | stats sum(demo)'
+  - match: {"total": 1}
+  - match: {"schema": [{"name": "sum(demo)", "type": "bigint"}]}
+  - match: {"datarows": [[40000]]}