[SPARK-40403][SQL] Calculate unsafe array size using longs to avoid negative size in error message

core/src/main/resources/error/error-classes.json

@@ -423,6 +423,11 @@
     ],
     "sqlState" : "22023"
   },
+  "TOO_MANY_ARRAY_ELEMENTS" : {
+    "message" : [
+      "Cannot initialize array with <numElements> elements of size <size>"
+    ]
+  },
   "UNABLE_TO_ACQUIRE_MEMORY" : {
     "message" : [
       "Unable to acquire <requestedBytes> bytes of memory, got <receivedBytes>"

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/codegen/UnsafeArrayWriter.java

@@ -17,6 +17,7 @@
 
 package org.apache.spark.sql.catalyst.expressions.codegen;
 
+import org.apache.spark.sql.errors.QueryExecutionErrors;
 import org.apache.spark.sql.types.Decimal;
 import org.apache.spark.unsafe.Platform;
 import org.apache.spark.unsafe.array.ByteArrayMethods;
@@ -55,10 +56,18 @@ public void initialize(int numElements) {
 
     this.startingOffset = cursor();
 
+    long fixedPartInBytesLong =
+      ByteArrayMethods.roundNumberOfBytesToNearestWord((long) elementSize * numElements);
+    long totalInitialSize = headerInBytes + fixedPartInBytesLong;
+
+    if (totalInitialSize > Integer.MAX_VALUE) {
+      throw QueryExecutionErrors.tooManyArrayElementsError(numElements, elementSize);
+    }
+
+    // it's now safe to cast fixedPartInBytesLong and totalInitialSize to int
+    int fixedPartInBytes = (int) fixedPartInBytesLong;
     // Grows the global buffer ahead for header and fixed size data.
-    int fixedPartInBytes =
-      ByteArrayMethods.roundNumberOfBytesToNearestWord(elementSize * numElements);
-    holder.grow(headerInBytes + fixedPartInBytes);
+    holder.grow((int)totalInitialSize);
 
     // Write numElements and clear out null bits to header
     Platform.putLong(getBuffer(), startingOffset, numElements);

sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala

@@ -2122,4 +2122,14 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase {
         "functionName" -> toSQLId(funcName),
         "expected" -> pattern))
   }
+
+  def tooManyArrayElementsError(
+      numElements: Int,
+      elementSize: Int): SparkIllegalArgumentException = {
+    new SparkIllegalArgumentException(
+      errorClass = "TOO_MANY_ARRAY_ELEMENTS",
+      messageParameters = Map(
+        "numElements" -> numElements.toString,
+        "size" -> elementSize.toString))
+  }
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/codegen/UnsafeArrayWriterSuite.scala

@@ -0,0 +1,31 @@
+/*
+ * 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.codegen
+
+import org.apache.spark.{SparkFunSuite, SparkIllegalArgumentException}
+
+class UnsafeArrayWriterSuite extends SparkFunSuite {
+  test("SPARK-40403: don't print negative number when array is too big") {
+    val rowWriter = new UnsafeRowWriter(1)
+    rowWriter.resetRowWriter()
+    val arrayWriter = new UnsafeArrayWriter(rowWriter, 8)
+    assert(intercept[SparkIllegalArgumentException] {
+      arrayWriter.initialize(268271216)
+    }.getMessage.contains("Cannot initialize array with 268271216 elements of size 8"))
+  }
+}