fix: [bug] narrow BIGINT to INTEGER for int-domain function arguments (#5660)

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

@@ -733,6 +733,102 @@ private static boolean containsSubQuery(RexNode node) {
     return false;
   }
 
+  /**
+   * Narrows BIGINT arguments to INTEGER at operand positions whose registered type checker declares
+   * a strict INTEGER family.
+   *
+   * <p>PPL's integer arithmetic widening (#5603) makes all integer expressions produce BIGINT.
+   * However, many Calcite runtime methods ({@code SqlFunctions.arrayItemOptional}, {@code left},
+   * {@code right}, {@code round}, {@code truncate}, etc.) require Java {@code int} parameters.
+   * Since {@code SqlTypeFamily.INTEGER} accepts BIGINT at the type-checking level, the call passes
+   * validation but code generation fails when the JVM cannot auto-narrow {@code long} to {@code
+   * int}.
+   *
+   * <p>The int-domain positions are derived from the operator's own registered type checker rather
+   * than a hardcoded per-function list: a position is narrowed only when every allowed type at that
+   * position is INTEGER (the INTEGER family maps to {@code [INTEGER]}). Value positions that
+   * legitimately accept BIGINT are unaffected because their family (e.g. NUMERIC maps to {@code
+   * [INTEGER, DOUBLE]}) is not strictly INTEGER — so {@code round(bigint_value, 2)} keeps its
+   * BIGINT first operand and only narrows the precision.
+   */
+  private static RexNode[] narrowBigintArgs(
+      RexBuilder builder, PPLTypeChecker typeChecker, RexNode... args) {
+    if (typeChecker == null) {
+      return args;
+    }
+    // Fast path: only inspect the signature when at least one argument is BIGINT. The vast
+    // majority of operator calls have no BIGINT operand and return here without building the
+    // (potentially expensive) parameter-type signature list.
+    boolean hasBigint = false;
+    for (RexNode arg : args) {
+      if (arg.getType().getSqlTypeName() == SqlTypeName.BIGINT) {
+        hasBigint = true;
+        break;
+      }
+    }
+    if (!hasBigint) {
+      return args;
+    }
+    List<List<RelDataType>> parameterTypes;
+    try {
+      parameterTypes = typeChecker.getParameterTypes();
+    } catch (IllegalArgumentException e) {
+      // A type checker whose composition is not purely family-based cannot enumerate per-position
+      // parameter types (see PPLTypeChecker composite implementations). We cannot reason about
+      // int-domain positions, so skip narrowing rather than fail an otherwise-valid query.
+      return args;
+    }
+    if (parameterTypes == null || parameterTypes.isEmpty()) {
+      return args;
+    }
+    RexNode[] narrowed = null;
+    for (int pos = 0; pos < args.length; pos++) {
+      if (args[pos].getType().getSqlTypeName() == SqlTypeName.BIGINT
+          && expectsStrictInteger(parameterTypes, pos)) {
+        if (narrowed == null) {
+          narrowed = args.clone();
+        }
+        RelDataType intType =
+            TYPE_FACTORY.createTypeWithNullability(
+                TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER), args[pos].getType().isNullable());
+        narrowed[pos] = builder.makeCast(intType, args[pos]);
+      }
+    }
+    return narrowed != null ? narrowed : args;
+  }
+
+  /**
+   * Returns true when the operand at {@code position} should be narrowed to INTEGER. This holds
+   * when INTEGER is the only numeric type the operator accepts at that position: at least one
+   * allowed type combination declares INTEGER there, and none declares a wider numeric type (BIGINT
+   * or DOUBLE).
+   *
+   * <p>Composite (OR) type checkers may list several combinations for different operand shapes
+   * (e.g. {@code ITEM} accepts {@code [ARRAY, INTEGER]} or {@code [MAP, ANY]}); ANY at a position
+   * is tolerated because it belongs to a non-numeric shape and never applies to a BIGINT argument.
+   * Value positions that accept a wider numeric type (e.g. {@code ROUND}'s NUMERIC first operand,
+   * which maps to {@code [INTEGER, DOUBLE]}) are therefore never narrowed.
+   */
+  private static boolean expectsStrictInteger(
+      List<List<RelDataType>> parameterTypes, int position) {
+    boolean sawInteger = false;
+    for (List<RelDataType> combination : parameterTypes) {
+      if (position >= combination.size()) {
+        continue;
+      }
+      SqlTypeName typeName = combination.get(position).getSqlTypeName();
+      if (typeName == SqlTypeName.INTEGER) {
+        sawInteger = true;
+      } else if (typeName == SqlTypeName.BIGINT || typeName == SqlTypeName.DOUBLE) {
+        // A wider numeric type is legitimately accepted here — this is a value position.
+        return false;
+      }
+      // Other types (ANY, ARRAY, CHARACTER, ...) belong to non-numeric operand shapes and do not
+      // apply to a BIGINT argument, so they are ignored.
+    }
+    return sawInteger;
+  }
+
   /**
    * Ad-hoc coercion for some functions that require specific casting of arguments. Now it only
    * applies to the REDUCE function.
@@ -836,7 +932,8 @@ protected void registerOperator(
         BuiltinFunctionName functionName, SqlOperator operator, PPLTypeChecker typeChecker) {
       register(
           functionName,
-          (RexBuilder builder, RexNode... args) -> builder.makeCall(operator, args),
+          (RexBuilder builder, RexNode... args) ->
+              builder.makeCall(operator, narrowBigintArgs(builder, typeChecker, args)),
           typeChecker);
     }
 
@@ -970,7 +1067,14 @@ void populate() {
           PPLTypeChecker.family(SqlTypeFamily.DATETIME, SqlTypeFamily.DATETIME));
       registerWideningIntegerOperator(MULTIPLY, SqlStdOperatorTable.MULTIPLY);
       registerWideningIntegerOperator(MULTIPLYFUNCTION, SqlStdOperatorTable.MULTIPLY);
-      registerOperator(TRUNCATE, SqlStdOperatorTable.TRUNCATE);
+      registerOperator(
+          TRUNCATE,
+          SqlStdOperatorTable.TRUNCATE,
+          PPLTypeChecker.wrapComposite(
+              (CompositeOperandTypeChecker)
+                  OperandTypes.NUMERIC.or(
+                      OperandTypes.family(SqlTypeFamily.NUMERIC, SqlTypeFamily.INTEGER)),
+              false));
       registerOperator(ASCII, SqlStdOperatorTable.ASCII);
       registerOperator(LENGTH, SqlStdOperatorTable.CHAR_LENGTH);
       registerOperator(LOWER, SqlStdOperatorTable.LOWER);
@@ -1068,7 +1172,13 @@ void populate() {
       registerOperator(POW, SqlStdOperatorTable.POWER);
       registerOperator(POWER, SqlStdOperatorTable.POWER);
       registerOperator(RADIANS, SqlStdOperatorTable.RADIANS);
-      registerOperator(RAND, SqlStdOperatorTable.RAND);
+      registerOperator(
+          RAND,
+          SqlStdOperatorTable.RAND,
+          PPLTypeChecker.wrapComposite(
+              (CompositeOperandTypeChecker)
+                  OperandTypes.NILADIC.or(OperandTypes.family(SqlTypeFamily.INTEGER)),
+              false));
       // TODO, workaround to support sequence CompositeOperandTypeChecker.
       registerOperator(
           ROUND,
@@ -1299,7 +1409,10 @@ void populate() {
       registerOperator(MAP_CONCAT, SqlLibraryOperators.MAP_CONCAT);
       registerOperator(MAP_REMOVE, PPLBuiltinOperators.MAP_REMOVE);
       registerOperator(ARRAY_LENGTH, SqlLibraryOperators.ARRAY_LENGTH);
-      registerOperator(ARRAY_SLICE, SqlLibraryOperators.ARRAY_SLICE);
+      registerOperator(
+          ARRAY_SLICE,
+          SqlLibraryOperators.ARRAY_SLICE,
+          PPLTypeChecker.family(SqlTypeFamily.ARRAY, SqlTypeFamily.INTEGER, SqlTypeFamily.INTEGER));
       registerOperator(ARRAY_COMPACT, SqlLibraryOperators.ARRAY_COMPACT);
       registerOperator(FORALL, PPLBuiltinOperators.FORALL);
       registerOperator(EXISTS, PPLBuiltinOperators.EXISTS);

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

@@ -580,6 +580,78 @@ public void testMvindexWithStatsAggregationPushdown() throws IOException {
     verifySchema(actual, schema("count()", "bigint"), schema("e", "string"));
   }
 
+  @Test
+  public void testMvindexWithLiteralArithmeticIndex() throws IOException {
+    // mvindex with user arithmetic as index (1 + 1). PPL widens arithmetic to BIGINT,
+    // so the ITEM index must be narrowed back to INTEGER at the serialization boundary.
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | eval arr = array('a', 'b', 'c'), result = mvindex(arr, 1 + 1)"
+                    + " | head 1 | fields result",
+                TEST_INDEX_BANK));
+
+    verifySchema(actual, schema("result", "string"));
+    verifyDataRows(actual, rows("c"));
+  }
+
+  @Test
+  public void testMvindexWithEvalDerivedArithmeticIndex() throws IOException {
+    // mvindex where the index comes from an eval-derived arithmetic value (BIGINT).
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | eval a = 1 + 1, arr = array('a', 'b', 'c'), result = mvindex(arr, a)"
+                    + " | head 1 | fields result",
+                TEST_INDEX_BANK));
+
+    verifySchema(actual, schema("result", "string"));
+    verifyDataRows(actual, rows("c"));
+  }
+
+  @Test
+  public void testMvindexWithFieldDerivedArithmeticIndex() throws IOException {
+    // mvindex where the index is derived from an integer field via arithmetic (age - N).
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | where age = 32 | eval arr = array('a', 'b', 'c', 'd', 'e'),"
+                    + " result = mvindex(arr, age - 30) | head 1 | fields result",
+                TEST_INDEX_BANK));
+
+    verifySchema(actual, schema("result", "string"));
+    verifyDataRows(actual, rows("c"));
+  }
+
+  @Test
+  public void testMvindexRangeWithArithmeticIndices() throws IOException {
+    // mvindex range access with arithmetic start and end indices (BIGINT).
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | eval arr = array(1, 2, 3, 4, 5), result = mvindex(arr, 1 + 0, 2 + 1)"
+                    + " | head 1 | fields result",
+                TEST_INDEX_BANK));
+
+    verifySchema(actual, schema("result", "array"));
+    verifyDataRows(actual, rows(List.of(2, 3, 4)));
+  }
+
+  @Test
+  public void testMvindexWithCastLongIndex() throws IOException {
+    // mvindex with an explicit cast(x as long) index. Pre-existing bug: a genuinely-BIGINT value
+    // handed to ITEM's int parameter must be narrowed to INTEGER.
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | eval arr = array('a', 'b', 'c'), result = mvindex(arr, cast(1 as"
+                    + " long)) | head 1 | fields result",
+                TEST_INDEX_BANK));
+
+    verifySchema(actual, schema("result", "string"));
+    verifyDataRows(actual, rows("b"));
+  }
+
   @Test
   public void testMvfindWithMatch() throws IOException {
     JSONObject actual =

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

@@ -5,6 +5,12 @@
 
 package org.opensearch.sql.calcite.remote;
 
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK;
+import static org.opensearch.sql.util.MatcherUtils.*;
+
+import java.io.IOException;
+import org.json.JSONObject;
+import org.junit.jupiter.api.Test;
 import org.opensearch.sql.ppl.MathematicalFunctionIT;
 
 public class CalciteMathematicalFunctionIT extends MathematicalFunctionIT {
@@ -13,4 +19,72 @@ public void init() throws Exception {
     super.init();
     enableCalcite();
   }
+
+  @Test
+  public void testRoundWithArithmeticPrecision() throws IOException {
+    // ROUND with arithmetic precision argument. PPL arithmetic widens to BIGINT but ROUND
+    // expects int.
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | eval result = round(123.456, 1 + 0) | head 1 | fields result",
+                TEST_INDEX_BANK));
+
+    verifySchema(actual, schema("result", null, "double"));
+    verifyDataRows(actual, rows(123.5));
+  }
+
+  @Test
+  public void testTruncateWithArithmeticPrecision() throws IOException {
+    // TRUNCATE with arithmetic precision argument. PPL arithmetic widens to BIGINT but TRUNCATE
+    // expects int.
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | eval result = truncate(123.456, 1 + 0) | head 1 | fields result",
+                TEST_INDEX_BANK));
+
+    verifySchema(actual, schema("result", null, "double"));
+    verifyDataRows(actual, rows(123.4));
+  }
+
+  @Test
+  public void testRoundWithCastLongPrecision() throws IOException {
+    // ROUND with an explicit cast(x as long) precision. Pre-existing bug: a genuinely-BIGINT
+    // value handed to ROUND's int parameter must be narrowed to INTEGER.
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | eval result = round(123.456, cast(1 as long)) | head 1 | fields"
+                    + " result",
+                TEST_INDEX_BANK));
+
+    verifySchema(actual, schema("result", null, "double"));
+    verifyDataRows(actual, rows(123.5));
+  }
+
+  @Test
+  public void testConvWithArithmeticBases() throws IOException {
+    // CONV with arithmetic base arguments. Both widen to BIGINT but CONV expects int radixes.
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | eval result = conv('11', 1 + 1, 8 + 2) | head 1 | fields result",
+                TEST_INDEX_BANK));
+
+    verifySchema(actual, schema("result", null, "string"));
+    verifyDataRows(actual, rows("3"));
+  }
+
+  @Test
+  public void testSha2WithArithmeticBitLength() throws IOException {
+    // SHA2 with arithmetic bit-length argument. Widens to BIGINT but SHA2 expects int.
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | eval result = sha2('abc', 128 + 128) | head 1 | fields result",
+                TEST_INDEX_BANK));
+
+    verifySchema(actual, schema("result", null, "string"));
+  }
 }

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

@@ -341,6 +341,21 @@ public void testRand() throws IOException {
     verifyDataRows(actual, rows("Jake"), rows("Hello"), rows("Jane"), rows("John"));
   }
 
+  @Test
+  public void testRandWithArithmeticSeed() throws IOException {
+    // RAND with an arithmetic seed. PPL arithmetic widens 1 + 1 to BIGINT, but RAND's seed is a
+    // Java int parameter — the seed must be narrowed to INTEGER or codegen fails.
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | eval rand = rand(1 + 1) | where rand >= 0 | where rand < 1 | fields"
+                    + " name",
+                TEST_INDEX_STATE_COUNTRY));
+
+    verifySchema(actual, schema("name", "string"));
+    verifyDataRows(actual, rows("Jake"), rows("Hello"), rows("Jane"), rows("John"));
+  }
+
   @Test
   public void testPowInvalidArgShouldReturnNull() throws IOException {
     JSONObject actual =