Detect BIGINT overflow in SUM and fix wrong AVG on the Calcite engine

SUM and AVG over a BIGINT (long) column near 2^63 silently produced wrong
results on the Calcite engine:
  - SUM(long): the enumerable accumulator is a plain long, so a running sum
    past 2^63 wrapped to a negative value (e.g. SUM(UserID) returned
    -5594372458244005145 instead of erroring).
  - AVG(long): AVG reduces to SUM(field)/COUNT(field); the intermediate long
    SUM wrapped the same way, so AVG returned a wrong (often negative) result.

Fix, applied in the SQL-plugin lowering so both the DSL-pushdown and the
in-memory (no-pushdown) paths behave identically:
  - SUM(long): the aggregate argument (a bare BIGINT column) is summed in
    DECIMAL so it cannot wrap, and the result is narrowed back to BIGINT with
    CHECKED_LONG_NARROW. Narrowing errors (HTTP 4xx) when the sum clearly
    overflowed 2^63, and otherwise saturates. The output type stays bigint.
  - AVG(long): the bare BIGINT argument is averaged in DOUBLE, which holds the
    true average without an intermediate long wrap. The DOUBLE output type is
    unchanged.

Only bare BIGINT column references are widened; expressions such as
sum(field + 1) are left untouched so PPLAggregateConvertRule can still rewrite
them to pushdown-friendly form. Narrower integer sums (byte/short/int) already
widen to a BIGINT accumulator and cannot overflow, so they are unchanged.

Boundary note: OpenSearch computes pushed-down sums in double, and
(double) Long.MAX_VALUE rounds to exactly 2^63, indistinguishable from a small
overflow. To avoid erroring on a legitimate near-Long.MAX_VALUE sum, only
magnitudes strictly beyond 2^63 are treated as overflow; a genuine overflow
lands far outside that boundary, and the in-memory path detects it exactly.

AVG pushdown note: casting the AVG argument to DOUBLE keeps native avg pushdown
in most cases, but an AVG with a preceding sort can fall back to an in-memory
sum/count reduction instead of a pushed-down native avg. The result stays
correct; only that narrow case loses aggregate pushdown.

Adds CalcitePPLAggregationIT coverage (overflow -> 4xx, avg correct, in-range
and Long.MAX sums), a REST yaml test (issues/5164_agg.yml), and regenerates the
affected explain fixtures.

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

Author: ahkcs

core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java

@@ -188,6 +188,7 @@
 import org.opensearch.sql.datasource.DataSourceService;
 import org.opensearch.sql.exception.CalciteUnsupportedException;
 import org.opensearch.sql.exception.SemanticCheckException;
+import org.opensearch.sql.executor.OpenSearchTypeSystem;
 import org.opensearch.sql.expression.HighlightExpression;
 import org.opensearch.sql.expression.function.BuiltinFunctionName;
 import org.opensearch.sql.expression.function.PPLBuiltinOperators;
