Detect integral SUM overflow and fix BIGINT AVG

Integral SUM used Calcite's unchecked long accumulator and could
silently wrap. BIGINT AVG could likewise overflow its intermediate
long sum before division.

Use CHECKED_LONG_SUM for integral inputs and Math.addExact during
enumerable accumulation. Preserve SqlKind.SUM so planner rewrites and
native OpenSearch sum pushdown remain unchanged. Range-check and narrow
pushed results while retaining the backend's double semantics.

Use a double sum and long count for enumerable BIGINT AVG. Preserve
SqlKind.AVG and native avg pushdown. Add unit, integration, REST,
documentation, and explain coverage for both execution paths.

Signed-off-by: Kai Huang <ahkcs@amazon.com>

Author: ahkcs

api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlV2Test.java

@@ -286,7 +286,7 @@ SELECT department, SUM(age) AS total FROM catalog.employees GROUP BY department
         .assertPlan(
             """
             LogicalProject(department=[$0], total=[$1])
-              LogicalAggregate(group=[{0}], SUM(age)=[SUM($1)])
+              LogicalAggregate(group=[{0}], SUM(age)=[CHECKED_LONG_SUM($1)])
                 LogicalProject(department=[$3], age=[$2])
                   LogicalTableScan(table=[[catalog, employees]])
             """);
@@ -366,7 +366,7 @@ SELECT department, SUM(age) FILTER(WHERE age > 30) FROM catalog.employees
             """)
         .assertPlan(
             """
-            LogicalAggregate(group=[{0}], SUM(age) FILTER(WHERE age > 30)=[SUM($1) FILTER $2])
+            LogicalAggregate(group=[{0}], SUM(age) FILTER(WHERE age > 30)=[CHECKED_LONG_SUM($1) FILTER $2])
               LogicalProject(department=[$3], age=[$2], $f3=[>($2, 30)])
                 LogicalTableScan(table=[[catalog, employees]])
             """);
@@ -487,7 +487,7 @@ SELECT name, SUM(age) OVER(PARTITION BY department ORDER BY age) FROM catalog.em
             """)
         .assertPlan(
             """
-            LogicalProject(name=[$1], SUM(age) OVER(PARTITION BY department ORDER BY age)=[SUM($2) OVER (PARTITION BY $3 ORDER BY $2 NULLS FIRST)])
+            LogicalProject(name=[$1], SUM(age) OVER(PARTITION BY department ORDER BY age)=[CHECKED_LONG_SUM($2) OVER (PARTITION BY $3 ORDER BY $2 NULLS FIRST)])
               LogicalTableScan(table=[[catalog, employees]])
             """);
   }

core/src/main/java/org/opensearch/sql/calcite/udf/udaf/BigintAvgAggFunction.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite.udf.udaf;
+
+/** BIGINT average aggregate that accumulates in double to avoid an intermediate long overflow. */
+public class BigintAvgAggFunction {
+
+  public static Accumulator init() {
+    return new Accumulator();
+  }
+
+  public static Accumulator add(Accumulator accumulator, Long value) {
+    if (value != null) {
+      accumulator.sum += value;
+      accumulator.count++;
+    }
+    return accumulator;
+  }
+
+  public static Double result(Accumulator accumulator) {
+    return accumulator.count == 0 ? null : accumulator.sum / accumulator.count;
+  }
+
+  public static class Accumulator {
+    private double sum;
+    private long count;
+  }
+}

core/src/main/java/org/opensearch/sql/calcite/udf/udaf/CheckedLongSumAggFunction.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite.udf.udaf;
+
+/** BIGINT sum aggregate that throws when its running long accumulator overflows. */
+public class CheckedLongSumAggFunction {
+
+  public static long init() {
+    return 0L;
+  }
+
+  public static long add(long accumulator, long value) {
+    return Math.addExact(accumulator, value);
+  }
+
+  public static long result(long accumulator) {
+    return accumulator;
+  }
+}

core/src/main/java/org/opensearch/sql/calcite/utils/UserDefinedFunctionUtils.java

