[SPARK-59598][SQL] GetMapValue and ElementAt must not match a null map key

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

@@ -547,9 +547,13 @@ trait GetMapValueUtil extends BinaryExpression with ImplicitCastInputTypes {
     val hm = new java.util.HashMap[Any, Int]((len * 1.5).toInt)
     var i = 0
     while (i < len) {
-      // putIfAbsent preserves first-match semantics for maps with duplicate keys (allowed at
-      // the physical level by [[ArrayBasedMapData]]), matching the linear scan path.
-      hm.putIfAbsent(keys.get(i, keyType), i)
+      // Null keys are skipped: the lookup key is never null (both eval and codegen are
+      // null-safe on the ordinal), so a null key can never match. See [[LinearExecutor]].
+      if (!keys.isNullAt(i)) {
+        // putIfAbsent preserves first-match semantics for maps with duplicate keys (allowed at
+        // the physical level by [[ArrayBasedMapData]]), matching the linear scan path.
+        hm.putIfAbsent(keys.get(i, keyType), i)
+      }
       i += 1
     }
     hm
@@ -572,12 +576,18 @@ trait GetMapValueUtil extends BinaryExpression with ImplicitCastInputTypes {
     val mask = cap - 1
     var i = 0
     while (i < len) {
-      var h = hashKeyOnDriver(keys.get(i, keyType), keyType) & mask
-      // Open addressing with linear probing; duplicates take the next free slot so that the
-      // lookup (which stops at the first match) returns the first-inserted index -- matches
-      // [[buildHashIndex]] / [[ArrayBasedMapData]] first-wins semantics.
-      while (buckets(h) != -1) h = (h + 1) & mask
-      buckets(h) = i
+      // Null keys are skipped, as in [[buildHashIndex]]. This is load-bearing here: the
+      // generated probe reads a candidate key with a primitive getter, which on a null slot
+      // returns the type's zero value, so a bucket for a null key would let a lookup of 0
+      // match it.
+      if (!keys.isNullAt(i)) {
+        var h = hashKeyOnDriver(keys.get(i, keyType), keyType) & mask
+        // Open addressing with linear probing; duplicates take the next free slot so that the
+        // lookup (which stops at the first match) returns the first-inserted index -- matches
+        // [[buildHashIndex]] / [[ArrayBasedMapData]] first-wins semantics.
+        while (buckets(h) != -1) h = (h + 1) & mask
+        buckets(h) = i
+      }
       i += 1
     }
     (buckets, mask)
@@ -651,7 +661,12 @@ trait GetMapValueUtil extends BinaryExpression with ImplicitCastInputTypes {
       var i = 0
       var found = false
       while (i < length && !found) {
-        if (ordering.equiv(keys.get(i, keyType), ordinal)) {
+        // Skip null keys. `ordinal` is never null here (the caller is null-safe on both
+        // inputs), so a null key can never be the one being looked up. The guard is required
+        // rather than cosmetic: `ordering` is the key type's natural ordering, and for a
+        // primitive key type it unboxes its arguments, turning a null key into the type's
+        // zero value -- which would make a lookup of 0 match it.
+        if (!keys.isNullAt(i) && ordering.equiv(keys.get(i, keyType), ordinal)) {
           found = true
         } else {
           i += 1
@@ -693,6 +708,11 @@ trait GetMapValueUtil extends BinaryExpression with ImplicitCastInputTypes {
            |int $index = -1;
            |
            |for (int $i = 0; $i < $length; $i++) {
+           |  // Skip null keys: the getter below is not null-aware, so a null key would read
+           |  // back as the java type's default value and could match the lookup key.
+           |  if ($keys.isNullAt($i)) {
+           |    continue;
+           |  }
            |  $keyJavaType $loopKey = ${CodeGenerator.getValue(keys, keyType, i)};
            |  if (${ctx.genEqual(keyType, loopKey, eval2)}) {
            |    $index = $i;

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ComplexTypeSuite.scala

@@ -399,6 +399,63 @@ class ComplexTypeSuite extends SparkFunSuite with ExpressionEvalHelper {
     }
   }
 
+  test("SPARK-59598: map lookup must not match a null key") {
+    // A map whose key array contains a null. `ArrayBasedMapBuilder` rejects null keys, but the
+    // file-format readers build `ArrayBasedMapData` directly and do not: see the "the parquet
+    // map may contains null or duplicated map keys" note in `ParquetRowConverter`. A lookup
+    // must never match such a key, whichever executor and code path is used.
+    def mapWithNullKey(keys: Array[Any], values: Array[Any]): ArrayBasedMapData =
+      new ArrayBasedMapData(new GenericArrayData(keys), new GenericArrayData(values))
+
+    // Primitive keys are the dangerous case: a null slot read with a primitive getter (codegen)
+    // or unboxed by the natural ordering (interpreted) yields 0, so `m[0]` used to return the
+    // null key's value instead of null.
+    val intKeyMap = mapWithNullKey(Array(null, 1), Array(10, 20))
+    val intMapType = MapType(IntegerType, IntegerType)
+
+    // Non-foldable input -> LinearExecutor, regardless of threshold.
+    withSQLConf(SQLConf.MAP_LOOKUP_HASH_THRESHOLD.key -> "0") {
+      val mapRef = BoundReference(0, intMapType, nullable = true)
+      val row = create_row(intKeyMap)
+      assert(!GetMapValue(mapRef, Literal(0)).usesFoldableHashLookup)
+
+      checkEvaluation(GetMapValue(mapRef, Literal(0)), null, row)
+      checkEvaluation(ElementAt(mapRef, Literal(0)), null, row)
+      // Non-null keys in the same map still resolve, i.e. the scan skips the null slot rather
+      // than stopping at it.
+      checkEvaluation(GetMapValue(mapRef, Literal(1)), 20, row)
+      checkEvaluation(ElementAt(mapRef, Literal(1)), 20, row)
+    }
+
+    // Foldable input above the threshold -> PrebuiltHashExecutor. The generated probe reads a
+    // candidate key with a primitive getter, so a bucket built for a null key would also match
+    // a lookup of 0.
+    withSQLConf(SQLConf.MAP_LOOKUP_HASH_THRESHOLD.key -> "0") {
+      val foldable = Literal.create(intKeyMap, intMapType)
+      assert(GetMapValue(foldable, Literal(0)).usesFoldableHashLookup)
+
+      checkEvaluation(GetMapValue(foldable, Literal(0)), null)
+      checkEvaluation(ElementAt(foldable, Literal(0)), null)
+      checkEvaluation(GetMapValue(foldable, Literal(1)), 20)
+      checkEvaluation(ElementAt(foldable, Literal(1)), 20)
+    }
+
+    // Non-primitive keys: the interpreted ordering would compare against null and the generated
+    // code would call `equals` on a value read from a null slot.
+    val stringKeyMap = mapWithNullKey(
+      Array(null, UTF8String.fromString("a")),
+      Array(UTF8String.fromString("x"), UTF8String.fromString("y")))
+    val stringMapType = MapType(StringType, StringType)
+    withSQLConf(SQLConf.MAP_LOOKUP_HASH_THRESHOLD.key -> "0") {
+      val mapRef = BoundReference(0, stringMapType, nullable = true)
+      val row = create_row(stringKeyMap)
+      checkEvaluation(GetMapValue(mapRef, Literal("a")), "y", row)
+      checkEvaluation(GetMapValue(mapRef, Literal("missing")), null, row)
+      checkEvaluation(ElementAt(mapRef, Literal("a")), "y", row)
+      checkEvaluation(ElementAt(mapRef, Literal("missing")), null, row)
+    }
+  }
+
   test("GetMapValue - strategy choice for foldable maps") {
     // Build a foldable map literal large enough to clear the default threshold. The
     // strategy assertions here pair with the non-foldable test above: together they lock in