Skip to content

Commit 12ab2c5

Browse files
gengliangwanggengliangwang
committed
[SPARK-56935][SQL] Simplify GetArrayItem codegen and consolidate ElementAtUtils into ArrayExpressionUtils
### What changes were proposed in this pull request? Two related changes: 1. Fold `ElementAtUtils.resolveArrayIndex` into the existing `ArrayExpressionUtils.java`, and remove `ElementAtUtils.java`. The per-expression naming chosen in SPARK-56916 didn't match the codebase's category-scoped utility-class convention (`ArrayExpressionUtils`, `BitmapExpressionUtils`, `ExpressionImplUtils`, ...) and there's now a natural home for any future array-expression ANSI helper. 2. Refactor `GetArrayItem`'s ANSI codegen + eval paths to use a new `ArrayExpressionUtils.checkArrayIndex(int length, int index, QueryContext context)` helper, mirroring how `ElementAt` uses `resolveArrayIndex`. The helper throws `invalidArrayIndexError` for negative / out-of-bound ANSI indices and returns the validated 0-based position so the caller chains into `arr.get(idx, dataType)`. The non-ANSI branch keeps its inline form because it must return `null` (not throw) on out-of-bound. Net effect: the existing per-expression `ElementAtUtils.java` is removed; the existing `ArrayExpressionUtils.java` grows two `*ArrayIndex` helpers used by `ElementAt` and `GetArrayItem` codegen + eval. ### Why are the changes needed? Part of SPARK-56908 (umbrella). `arr[idx]` and `element_at(arr, idx)` share the same ANSI out-of-bound error shape; collapsing both into one-line helper calls keeps the codegen size small and avoids maintaining two parallel inline forms. ### Does this PR introduce _any_ user-facing change? No. The compiled behavior is identical; only the emitted Java source text changes. ### How was this patch tested? ``` build/sbt "catalyst/testOnly *ComplexTypeSuite *CollectionExpressionsSuite" build/sbt "sql/testOnly *QueryExecutionAnsiErrorsSuite" ``` All pass (83/83 catalyst, 21/21 sql). ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Cursor 1.x Closes #55973 from gengliangwang/SPARK-56935-getarrayitem. Authored-by: Gengliang Wang <gengliang@apache.org> Signed-off-by: Gengliang Wang <gengliang@apache.org>
1 parent af2cfc0 commit 12ab2c5

4 files changed

Lines changed: 95 additions & 89 deletions

File tree

‎sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ArrayExpressionUtils.java‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,53 @@
1919
import java.util.Arrays;
2020
import java.util.Comparator;
2121

22+
import org.apache.spark.QueryContext;
2223
import org.apache.spark.sql.catalyst.util.SQLOrderingUtil;
24+
import org.apache.spark.sql.errors.QueryExecutionErrors;
2325

