[SPARK-33391][SQL] element_at with CreateArray not respect one based index.

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

@@ -1966,7 +1966,20 @@ case class ElementAt(left: Expression, right: Expression)
   }
 
   override def nullable: Boolean = left.dataType match {
-    case _: ArrayType => computeNullabilityFromArray(left, right)
+    case _: ArrayType =>
+      def specialNormalizeIndex: (Int, Int) => Int = {
+        (arrayLength: Int, index: Int) => {
+          if (index < 0) {
+            arrayLength + index
+          } else if (index == 0) {
+            // make it default TRUE.
+            arrayLength
+          } else {
+            index - 1
+          }
+        }
+      }
+      computeNullabilityFromArray(left, right, normalizeIndex = specialNormalizeIndex)
     case _: MapType => true
   }
 

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

@@ -273,15 +273,19 @@ case class GetArrayItem(child: Expression, ordinal: Expression)
 trait GetArrayItemUtil {
 
   /** `Null` is returned for invalid ordinals. */
-  protected def computeNullabilityFromArray(child: Expression, ordinal: Expression): Boolean = {
+  protected def computeNullabilityFromArray(
+      child: Expression,
+      ordinal: Expression,
+      normalizeIndex: (Int, Int) => Int = (_: Int, index: Int) => index): Boolean = {
+
     if (ordinal.foldable && !ordinal.nullable) {
       val intOrdinal = ordinal.eval().asInstanceOf[Number].intValue()
       child match {
-        case CreateArray(ar, _) if intOrdinal < ar.length =>
-          ar(intOrdinal).nullable
+        case CreateArray(ar, _) if normalizeIndex(ar.length, intOrdinal) < ar.length =>
+          ar(normalizeIndex(ar.length, intOrdinal)).nullable
         case GetArrayStructFields(CreateArray(elements, _), field, _, _, _)
-          if intOrdinal < elements.length =>
-          elements(intOrdinal).nullable || field.nullable
+          if normalizeIndex(elements.length, intOrdinal) < elements.length =>
+          elements(normalizeIndex(elements.size, intOrdinal)).nullable || field.nullable
         case _ =>
           true
       }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala

@@ -1122,9 +1122,9 @@ class CollectionExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper
     val a = AttributeReference("a", IntegerType, nullable = false)()
     val b = AttributeReference("b", IntegerType, nullable = true)()
     val array = CreateArray(a :: b :: Nil)
-    assert(!ElementAt(array, Literal(0)).nullable)
-    assert(ElementAt(array, Literal(1)).nullable)
-    assert(!ElementAt(array, Subtract(Literal(2), Literal(2))).nullable)
+    assert(!ElementAt(array, Literal(1)).nullable)
+    assert(ElementAt(array, Literal(2)).nullable)
+    assert(!ElementAt(array, Subtract(Literal(2), Literal(1))).nullable)
     assert(ElementAt(array, AttributeReference("ordinal", IntegerType)()).nullable)
 
     // GetArrayStructFields case
@@ -1135,19 +1135,19 @@ class CollectionExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper
     val inputArray1 = CreateArray(c :: Nil)
     val inputArray1ContainsNull = c.nullable
     val stArray1 = GetArrayStructFields(inputArray1, f1, 0, 2, inputArray1ContainsNull)
-    assert(!ElementAt(stArray1, Literal(0)).nullable)
+    assert(!ElementAt(stArray1, Literal(1)).nullable)
     val stArray2 = GetArrayStructFields(inputArray1, f2, 1, 2, inputArray1ContainsNull)
-    assert(ElementAt(stArray2, Literal(0)).nullable)
+    assert(ElementAt(stArray2, Literal(1)).nullable)
 
     val d = AttributeReference("d", structType, nullable = true)()
     val inputArray2 = CreateArray(c :: d :: Nil)
     val inputArray2ContainsNull = c.nullable || d.nullable
     val stArray3 = GetArrayStructFields(inputArray2, f1, 0, 2, inputArray2ContainsNull)
-    assert(!ElementAt(stArray3, Literal(0)).nullable)
-    assert(ElementAt(stArray3, Literal(1)).nullable)
+    assert(!ElementAt(stArray3, Literal(1)).nullable)
+    assert(ElementAt(stArray3, Literal(2)).nullable)
     val stArray4 = GetArrayStructFields(inputArray2, f2, 1, 2, inputArray2ContainsNull)
-    assert(ElementAt(stArray4, Literal(0)).nullable)
     assert(ElementAt(stArray4, Literal(1)).nullable)
+    assert(ElementAt(stArray4, Literal(2)).nullable)
   }
 
   test("Concat") {

sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala

@@ -1401,6 +1401,40 @@ class DataFrameFunctionsSuite extends QueryTest with SharedSparkSession {
     assert(e3.message.contains(errorMsg3))
   }
 
+  test("SPARK-33391: element_at with CreateArray") {
+    // element_at should use one-based index and support negative index.
+    // valid index for array(1, 2, 3) should be 1,2,3,-1,-2,-3
+    var df = OneRowRelation().selectExpr("element_at(array(1, 2, 3), 1)")
+    assert(!df.schema.head.nullable)
+    checkAnswer(
+      df,
+      Seq(Row(1))
+    )
+
+    df = OneRowRelation().selectExpr("element_at(array(1, 2, 3), -1)")
+    assert(!df.schema.head.nullable)
+    checkAnswer(
+      df,
+      Seq(Row(3))
+    )
+
+    df = OneRowRelation().selectExpr("element_at(array(1, 2, 3), 3)")
+    assert(!df.schema.head.nullable)
+    checkAnswer(
+      df,
+      Seq(Row(3))
+    )
+
+    // 0 is not a valid index, return default nullable which is 'TRUE'.
+    df = OneRowRelation().selectExpr("element_at(array(1, 2, 3), 0)")
+    assert(df.schema.head.nullable)
+
+    val ex = intercept[ArrayIndexOutOfBoundsException] {
+      df.collect()
+    }
+    assert(ex.getMessage.contains("SQL array indices start at 1"))
+  }
+
   test("array_union functions") {
     val df1 = Seq((Array(1, 2, 3), Array(4, 2))).toDF("a", "b")
     val ans1 = Row(Seq(1, 2, 3, 4))