[SPARK-56916][SQL] Simplify ElementAt array codegen under ANSI mode

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ArrayUtils.java

@@ -0,0 +1,56 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions;
+
+import org.apache.spark.QueryContext;
+import org.apache.spark.sql.errors.QueryExecutionErrors;
+
+/**
+ * Static helpers for array-related expressions invoked from
+ * {@code doGenCode} and {@code eval} paths.
+ *
+ * Currently provides the ANSI-mode index validation for
+ * {@link ElementAt} on {@code ArrayType}: a single call replaces ~12 lines
+ * of inline length / zero / sign-normalization codegen with a return of
+ * the normalized array position (0-based).
+ */
+public final class ArrayUtils {
+
+  private ArrayUtils() {}
+
+  /**
+   * Validates a 1-based {@code element_at} index against the array length
+   * and returns the 0-based position. Throws when the absolute index
+   * exceeds the array length (ANSI out-of-bounds) or when {@code index} is
+   * zero (always invalid).
+   *
+   * @param length  the array length
+   * @param index   the 1-based index supplied by the user (positive or negative)
+   * @param context the query context attached to the error
+   * @return        the validated 0-based position
+   */
+  public static int elementAtIndexExact(int length, int index, QueryContext context) {
+    if (length < Math.abs(index)) {
+      throw QueryExecutionErrors.invalidElementAtIndexError(index, length, context);
+    }
+    if (index == 0) {
+      throw QueryExecutionErrors.invalidIndexOfZeroError(context);
+    }
+    return index > 0 ? index - 1 : length + index;
+  }
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala

@@ -2738,19 +2738,21 @@ case class ElementAt(
   override def nullSafeEval(value: Any, ordinal: Any): Any = doElementAt(value, ordinal)
 
   @transient private lazy val doElementAt: (Any, Any) => Any = left.dataType match {
+    case _: ArrayType if failOnError =>
+      (value, ordinal) => {
+        val array = value.asInstanceOf[ArrayData]
+        val idx = ArrayUtils.elementAtIndexExact(
+          array.numElements(), ordinal.asInstanceOf[Int], getContextOrNull())
+        if (arrayElementNullable && array.isNullAt(idx)) null else array.get(idx, dataType)
+      }
     case _: ArrayType =>
       (value, ordinal) => {
         val array = value.asInstanceOf[ArrayData]
         val index = ordinal.asInstanceOf[Int]
         if (array.numElements() < math.abs(index)) {
-          if (failOnError) {
-            throw QueryExecutionErrors.invalidElementAtIndexError(
-              index, array.numElements(), getContextOrNull())
-          } else {
-            defaultValueOutOfBound match {
-              case Some(value) => value.eval()
-              case None => null
-            }
+          defaultValueOutOfBound match {
+            case Some(value) => value.eval()
+            case None => null
           }
         } else {
           val idx = if (index == 0) {
@@ -2773,7 +2775,7 @@ case class ElementAt(
 
   override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
     left.dataType match {
-      case _: ArrayType =>
+      case _: ArrayType if failOnError =>
         nullSafeCodeGen(ctx, ev, (eval1, eval2) => {
           val index = ctx.freshName("elementAtIndex")
           val nullCheck = if (arrayElementNullable) {
@@ -2786,21 +2788,38 @@ case class ElementAt(
             ""
           }
           val errorContext = getContextOrNullCode(ctx)
-          val indexOutOfBoundBranch = if (failOnError) {
-            // scalastyle:off line.size.limit
-            s"throw QueryExecutionErrors.invalidElementAtIndexError($index, $eval1.numElements(), $errorContext);"
-            // scalastyle:on line.size.limit
+          val utils = classOf[ArrayUtils].getName
+          s"""
+             |int $index = $utils.elementAtIndexExact(
+             |  $eval1.numElements(), (int) $eval2, $errorContext);
+             |$nullCheck
+             |{
+             |  ${ev.value} = ${CodeGenerator.getValue(eval1, dataType, index)};
+             |}
+           """.stripMargin
+        })
+      case _: ArrayType =>
+        nullSafeCodeGen(ctx, ev, (eval1, eval2) => {
+          val index = ctx.freshName("elementAtIndex")
+          val nullCheck = if (arrayElementNullable) {
+            s"""
+               |if ($eval1.isNullAt($index)) {
+               |  ${ev.isNull} = true;
+               |} else
+             """.stripMargin
           } else {
-            defaultValueOutOfBound match {
-              case Some(value) =>
-                val defaultValueEval = value.genCode(ctx)
-                s"""
-                  ${defaultValueEval.code}
-                  ${ev.isNull} = ${defaultValueEval.isNull};
-                  ${ev.value} = ${defaultValueEval.value};
-                """.stripMargin
-              case None => s"${ev.isNull} = true;"
-            }
+            ""
+          }
+          val errorContext = getContextOrNullCode(ctx)
+          val indexOutOfBoundBranch = defaultValueOutOfBound match {
+            case Some(value) =>
+              val defaultValueEval = value.genCode(ctx)
+              s"""
+                ${defaultValueEval.code}
+                ${ev.isNull} = ${defaultValueEval.isNull};
+                ${ev.value} = ${defaultValueEval.value};
+              """.stripMargin
+            case None => s"${ev.isNull} = true;"
           }
 
           s"""