@@ -94,9 +94,32 @@ public static SqlUserDefinedAggFunction createUserDefinedAggFunction(
       String functionName,
       SqlReturnTypeInference returnType,
       @Nullable UDFOperandMetadata operandMetadata) {
+    return createReflectiveAggFunction(udafClass, functionName, returnType, operandMetadata);
+  }
+
+  /** Creates an aggregate function from a class following Calcite's reflective UDAF convention. */
+  public static SqlUserDefinedAggFunction createReflectiveAggFunction(
+      Class<?> udafClass,
+      String functionName,
+      SqlReturnTypeInference returnType,
+      @Nullable UDFOperandMetadata operandMetadata) {
+    return createReflectiveAggFunction(
+        udafClass, functionName, SqlKind.OTHER_FUNCTION, returnType, operandMetadata);
+  }
+
+  /**
+   * Creates an aggregate function whose kind remains visible to planner rules while execution uses
+   * the supplied reflective UDAF.
+   */
+  public static SqlUserDefinedAggFunction createReflectiveAggFunction(
+      Class<?> udafClass,
+      String functionName,
+      SqlKind kind,
+      SqlReturnTypeInference returnType,
+      @Nullable UDFOperandMetadata operandMetadata) {
     return new SqlUserDefinedAggFunction(
         new SqlIdentifier(functionName, SqlParserPos.ZERO),
-        SqlKind.OTHER_FUNCTION,
+        kind,
         returnType,
         null,
         operandMetadata,

core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java

@@ -8,6 +8,7 @@
 import static org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils.adaptExprMethodToUDF;
 import static org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils.adaptExprMethodWithPropertiesToUDF;
 import static org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils.adaptMathFunctionToUDF;
+import static org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils.createReflectiveAggFunction;
 import static org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils.createUserDefinedAggFunction;
 
 import com.google.common.base.Suppliers;
@@ -29,6 +30,8 @@
 import org.apache.calcite.sql.type.SqlTypeTransforms;
 import org.apache.calcite.sql.util.ReflectiveSqlOperatorTable;
 import org.apache.calcite.util.BuiltInMethod;
+import org.opensearch.sql.calcite.udf.udaf.BigintAvgAggFunction;
+import org.opensearch.sql.calcite.udf.udaf.CheckedLongSumAggFunction;
 import org.opensearch.sql.calcite.udf.udaf.DistinctCountApproxLogicalAggFunction;
 import org.opensearch.sql.calcite.udf.udaf.FirstAggFunction;
 import org.opensearch.sql.calcite.udf.udaf.LastAggFunction;
@@ -449,7 +452,6 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
       new NumberToStringFunction().toUDF("NUMBER_TO_STRING");
   public static final SqlOperator TONUMBER = new ToNumberFunction().toUDF("TONUMBER");
   public static final SqlOperator TOSTRING = new ToStringFunction().toUDF("TOSTRING");
-
   // PPL Convert command functions
   public static final SqlOperator AUTO = new AutoConvertFunction().toUDF("AUTO");
   public static final SqlOperator NUM = new NumConvertFunction().toUDF("NUM");
@@ -488,6 +490,20 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
       new NullableSqlAvgAggFunction(SqlKind.VAR_POP);
   public static final SqlAggFunction VAR_SAMP_NULLABLE =
       new NullableSqlAvgAggFunction(SqlKind.VAR_SAMP);
+  public static final SqlAggFunction CHECKED_LONG_SUM =
+      createReflectiveAggFunction(
+          CheckedLongSumAggFunction.class,
+          "CHECKED_LONG_SUM",
+          SqlKind.SUM,
+          ReturnTypes.BIGINT_FORCE_NULLABLE,
+          PPLOperandTypes.NUMERIC);
+  public static final SqlAggFunction BIGINT_AVG =
+      createReflectiveAggFunction(
+          BigintAvgAggFunction.class,
+          "AVG",
+          SqlKind.AVG,
+          ReturnTypes.DOUBLE_NULLABLE,
+          PPLOperandTypes.NUMERIC);
   public static final SqlAggFunction TAKE =
       createUserDefinedAggFunction(
           TakeAggFunction.class,

core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java

@@ -283,6 +283,7 @@
 import java.util.Optional;
 import java.util.StringJoiner;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
 import java.util.regex.Pattern;
 import java.util.regex.PatternSyntaxException;
 import java.util.stream.Collectors;
@@ -1595,23 +1596,45 @@ void register(
     }
 
     void registerOperator(BuiltinFunctionName functionName, SqlAggFunction aggFunction) {
-      SqlOperandTypeChecker innerTypeChecker = extractTypeCheckerFromUDF(aggFunction);
+      registerOperator(functionName, aggFunction, field -> aggFunction);
+    }
+
+    void registerOperator(
+        BuiltinFunctionName functionName,
+        SqlAggFunction typeCheckerSource,
+        Function<RexNode, SqlAggFunction> aggFunctionSelector) {
+      SqlOperandTypeChecker innerTypeChecker = extractTypeCheckerFromUDF(typeCheckerSource);
       PPLTypeChecker typeChecker =
           wrapSqlOperandTypeChecker(innerTypeChecker, functionName.name(), true);
       AggHandler handler =
           (distinct, field, argList, ctx) -> {
             List<RexNode> newArgList =
                 argList.stream().map(PlanUtils::derefMapCall).collect(Collectors.toList());
             return UserDefinedFunctionUtils.makeAggregateCall(
-                aggFunction, List.of(field), newArgList, ctx.relBuilder);
+                aggFunctionSelector.apply(field), List.of(field), newArgList, ctx.relBuilder);
           };
       register(functionName, handler, typeChecker);
     }
 
+    /** Registers checked integral sums while retaining standard SUM behavior for other types. */
+    void registerSumOperator() {
+      registerOperator(
+          SUM,
+          SqlStdOperatorTable.SUM,
+          field ->
+              isIntegral(field.getType().getSqlTypeName())
+                  ? PPLBuiltinOperators.CHECKED_LONG_SUM
+                  : SqlStdOperatorTable.SUM);
+    }
+
+    private static boolean isIntegral(SqlTypeName typeName) {
+      return SqlTypeName.INT_TYPES.contains(typeName);
+    }
+
     void populate() {
       registerOperator(MAX, SqlStdOperatorTable.MAX);
       registerOperator(MIN, SqlStdOperatorTable.MIN);
-      registerOperator(SUM, SqlStdOperatorTable.SUM);
+      registerSumOperator();
       registerOperator(VARSAMP, PPLBuiltinOperators.VAR_SAMP_NULLABLE);
       registerOperator(VARPOP, PPLBuiltinOperators.VAR_POP_NULLABLE);
       registerOperator(STDDEV_SAMP, PPLBuiltinOperators.STDDEV_SAMP_NULLABLE);
@@ -1630,7 +1653,14 @@ void populate() {
 
       register(
           AVG,
-          (distinct, field, argList, ctx) -> ctx.relBuilder.avg(distinct, null, field),
+          (distinct, field, argList, ctx) -> {
+            if (field.getType().getSqlTypeName() == SqlTypeName.BIGINT) {
+              return ctx.relBuilder
+                  .aggregateCall(PPLBuiltinOperators.BIGINT_AVG, field)
+                  .distinct(distinct);
+            }
+            return ctx.relBuilder.avg(distinct, null, field);
+          },
           wrapSqlOperandTypeChecker(
               SqlStdOperatorTable.AVG.getOperandTypeChecker(), AVG.name(), false));
 

core/src/test/java/org/opensearch/sql/calcite/udf/udaf/BigintAvgAggFunctionTest.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite.udf.udaf;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+import org.apache.calcite.sql.SqlKind;
+import org.junit.jupiter.api.Test;
+import org.opensearch.sql.expression.function.PPLBuiltinOperators;
+
+class BigintAvgAggFunctionTest {
+
+  @Test
+  void retainsAvgKindForPushdownRules() {
+    assertEquals(SqlKind.AVG, PPLBuiltinOperators.BIGINT_AVG.getKind());
+  }
+
+  @Test
+  void averagesWithoutLongOverflow() {
+    BigintAvgAggFunction.Accumulator accumulator = BigintAvgAggFunction.init();
+    accumulator = BigintAvgAggFunction.add(accumulator, Long.MAX_VALUE);
+    accumulator = BigintAvgAggFunction.add(accumulator, Long.MAX_VALUE);
+
+    assertEquals((double) Long.MAX_VALUE, BigintAvgAggFunction.result(accumulator));
+  }
+
+  @Test
+  void ignoresNulls() {
+    BigintAvgAggFunction.Accumulator accumulator = BigintAvgAggFunction.init();
+    accumulator = BigintAvgAggFunction.add(accumulator, null);
+    accumulator = BigintAvgAggFunction.add(accumulator, 10L);
+
+    assertEquals(10D, BigintAvgAggFunction.result(accumulator));
+  }
+
+  @Test
+  void returnsNullForEmptyInput() {
+    assertNull(BigintAvgAggFunction.result(BigintAvgAggFunction.init()));
+  }
+}

core/src/test/java/org/opensearch/sql/calcite/udf/udaf/CheckedLongSumAggFunctionTest.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite.udf.udaf;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import org.apache.calcite.sql.SqlKind;
+import org.junit.jupiter.api.Test;
+import org.opensearch.sql.expression.function.PPLBuiltinOperators;
+
+class CheckedLongSumAggFunctionTest {
+
+  @Test
+  void retainsSumKindForPlannerRules() {
+    assertEquals(SqlKind.SUM, PPLBuiltinOperators.CHECKED_LONG_SUM.getKind());
+  }
+
+  @Test
+  void sumsExactly() {
+    long accumulator = CheckedLongSumAggFunction.init();
+    accumulator = CheckedLongSumAggFunction.add(accumulator, 1L << 62);
+    accumulator = CheckedLongSumAggFunction.add(accumulator, 1L);
+
+    assertEquals((1L << 62) + 1L, CheckedLongSumAggFunction.result(accumulator));
+  }
+
+  @Test
+  void throwsOnPositiveOverflow() {
+    assertThrows(
+        ArithmeticException.class, () -> CheckedLongSumAggFunction.add(Long.MAX_VALUE, 1L));
+  }
+
+  @Test
+  void throwsOnNegativeOverflow() {
+    assertThrows(
+        ArithmeticException.class, () -> CheckedLongSumAggFunction.add(Long.MIN_VALUE, -1L));
+  }
+}

docs/user/ppl/functions/expressions.md

@@ -13,7 +13,14 @@ Arithmetic expressions are formed by combining numeric literals and binary arith
 
 ### Overflow behavior
 
-Integer and long arithmetic operations (`+`, `-`, `*`) in `eval` expressions detect overflow and return an error instead of silently wrapping. For example, `eval x = int_field + 1` where `int_field` is `2147483647` (integer max) returns an error rather than `-2147483648`. Floating-point (`float`, `double`) arithmetic follows IEEE 754 and does not produce overflow errors.
+Long (`BIGINT`) arithmetic operations (`+`, `-`, `*`) in `eval` expressions detect overflow and return an error instead of silently wrapping. Narrower integer operands are widened before arithmetic, so crossing the 32-bit integer boundary does not overflow. Floating-point (`float`, `double`) arithmetic follows IEEE 754 and does not produce overflow errors.
+
+The accumulator used by `stats sum(integral_field)` depends on whether Calcite pushdown is enabled:
+
+- With `plugins.calcite.pushdown.enabled=true` (the default), OpenSearch uses its native double-based `sum`, then the result is checked and narrowed to `BIGINT`. Large in-range sums can lose low-order precision. A small overflow that rounds to a signed `BIGINT` boundary can also be indistinguishable from an in-range sum and may saturate at the boundary instead of returning an error.
+- With `plugins.calcite.pushdown.enabled=false`, Calcite uses an exact `BIGINT` accumulator and `Math.addExact` for every addition. It returns an error as soon as the running sum exceeds the `BIGINT` range.
+
+For example, summing `4611686018427387904` (`2^62`) and `1` returns the exact `4611686018427387905` without pushdown. With pushdown enabled, the double accumulator cannot represent the low-order `1`, so the result is `4611686018427387904`.
 
 ### Precedence
 
@@ -189,4 +196,3 @@ fetched rows / total rows = 2/2
 | 28  |
 +-----+
 ```
-