(spark) branch master updated: [SPARK-59992][SQL] Make element_at and array index nullable when the index is nullable under ANSI mode

[ge...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

gengliangwang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new ba9a922fbd14 [SPARK-59992][SQL] Make element_at and array index nullable when the index is nullable under ANSI mode
ba9a922fbd14 is described below

commit ba9a922fbd14da1b5186a91184cafa5b7c50b5db
Author: richardc-db <r....@databricks.com>
AuthorDate: Thu Oct 15 09:48:34 2026 -0700

    [SPARK-59992][SQL] Make element_at and array index nullable when the index is nullable under ANSI mode
    
    ### What changes were proposed in this pull request?
    Under ANSI mode, `GetArrayItemUtil.computeNullabilityFromArray` (shared by `ElementAt` and `GetArrayItem`) now also considers `ordinal.nullable` for a non-foldable index, so `element_at(arr, idx)` and `arr[idx]` are nullable whenever the index can be null.
    
    With a non-nullable index the result keeps the existing non-nullable type, so plans that do not involve a nullable index are unchanged.
    
    ### Why are the changes needed?
    Both expressions are null-intolerant and return NULL for a NULL index, but under ANSI they reported `nullable = false` whenever the array and its elements were non-nullable (e.g. `array('a', 'b')`).
    
    The optimizer trusts that flag: `NullPropagation` truncates a `Coalesce` after its first non-nullable child, so a fallback such as `coalesce(element_at(arr, idx), 'z')` is removed. The null then lands in a column the plan treats as non-nullable, so the row carries no null marker and readers see whatever is in the slot: the query silently returns a wrong element (e.g. `'b'` instead of `'z'`) or fails with `INVALID_ARRAY_INDEX`.
    
    ### Does this PR introduce _any_ user-facing change?
    Yes, a correctness fix. With `spark.sql.ansi.enabled=true`, `SELECT coalesce(element_at(array('a', 'b'), CAST(nullif(id, 0) AS INT)), 'z') FROM range(2)` previously returned `b, a`; it now returns `z, a`. `SELECT coalesce(array('a', 'b')[CAST(nullif(id, 0) AS INT)], 'z') FROM range(2)` previously failed with `INVALID_ARRAY_INDEX` (index -1); it now returns `z, b`.
    
    `element_at` / `arr[idx]` with a nullable index now report a nullable result type under ANSI, matching non-ANSI mode. Affects released versions since Spark 3.1 (SPARK-33386).
    
    ### How was this patch tested?
    New assertions in `CollectionExpressionsSuite` (`ElementAt`) and `ComplexTypeSuite` (`GetArrayItem`): a non-null array with a nullable index is nullable, and with a non-nullable index stays non-nullable under ANSI.
    
    New end-to-end test in `DataFrameFunctionsSuite` checking that `coalesce(element_at(...), 'z')` and `coalesce(arr[...], 'z')` return `'z'` for a null index under ANSI. All three tests fail without the fix and pass with it (`build/sbt 'catalyst/testOnly *ComplexTypeSuite *CollectionExpressionsSuite'`, `build/sbt 'sql/testOnly *DataFrameFunctionsSuite -- -z SPARK-59992'`).
    
    ### Was this patch authored or co-authored using generative AI tooling?
    Generated-by: Claude Opus 5.5
    
    Closes #59244 from richardc-db/SPARK-59992-element-at-nullable-ordinal.
    
    Authored-by: richardc-db <r....@databricks.com>
    Signed-off-by: Gengliang Wang <ge...@apache.org>
---
 .../sql/catalyst/expressions/complexTypeExtractors.scala      |  2 +-
 .../sql/catalyst/expressions/CollectionExpressionsSuite.scala |  5 +++++
 .../spark/sql/catalyst/expressions/ComplexTypeSuite.scala     |  8 ++++++++
 .../scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala  | 11 +++++++++++
 4 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeExtractors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeExtractors.scala
index 2f9c1e2dd704..7b824ea3c69d 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeExtractors.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeExtractors.scala
@@ -478,7 +478,7 @@ trait GetArrayItemUtil {
           true
       }
     } else {
-      if (failOnError) arrayElementNullable || child.nullable else true
+      if (failOnError) arrayElementNullable || child.nullable || ordinal.nullable else true
     }
   }
 }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala
index f3e868cdfafd..211a47282cb9 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala
@@ -2477,6 +2477,11 @@ class CollectionExpressionsSuite
         assert(ElementAt(array, Literal(-1)).nullable)
         assert(!ElementAt(array, Subtract(Literal(2), Literal(1))).nullable)
         assert(ElementAt(array, AttributeReference("ordinal", IntegerType)()).nullable)
+        // A null ordinal returns null even when neither the array nor its elements can be null.
+        val nonNullArray = CreateArray(a :: Nil)
+        assert(ElementAt(nonNullArray, AttributeReference("ordinal", IntegerType)()).nullable)
+        assert(ElementAt(nonNullArray,
+          AttributeReference("ordinal", IntegerType, nullable = false)()).nullable == !ansiEnabled)
 
         // CreateArray case invalid indices
         assert(!ElementAt(array, Literal(0)).nullable)
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ComplexTypeSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ComplexTypeSuite.scala
index abc28ff203f8..deed423b71c5 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ComplexTypeSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ComplexTypeSuite.scala
@@ -105,6 +105,14 @@ class ComplexTypeSuite extends SparkFunSuite with ExpressionEvalHelper {
     assert(GetArrayItem(array, Literal(1)).nullable)
     assert(!GetArrayItem(array, Subtract(Literal(2), Literal(2))).nullable)
     assert(GetArrayItem(array, AttributeReference("ordinal", IntegerType)()).nullable)
+    // A null ordinal returns null even when neither the array nor its elements can be null.
+    val nonNullArray = CreateArray(a :: Nil)
+    val nullableOrdinal = AttributeReference("ordinal", IntegerType)()
+    val nonNullOrdinal = AttributeReference("ordinal", IntegerType, nullable = false)()
+    Seq(true, false).foreach { failOnError =>
+      assert(GetArrayItem(nonNullArray, nullableOrdinal, failOnError).nullable)
+      assert(GetArrayItem(nonNullArray, nonNullOrdinal, failOnError).nullable == !failOnError)
+    }
 
     // GetArrayStructFields case
     val f1 = StructField("a", IntegerType, nullable = false)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
index 79ceb6782053..504a84668068 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
@@ -2536,6 +2536,17 @@ class DataFrameFunctionsSuite extends SharedSparkSession {
     )
   }
 
+  test("SPARK-59992: coalesce over element_at with a null index is not simplified away") {
+    // Under ANSI, element_at was marked non-nullable for a non-null array regardless of the
+    // index, so the optimizer dropped the coalesce and the null index leaked through.
+    withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+      checkAnswer(
+        sql("SELECT coalesce(element_at(array('a', 'b'), CAST(nullif(id, 0) AS INT)), 'z'), " +
+          "coalesce(array('a', 'b')[CAST(nullif(id, 0) AS INT) - 1], 'z') FROM range(2)"),
+        Seq(Row("z", "z"), Row("a", "a")))
+    }
+  }
+
   test("element_at function") {
     val df = Seq(
       (Seq[String]("1", "2", "3"), 1),


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org