@@ -1763,6 +1764,12 @@ private void visitAggregation(
     // Add aggregation results first
     List<RexNode> aggRexList =
         outputFields.subList(numOfOutputFields - numOfAggList, numOfOutputFields);
+    // SUM(bigint) accumulates in DECIMAL (OpenSearchTypeSystem#deriveSumType) so a running sum near
+    // 2^63 cannot silently wrap to a negative long. Cast the user-visible result back to BIGINT:
+    // the cast errors (ArithmeticException -> 4xx) on genuine overflow while keeping the output
+    // type bigint. AVG is unaffected (its output stays DOUBLE and its reduced internal SUM is not
+    // surfaced here).
+    aggRexList = castBigintSumResults(context, aggExprList, aggRexList);
     List<RexNode> aliasedGroupByList =
         aggregationAttributes.getLeft().stream()
             .map(this::extractAliasLiteral)
@@ -1812,6 +1819,60 @@ private static AggregateFunction extractAggregateFunction(UnresolvedExpression e
     return null;
   }
 
+  /**
+   * Narrow the results of {@code SUM(bigint)} aggregations back to BIGINT. {@code SUM} over a
+   * BIGINT column is accumulated in a wider type (its argument is cast to DECIMAL in {@code
+   * PPLFuncImpTable}) so a running sum near 2^63 cannot silently wrap to a negative long. This
+   * restores the declared BIGINT output type via {@link PPLBuiltinOperators#CHECKED_LONG_NARROW},
+   * which errors on a genuine overflow (surfaced as a 4xx) while saturating a value that merely
+   * rounds to 2^63 on the double-based pushdown path, so both execution paths behave identically.
+   * Non-SUM aggregations and sums whose accumulator is not the widened DECIMAL (e.g. {@code
+   * SUM(double)}, user {@code SUM(decimal)}) are returned unchanged.
+   */
+  private static List<RexNode> castBigintSumResults(
+      CalcitePlanContext context,
+      List<UnresolvedExpression> aggExprList,
+      List<RexNode> aggRexList) {
+    if (aggExprList.size() != aggRexList.size()) {
+      return aggRexList;
+    }
+    List<RexNode> result = new ArrayList<>(aggRexList.size());
+    for (int i = 0; i < aggRexList.size(); i++) {
+      RexNode aggRex = aggRexList.get(i);
+      AggregateFunction aggFunc = extractAggregateFunction(aggExprList.get(i));
+      if (aggFunc != null
+          && BuiltinFunctionName.ofAggregation(aggFunc.getFuncName())
+              .filter(name -> name == BuiltinFunctionName.SUM)
+              .isPresent()
+          && isWidenedBigintSumType(aggRex.getType())) {
+        RexNode narrowed =
+            context.rexBuilder.makeCall(PPLBuiltinOperators.CHECKED_LONG_NARROW, aggRex);
+        // Wrapping in the narrowing UDF drops the aggregate's output name; restore it so downstream
+        // projection and the result schema keep the original "sum(...)" column name.
+        result.add(
+            aggRex instanceof RexInputRef ref
+                ? context.relBuilder.alias(
+                    narrowed,
+                    context.relBuilder.peek().getRowType().getFieldNames().get(ref.getIndex()))
+                : narrowed);
+      } else {
+        result.add(aggRex);
+      }
+    }
+    return result;
+  }
+
+  /**
+   * Whether the type is the DECIMAL accumulator produced for {@code SUM(bigint)} (a scale-0 decimal
+   * at max precision, from the argument widening in {@code PPLFuncImpTable}), as opposed to a
+   * genuine user-supplied {@code SUM(decimal)} with a fractional scale.
+   */
+  private static boolean isWidenedBigintSumType(RelDataType type) {
+    return type.getSqlTypeName() == SqlTypeName.DECIMAL
+        && type.getScale() == 0
+        && type.getPrecision() == OpenSearchTypeSystem.MAX_PRECISION;
+  }
+
   private static Function extractFunction(UnresolvedExpression expr) {
     if (expr instanceof Function f) return f;
     if (expr instanceof Alias alias) return extractFunction(alias.getDelegated());

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

@@ -65,6 +65,7 @@
 import org.opensearch.sql.expression.function.jsonUDF.JsonSetFunctionImpl;
 import org.opensearch.sql.expression.function.udf.AutoConvertFunction;
 import org.opensearch.sql.expression.function.udf.CTimeConvertFunction;
+import org.opensearch.sql.expression.function.udf.CheckedLongNarrowFunction;
 import org.opensearch.sql.expression.function.udf.CryptographicFunction;
 import org.opensearch.sql.expression.function.udf.Dur2SecConvertFunction;
 import org.opensearch.sql.expression.function.udf.MemkConvertFunction;
@@ -436,6 +437,8 @@ 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");
+  public static final SqlOperator CHECKED_LONG_NARROW =
+      new CheckedLongNarrowFunction().toUDF("CHECKED_LONG_NARROW");
 
   // PPL Convert command functions
   public static final SqlOperator AUTO = new AutoConvertFunction().toUDF("AUTO");

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

@@ -287,6 +287,7 @@
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexInputRef;
 import org.apache.calcite.rex.RexLambda;
 import org.apache.calcite.rex.RexLambdaRef;
 import org.apache.calcite.rex.RexLiteral;
@@ -1537,10 +1538,68 @@ void registerOperator(BuiltinFunctionName functionName, SqlAggFunction aggFuncti
       register(functionName, handler, typeChecker);
     }
 
+    /**
+     * Registers {@code SUM} with the extra behaviour that a BIGINT (long) column is summed in
+     * DECIMAL rather than long. A running long sum near 2^63 silently wraps to a negative value
+     * (the enumerable {@code SumImplementor} generates a plain {@code long + long}); accumulating
+     * in DECIMAL is exact. {@link CalciteRelNodeVisitor} casts the DECIMAL result back to BIGINT so
+     * the output type is unchanged and a genuine overflow surfaces as a client error instead of a
+     * wrong value.
+     *
+     * <p>Only a bare BIGINT column reference is widened. Expressions like {@code sum(field + 1)}
+     * are left as-is so {@link org.opensearch.sql.calcite.plan.rule.PPLAggregateConvertRule} can
+     * still rewrite them to pushdown-friendly {@code sum(field) OP literal} form.
+     */
+    void registerSumOperator() {
+      SqlOperandTypeChecker innerTypeChecker = extractTypeCheckerFromUDF(SqlStdOperatorTable.SUM);
+      PPLTypeChecker typeChecker = wrapSqlOperandTypeChecker(innerTypeChecker, SUM.name(), true);
+      AggHandler handler =
+          (distinct, field, argList, ctx) -> {
+            List<RexNode> newArgList =
+                argList.stream().map(PlanUtils::derefMapCall).collect(Collectors.toList());
+            RexNode sumField = widenBigintColumnToDecimal(field, ctx);
+            return UserDefinedFunctionUtils.makeAggregateCall(
+                SqlStdOperatorTable.SUM, List.of(sumField), newArgList, ctx.relBuilder);
+          };
+      register(SUM, handler, typeChecker);
+    }
+
+    /**
+     * If {@code field} is a bare BIGINT column reference, cast it to DECIMAL so the aggregate that
+     * consumes it accumulates exactly instead of wrapping a long. Any other expression is returned
+     * unchanged.
+     */
+    private static RexNode widenBigintColumnToDecimal(RexNode field, CalcitePlanContext ctx) {
+      if (field instanceof RexInputRef && field.getType().getSqlTypeName() == SqlTypeName.BIGINT) {
+        int maxPrecision = TYPE_FACTORY.getTypeSystem().getMaxNumericPrecision();
+        RelDataType decimalType =
+            TYPE_FACTORY.createTypeWithNullability(
+                TYPE_FACTORY.createSqlType(SqlTypeName.DECIMAL, maxPrecision, 0),
+                field.getType().isNullable());
+        return ctx.rexBuilder.makeCast(decimalType, field);
+      }
+      return field;
+    }
+
+    /**
+     * If {@code field} is a bare BIGINT column reference, cast it to DOUBLE so that AVG's reduced
+     * intermediate SUM accumulates in double rather than a long that wraps near 2^63. Any other
+     * expression is returned unchanged.
+     */
+    private static RexNode widenBigintColumnToDouble(RexNode field, CalcitePlanContext ctx) {
+      if (field instanceof RexInputRef && field.getType().getSqlTypeName() == SqlTypeName.BIGINT) {
+        RelDataType doubleType =
+            TYPE_FACTORY.createTypeWithNullability(
+                TYPE_FACTORY.createSqlType(SqlTypeName.DOUBLE), field.getType().isNullable());
+        return ctx.rexBuilder.makeCast(doubleType, field);
+      }
+      return field;
+    }
+
     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);
@@ -1559,7 +1618,11 @@ void populate() {
 
       register(
           AVG,
-          (distinct, field, argList, ctx) -> ctx.relBuilder.avg(distinct, null, field),
+          (distinct, field, argList, ctx) ->
+              // AVG reduces to SUM(field)/COUNT(field); over a bare BIGINT column the intermediate
+              // long SUM wraps near 2^63 (returning a wrong negative average). Averaging in DOUBLE
+              // avoids the wrap and keeps the DOUBLE output type unchanged.
+              ctx.relBuilder.avg(distinct, null, widenBigintColumnToDouble(field, ctx)),
           wrapSqlOperandTypeChecker(
               SqlStdOperatorTable.AVG.getOperandTypeChecker(), AVG.name(), false));
 

core/src/main/java/org/opensearch/sql/expression/function/udf/CheckedLongNarrowFunction.java

@@ -0,0 +1,78 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.expression.function.udf;
+
+import java.math.BigDecimal;
+import java.util.List;
+import org.apache.calcite.adapter.enumerable.NotNullImplementor;
+import org.apache.calcite.adapter.enumerable.NullPolicy;
+import org.apache.calcite.adapter.enumerable.RexToLixTranslator;
+import org.apache.calcite.linq4j.tree.Expression;
+import org.apache.calcite.linq4j.tree.Expressions;
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.sql.type.ReturnTypes;
+import org.apache.calcite.sql.type.SqlReturnTypeInference;
+import org.opensearch.sql.calcite.utils.PPLOperandTypes;
+import org.opensearch.sql.expression.function.ImplementorUDF;
+import org.opensearch.sql.expression.function.UDFOperandMetadata;
+
+/**
+ * Narrows a numeric SUM accumulator (widened to DECIMAL/double so it does not wrap) back to BIGINT,
+ * raising an error when the value overflowed the BIGINT range instead of wrapping to a negative
+ * value.
+ *
+ * <p>SUM over a BIGINT column is accumulated as a wider type to avoid the silent {@code long}
+ * two's-complement wrap near 2^63. This function restores the declared BIGINT output type. A value
+ * whose magnitude clearly exceeds {@code 2^63} is treated as overflow and surfaces as a client
+ * error. Because OpenSearch computes pushed-down sums in {@code double}, {@code (double)
+ * Long.MAX_VALUE} rounds to exactly {@code 2^63} and cannot be distinguished from a small overflow;
+ * to avoid erroring on a legitimate near-{@code Long.MAX_VALUE} sum, only magnitudes strictly
+ * beyond {@code 2^63} are rejected, and in-range values saturate on narrowing rather than wrap.
+ * This makes the pushdown and in-memory paths behave identically.
+ */
+public class CheckedLongNarrowFunction extends ImplementorUDF {
+  public CheckedLongNarrowFunction() {
+    super(new CheckedLongNarrowImplementor(), NullPolicy.ANY);
+  }
+
+  @Override
+  public SqlReturnTypeInference getReturnTypeInference() {
+    return ReturnTypes.BIGINT_FORCE_NULLABLE;
+  }
+
+  @Override
+  public UDFOperandMetadata getOperandMetadata() {
+    return PPLOperandTypes.NUMERIC;
+  }
+
+  /** {@code 2^63}, the first magnitude beyond the signed BIGINT range. */
+  private static final double TWO_POW_63 = 0x1p63;
+
+  /** Narrow a widened sum value to a long, erroring when it overflowed the BIGINT range. */
+  public static long narrow(Object value) {
+    double magnitude = ((Number) value).doubleValue();
+    if (magnitude > TWO_POW_63 || magnitude < -TWO_POW_63) {
+      throw new ArithmeticException("BIGINT overflow in SUM");
+    }
+    if (value instanceof BigDecimal decimal) {
+      // Exact for the in-memory (DECIMAL) path; clamps a value that rounds to exactly 2^63.
+      return decimal
+          .min(BigDecimal.valueOf(Long.MAX_VALUE))
+          .max(BigDecimal.valueOf(Long.MIN_VALUE))
+          .longValue();
+    }
+    // double path (pushed-down sum): saturating narrow (JLS 5.1.3) clamps 2^63 to Long.MAX_VALUE.
+    return (long) magnitude;
+  }
+
+  public static class CheckedLongNarrowImplementor implements NotNullImplementor {
+    @Override
+    public Expression implement(
+        RexToLixTranslator translator, RexCall call, List<Expression> translatedOperands) {
+      return Expressions.call(CheckedLongNarrowFunction.class, "narrow", translatedOperands.get(0));
+    }
+  }
+}

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLAggregationIT.java

@@ -20,12 +20,16 @@
 import static org.opensearch.sql.util.MatcherUtils.verifyErrorMessageContains;
 import static org.opensearch.sql.util.MatcherUtils.verifySchema;
 import static org.opensearch.sql.util.MatcherUtils.verifySchemaInOrder;
+import static org.opensearch.sql.util.TestUtils.createIndexByRestClient;
+import static org.opensearch.sql.util.TestUtils.isIndexExist;
+import static org.opensearch.sql.util.TestUtils.performRequest;
 
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.List;
 import org.json.JSONObject;
 import org.junit.jupiter.api.Test;
+import org.opensearch.client.Request;
 import org.opensearch.sql.common.utils.StringUtils;
 import org.opensearch.sql.exception.SemanticCheckException;
 import org.opensearch.sql.ppl.PPLIntegTestCase;
@@ -93,6 +97,70 @@ public void testSumAvg() throws IOException {
     verifyDataRows(actual, rows(186973));
   }
 
+  @Test
+  public void testSumAvgLongOverflow() throws IOException {
+    String overflowIndex = "test_sum_long_overflow";
+    if (!isIndexExist(client(), overflowIndex)) {
+      createIndexByRestClient(
+          client(), overflowIndex, "{\"mappings\":{\"properties\":{\"v\":{\"type\":\"long\"}}}}");
+      Request bulk = new Request("POST", "/" + overflowIndex + "/_bulk?refresh=true");
+      bulk.setJsonEntity(
+          "{\"index\":{}}\n"
+              + "{\"v\":9223372036854775807}\n"
+              + "{\"index\":{}}\n"
+              + "{\"v\":9223372036854775807}\n"
+              + "{\"index\":{}}\n"
+              + "{\"v\":9223372036854775807}\n");
+      performRequest(client(), bulk);
+    }
+    String inRangeIndex = "test_sum_long_in_range";
+    if (!isIndexExist(client(), inRangeIndex)) {
+      createIndexByRestClient(
+          client(), inRangeIndex, "{\"mappings\":{\"properties\":{\"v\":{\"type\":\"long\"}}}}");
+      Request bulk = new Request("POST", "/" + inRangeIndex + "/_bulk?refresh=true");
+      bulk.setJsonEntity(
+          "{\"index\":{}}\n"
+              + "{\"v\":1000000000000}\n"
+              + "{\"index\":{}}\n"
+              + "{\"v\":2000000000000}\n"
+              + "{\"index\":{}}\n"
+              + "{\"v\":3000000000000}\n");
+      performRequest(client(), bulk);
+    }
+    String boundaryIndex = "test_sum_long_boundary";
+    if (!isIndexExist(client(), boundaryIndex)) {
+      createIndexByRestClient(
+          client(), boundaryIndex, "{\"mappings\":{\"properties\":{\"v\":{\"type\":\"long\"}}}}");
+      Request bulk = new Request("POST", "/" + boundaryIndex + "/_bulk?refresh=true");
+      bulk.setJsonEntity("{\"index\":{}}\n" + "{\"v\":9223372036854775807}\n");
+      performRequest(client(), bulk);
+    }
+
+    // SUM overflows the BIGINT range (3 * (2^63 - 1)); surfaced as a client error rather than
+    // silently wrapping to a negative value.
+    Throwable error =
+        assertThrowsWithReplace(
+            RuntimeException.class,
+            () -> executeQuery(String.format("source=%s | stats sum(v)", overflowIndex)));
+    verifyErrorMessageContains(error, "verflow");
+
+    // AVG is averaged in DOUBLE, so it holds the true average (the shared value) without wrapping.
+    JSONObject avg = executeQuery(String.format("source=%s | stats avg(v)", overflowIndex));
+    verifySchema(avg, schema("avg(v)", "double"));
+    verifyDataRows(avg, rows(9.223372036854776e18));
+
+    // A sum well within the BIGINT range returns the exact value with no error.
+    JSONObject inRange = executeQuery(String.format("source=%s | stats sum(v)", inRangeIndex));
+    verifySchema(inRange, schema("sum(v)", "bigint"));
+    verifyDataRows(inRange, rows(6000000000000L));
+
+    // A single Long.MAX_VALUE is a valid (non-overflowing) sum and must not error, even though the
+    // pushed-down double result rounds to exactly 2^63 — the narrow saturates instead of erroring.
+    JSONObject boundary = executeQuery(String.format("source=%s | stats sum(v)", boundaryIndex));
+    verifySchema(boundary, schema("sum(v)", "bigint"));
+    verifyDataRows(boundary, rows(9223372036854775807L));
+  }
+
   @Test
   public void testAsExistedField() throws IOException {
     JSONObject actual =