Skip to content

Commit 7f002fe

Browse files
SEPURI-SAI-KRISHNAuros-b
SEPURI-SAI-KRISHNA
authored and
uros-b
committed
[SPARK-59374][SQL] Check the array size limit in concat codegen
### What changes were proposed in this pull request? `Concat.genCodeForNumberOfElements` accumulates the result element count into a `long` and hands it to `CodeGenerator.createArrayData` without checking it against `ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH`. Interpreted evaluation does check. This PR adds the same check to the generated code so both paths raise the same error. This is the `concat` half of the pattern that [SPARK-58631](https://issues.apache.org/jira/browse/SPARK-58631) (#57832) fixed for `array_repeat`. `flatten` has the identical gap and is tracked separately as [SPARK-59375](https://issues.apache.org/jira/browse/SPARK-59375). ### Why are the changes needed? The two evaluation paths report different error conditions for the same input: - Interpreted: `COLLECTION_SIZE_LIMIT_EXCEEDED.FUNCTION`, from `QueryExecutionErrors.arrayFunctionWithElementsExceedLimitError`. - Codegen: no check, so execution reaches `ArrayData.allocateArrayData`, whose own bound check raises the internal `_LEGACY_ERROR_TEMP_2176` through `QueryExecutionErrors.cannotCreateArrayWithElementsExceedLimitError`. An internal error condition is not meant to be user facing, and which one a query gets should not depend on whether codegen was used. Worth noting for reviewers that this is latent rather than something a user hits today. Unlike `array_repeat`, where the element count comes straight from an argument, `concat` needs its inputs to actually hold that many elements, so an executor will normally run out of memory before the limit is reached. The existing test reaches it only by using a `ConstantColumnVector`, which reports its length without allocating. ### Does this PR introduce _any_ user-facing change? Yes, in the sense that the error condition changes for an input that already fails. `concat` with more than `MAX_ROUNDED_ARRAY_LENGTH` total elements now raises `COLLECTION_SIZE_LIMIT_EXCEEDED.FUNCTION` under codegen, where it previously raised the internal `_LEGACY_ERROR_TEMP_2176`. That matches what interpreted evaluation has always raised. The set of inputs that fail is unchanged. Added a migration guide entry alongside the one #57832 added for `array_repeat` and `array_insert`, since this is the same kind of change. Happy to drop it if you would rather not document a case that is this hard to reach. ### How was this patch tested? Extended the existing `Elements exceed limit for concat()` test in `QueryExecutionErrorsSuite` to run the expression through `GenerateUnsafeProjection` as well as `eval`, asserting the same error from both. That test already builds a large array cheaply with `ConstantColumnVector`, which reports its length without allocating, so nothing large is allocated. It sits in `QueryExecutionErrorsSuite` rather than in `CollectionExpressionsSuite`, where #57832 put its tests, because `ConstantColumnVector` lives in `sql/core` and is not reachable from a catalyst suite. Extending the existing test also seemed better than adding a second one: the test was already there and simply never exercised codegen, which is how the gap survived. Reverting the change makes the new codegen case fail with `_LEGACY_ERROR_TEMP_2176`, which is the behaviour being fixed. ``` build/sbt "sql/testOnly org.apache.spark.sql.errors.QueryExecutionErrorsSuite -- -z \"Elements exceed limit\"" ``` ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 5) Closes #58703 from SEPURI-SAI-KRISHNA/SPARK-59374-concat-codegen-limit. Authored-by: Sepuri Sai Krishna <[email protected]> Signed-off-by: Uros Bojanic <[email protected]> (cherry picked from commit fc091f5) Signed-off-by: Uros Bojanic <[email protected]>
1 parent 9332d4a commit 7f002fe

2 files changed

Lines changed: 26 additions & 11 deletions

File tree

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3372,11 +3372,17 @@ case class Concat(children: Seq[Expression]) extends ComplexTypeMergingExpressio
33723372
private def genCodeForNumberOfElements(ctx: CodegenContext) : (String, String) = {
33733373
val numElements = ctx.freshName("numElements")
33743374
val z = ctx.freshName("z")
3375+
// The upper bound is checked here rather than left to the array allocation so that this path
3376+
// reports the same error as `eval`. Without it the allocation fails with an internal error.
33753377
val code = s"""
33763378
|long $numElements = 0L;
33773379
|for (int $z = 0; $z < ${children.length}; $z++) {
33783380
| $numElements += args[$z].numElements();
33793381
|}
3382+
|if ($numElements > ${ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH}) {
3383+
| throw QueryExecutionErrors.arrayFunctionWithElementsExceedLimitError(
3384+
| "$prettyName", $numElements);
3385+
|}
33803386
""".stripMargin
33813387

33823388
(code, numElements)

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

Lines changed: 20 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ import org.apache.spark.sql.catalyst.analysis.{NamedParameter, UnresolvedGenerat
3939
import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, RowEncoder}
4040
import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Concat, CreateArray, EmptyRow, Expression, Flatten, Grouping, Literal, RowNumber, UnaryExpression, Years}
4141
import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._
42-
import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode}
42+
import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode, GenerateUnsafeProjection}
4343
import org.apache.spark.sql.catalyst.expressions.objects.InitializeJavaBean
4444
import org.apache.spark.sql.catalyst.rules.RuleIdCollection
4545
import org.apache.spark.sql.catalyst.util.TypeUtils.toSQLExpr
@@ -1340,18 +1340,27 @@ class QueryExecutionErrorsSuite
13401340
test("Elements exceed limit for concat()") {
13411341
val array = new ColumnarArray(
13421342
new ConstantColumnVector(Int.MaxValue, BooleanType), 0, Int.MaxValue)
1343+
val expr = Concat(Seq(Literal.create(array, ArrayType(BooleanType))))
13431344

1344-
checkError(
1345-
exception = intercept[SparkRuntimeException] {
1346-
Concat(Seq(Literal.create(array, ArrayType(BooleanType)))).eval(EmptyRow)
1347-
},
1348-
condition = "COLLECTION_SIZE_LIMIT_EXCEEDED.FUNCTION",
1349-
parameters = Map(
1350-
"numberOfElements" -> Int.MaxValue.toString,
1351-
"maxRoundedArrayLength" -> MAX_ROUNDED_ARRAY_LENGTH.toString,
1352-
"functionName" -> toSQLId("concat")
1345+
// SPARK-59374: the generated code has to reject the length as well. Interpreted evaluation
1346+
// has always checked it, but codegen used to reach `ArrayData.allocateArrayData` and fail
1347+
// there with an internal error instead.
1348+
Seq(
1349+
() => expr.eval(EmptyRow),
1350+
() => GenerateUnsafeProjection.generate(Seq(expr)).apply(EmptyRow)
1351+
).foreach { evaluate =>
1352+
checkError(
1353+
exception = intercept[SparkRuntimeException] {
1354+
evaluate()
1355+
},
1356+
condition = "COLLECTION_SIZE_LIMIT_EXCEEDED.FUNCTION",
1357+
parameters = Map(
1358+
"numberOfElements" -> Int.MaxValue.toString,
1359+
"maxRoundedArrayLength" -> MAX_ROUNDED_ARRAY_LENGTH.toString,
1360+
"functionName" -> toSQLId("concat")
1361+
)
13531362
)
1354-
)
1363+
}
13551364
}
13561365

13571366
test("Elements exceed limit for flatten()") {

0 commit comments

Comments
 (0)