2426
public class ArrayExpressionUtils {
2527

28+
// ANSI index helpers used by ArrayType expression codegen and eval paths.
29+
30+
/**
31+
* Resolves the user-supplied 1-based {@code element_at} index to a
32+
* 0-based array position. Throws when the absolute index exceeds the
33+
* array length (ANSI out-of-bounds) or when {@code index} is zero
34+
* (always invalid).
35+
*
36+
* @param length the array length
37+
* @param index the 1-based index supplied by the user (positive or negative)
38+
* @param context the query context attached to the error
39+
* @return the resolved 0-based position
40+
*/
41+
public static int resolveArrayIndex(int length, int index, QueryContext context) {
42+
if (length < Math.abs(index)) {
43+
throw QueryExecutionErrors.invalidElementAtIndexError(index, length, context);
44+
}
45+
if (index == 0) {
46+
throw QueryExecutionErrors.invalidIndexOfZeroError(context);
47+
}
48+
return index > 0 ? index - 1 : length + index;
49+
}
50+
51+
/**
52+
* Validates a 0-based {@code arr[idx]} index against the array length
53+
* under ANSI mode. Throws when {@code index} is negative or
54+
* {@code >= length}; otherwise returns {@code index} unchanged so the
55+
* caller can chain into {@code arr.get(idx, dataType)}.
56+
*
57+
* @param length the array length
58+
* @param index the 0-based index supplied by the user
59+
* @param context the query context attached to the error
60+
* @return the validated 0-based position (== {@code index})
61+
*/
62+
public static int checkArrayIndex(int length, int index, QueryContext context) {
63+
if (index < 0 || index >= length) {
64+
throw QueryExecutionErrors.invalidArrayIndexError(index, length, context);
65+
}
66+
return index;
67+
}
68+
2669
// comparator
2770
// Boolean ascending nullable comparator
2871
private static final Comparator<Boolean> booleanComp = (o1, o2) -> {

‎sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ElementAtUtils.java‎

Lines changed: 0 additions & 51 deletions
This file was deleted.

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2743,7 +2743,7 @@ case class ElementAt(
27432743
case _: ArrayType if failOnError =>
27442744
(value, ordinal) => {
27452745
val array = value.asInstanceOf[ArrayData]
2746-
val idx = ElementAtUtils.resolveArrayIndex(
2746+
val idx = ArrayExpressionUtils.resolveArrayIndex(
27472747
array.numElements(), ordinal.asInstanceOf[Int], getContextOrNull())
27482748
if (arrayElementNullable && array.isNullAt(idx)) null else array.get(idx, dataType)
27492749
}
@@ -2783,7 +2783,7 @@ case class ElementAt(
27832783
nullSafeCodeGen(ctx, ev, (eval1, eval2) => {
27842784
val index = ctx.freshName("elementAtIndex")
27852785
val errorContext = getContextOrNullCode(ctx)
2786-
val utils = classOf[ElementAtUtils].getName
2786+
val utils = classOf[ArrayExpressionUtils].getName
27872787
val assignment = s"${ev.value} = ${CodeGenerator.getValue(eval1, dataType, index)};"
27882788
val body = if (arrayElementNullable) {
27892789
s"""

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

Lines changed: 50 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ import org.apache.spark.sql.catalyst.analysis._
2525
import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodeGenerator, ExprCode}
2626
import org.apache.spark.sql.catalyst.trees.TreePattern.{EXTRACT_VALUE, TreePattern}
2727
import org.apache.spark.sql.catalyst.util.{quoteIdentifier, ArrayData, GenericArrayData, MapData, TypeUtils}
28-
import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase, QueryExecutionErrors}
28+
import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase}
2929
import org.apache.spark.sql.internal.SQLConf
3030
import org.apache.spark.sql.types._
3131

@@ -356,13 +356,11 @@ case class GetArrayItem(
356356
protected override def nullSafeEval(value: Any, ordinal: Any): Any = {
357357
val baseValue = value.asInstanceOf[ArrayData]
358358
val index = ordinal.asInstanceOf[Number].intValue()
359-
if (index >= baseValue.numElements() || index < 0) {
360-
if (failOnError) {
361-
throw QueryExecutionErrors.invalidArrayIndexError(
362-
index, baseValue.numElements(), getContextOrNull())
363-
} else {
364-
null
365-
}
359+
if (failOnError) {
360+
ArrayExpressionUtils.checkArrayIndex(baseValue.numElements(), index, getContextOrNull())
361+
if (baseValue.isNullAt(index)) null else baseValue.get(index, dataType)
362+
} else if (index >= baseValue.numElements() || index < 0) {
363+
null
366364
} else if (baseValue.isNullAt(index)) {
367365
null
368366
} else {
@@ -371,36 +369,52 @@ case class GetArrayItem(
371369
}
372370

373371
override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
374-
nullSafeCodeGen(ctx, ev, (eval1, eval2) => {
375-
val index = ctx.freshName("index")
376-
val childArrayElementNullable = child.dataType.asInstanceOf[ArrayType].containsNull
377-
val nullCheck = if (childArrayElementNullable) {
378-
s"""else if ($eval1.isNullAt($index)) {
379-
${ev.isNull} = true;
380-
}
381-
"""
382-
} else {
383-
""
384-
}
385-
386-
val indexOutOfBoundBranch = if (failOnError) {
372+
// ANSI (failOnError) and non-ANSI paths generate different codegen.
373+
if (failOnError) {
374+
nullSafeCodeGen(ctx, ev, (eval1, eval2) => {
375+
val index = ctx.freshName("index")
387376
val errorContext = getContextOrNullCode(ctx)
388-
// scalastyle:off line.size.limit
389-
s"throw QueryExecutionErrors.invalidArrayIndexError($index, $eval1.numElements(), $errorContext);"
390-
// scalastyle:on line.size.limit
391-
} else {
392-
s"${ev.isNull} = true;"
393-
}
394-
395-
s"""
396-
final int $index = (int) $eval2;
397-
if ($index >= $eval1.numElements() || $index < 0) {
398-
$indexOutOfBoundBranch
399-
} $nullCheck else {
400-
${ev.value} = ${CodeGenerator.getValue(eval1, dataType, index)};
377+
val utils = classOf[ArrayExpressionUtils].getName
378+
val childArrayElementNullable = child.dataType.asInstanceOf[ArrayType].containsNull
379+
val assignment = s"${ev.value} = ${CodeGenerator.getValue(eval1, dataType, index)};"
380+
val body = if (childArrayElementNullable) {
381+
s"""
382+
|if ($eval1.isNullAt($index)) {
383+
| ${ev.isNull} = true;
384+
|} else {
385+
| $assignment
386+
|}
387+
""".stripMargin
388+
} else {
389+
assignment
401390
}
402-
"""
403-
})
391+
s"""
392+
|int $index = $utils.checkArrayIndex($eval1.numElements(), (int) $eval2, $errorContext);
393+
|$body
394+
""".stripMargin
395+
})
396+
} else {
397+
nullSafeCodeGen(ctx, ev, (eval1, eval2) => {
398+
val index = ctx.freshName("index")
399+
val childArrayElementNullable = child.dataType.asInstanceOf[ArrayType].containsNull
400+
val nullCheck = if (childArrayElementNullable) {
401+
s"""else if ($eval1.isNullAt($index)) {
402+
${ev.isNull} = true;
403+
}
404+
"""
405+
} else {
406+
""
407+
}
408+
s"""
409+
final int $index = (int) $eval2;
410+
if ($index >= $eval1.numElements() || $index < 0) {
411+
${ev.isNull} = true;
412+
} $nullCheck else {
413+
${ev.value} = ${CodeGenerator.getValue(eval1, dataType, index)};
414+
}
415+
"""
416+
})
417+
}
404418
}
405419

406420
override protected def withNewChildrenInternal(

0 commit comments

Comments
 (0)