[SPARK-59374][SQL] Check the array size limit in concat codegen

docs/sql-migration-guide.md

@@ -29,6 +29,7 @@ license: |
 - Since Spark 4.4, a nondeterministic common expression that a conditional branch reads more than once is evaluated once, and every read sees that one value. `BETWEEN` and `NULLIF` read their input twice, so an expression such as `CASE WHEN c THEN monotonically_increasing_id() BETWEEN 3 AND 5 END` previously gave each read its own value and could return a different result. Queries whose results depended on the earlier behaviour change, and an optimized or physical plan may now show a `with` node for such a branch. A definition that is cheap to evaluate and reports itself deterministic is still substituted at every read, as before.
 - Since Spark 4.4, for storage-partitioned joins, `spark.sql.requireAllClusterKeysForCoPartition` requires every join key to be covered by some partition key instead of matching the partition keys positionally. As a result, a join-key column partitioned by more than one transform no longer prevents shuffle elimination, and `spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled` no longer additionally requires `spark.sql.requireAllClusterKeysForCoPartition` to be `false` when the join keys are a subset of the partition keys. As before, when the partition keys cover only part of the join keys, eliminating the shuffle still requires `spark.sql.requireAllClusterKeysForCoPartition` to be `false`.
 - Since Spark 4.4, when `array_repeat` or `array_insert` is asked to build an array larger than the maximum supported array length, generated code raises the same error as interpreted evaluation. `array_repeat` now fails with `COLLECTION_SIZE_LIMIT_EXCEEDED.PARAMETER` instead of the internal error `_LEGACY_ERROR_TEMP_2176`, and `array_insert` fails with `COLLECTION_SIZE_LIMIT_EXCEEDED.FUNCTION` instead of `COLLECTION_SIZE_LIMIT_EXCEEDED.PARAMETER`, which named a `count` parameter that `array_insert` does not have. Both functions raise an error under exactly the same conditions as before; only the reported error condition changes.
+- Since Spark 4.4, when `concat` is asked to build an array larger than the maximum supported array length, generated code raises the same error as interpreted evaluation, `COLLECTION_SIZE_LIMIT_EXCEEDED.FUNCTION`, instead of the internal error `_LEGACY_ERROR_TEMP_2176`. `concat` raises an error under exactly the same conditions as before; only the reported error condition changes.
 
 ## Upgrading from Spark SQL 4.2 to 4.3
 

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

@@ -3315,11 +3315,17 @@ case class Concat(children: Seq[Expression]) extends ComplexTypeMergingExpressio
   private def genCodeForNumberOfElements(ctx: CodegenContext) : (String, String) = {
     val numElements = ctx.freshName("numElements")
     val z = ctx.freshName("z")
+    // The upper bound is checked here rather than left to the array allocation so that this path
+    // reports the same error as `eval`. Without it the allocation fails with an internal error.
     val code = s"""
         |long $numElements = 0L;
         |for (int $z = 0; $z < ${children.length}; $z++) {
         |  $numElements += args[$z].numElements();
         |}
+        |if ($numElements > ${ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH}) {
+        |  throw QueryExecutionErrors.arrayFunctionWithElementsExceedLimitError(
+        |    "$prettyName", $numElements);
+        |}
       """.stripMargin
 
     (code, numElements)

sql/core/src/test/scala/org/apache/spark/sql/errors/QueryExecutionErrorsSuite.scala

@@ -39,7 +39,7 @@ import org.apache.spark.sql.catalyst.analysis.{NamedParameter, UnresolvedGenerat
 import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, RowEncoder}
 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Concat, CreateArray, EmptyRow, Expression, Flatten, Grouping, Literal, RowNumber, UnaryExpression, Years}
 import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._
-import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode}
+import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode, GenerateUnsafeProjection}
 import org.apache.spark.sql.catalyst.expressions.objects.InitializeJavaBean
 import org.apache.spark.sql.catalyst.rules.RuleIdCollection
 import org.apache.spark.sql.catalyst.util.TypeUtils.toSQLExpr
@@ -1340,18 +1340,27 @@ class QueryExecutionErrorsSuite
   test("Elements exceed limit for concat()") {
     val array = new ColumnarArray(
       new ConstantColumnVector(Int.MaxValue, BooleanType), 0, Int.MaxValue)
+    val expr = Concat(Seq(Literal.create(array, ArrayType(BooleanType))))
 
-    checkError(
-      exception = intercept[SparkRuntimeException] {
-        Concat(Seq(Literal.create(array, ArrayType(BooleanType)))).eval(EmptyRow)
-      },
-      condition = "COLLECTION_SIZE_LIMIT_EXCEEDED.FUNCTION",
-      parameters = Map(
-        "numberOfElements" -> Int.MaxValue.toString,
-        "maxRoundedArrayLength" -> MAX_ROUNDED_ARRAY_LENGTH.toString,
-        "functionName" -> toSQLId("concat")
+    // SPARK-59374: the generated code has to reject the length as well. Interpreted evaluation
+    // has always checked it, but codegen used to reach `ArrayData.allocateArrayData` and fail
+    // there with an internal error instead.
+    Seq(
+      () => expr.eval(EmptyRow),
+      () => GenerateUnsafeProjection.generate(Seq(expr)).apply(EmptyRow)
+    ).foreach { evaluate =>
+      checkError(
+        exception = intercept[SparkRuntimeException] {
+          evaluate()
+        },
+        condition = "COLLECTION_SIZE_LIMIT_EXCEEDED.FUNCTION",
+        parameters = Map(
+          "numberOfElements" -> Int.MaxValue.toString,
+          "maxRoundedArrayLength" -> MAX_ROUNDED_ARRAY_LENGTH.toString,
+          "functionName" -> toSQLId("concat")
+        )
       )
-    )
+    }
   }
 
   test("Elements exceed limit for flatten()") {