[SPARK-59653][SQL] Lazily materialize UnsafeMapData key/value arrays for escape analysis

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeMapData.java

@@ -58,18 +58,23 @@ public final class UnsafeMapData extends MapData implements Externalizable, Kryo
   public long getBaseOffset() { return baseOffset; }
   public int getSizeInBytes() { return sizeInBytes; }
 
-  private final UnsafeArrayData keys;
-  private final UnsafeArrayData values;
+  // The number of bytes of the key array, read from the 8-byte header in `pointTo`.
+  private long keyArraySize;
+
+  // The key/value array views are materialized lazily by `keyArray()`/`valueArray()`. Keeping them
+  // out of the constructor leaves a freshly pointed-to UnsafeMapData a flat object (primitives plus
+  // a base-object reference), which lets the JIT scalar-replace it via escape analysis for
+  // count-only accesses such as `numElements()` -- the common `size(map)` / `cardinality` path and
+  // the `size(x) > 0` filter inferred below explode/inline.
+  private UnsafeArrayData keys;
+  private UnsafeArrayData values;
 
   /**
    * Construct a new UnsafeMapData. The resulting UnsafeMapData won't be usable until
    * `pointTo()` has been called, since the value returned by this constructor is equivalent
    * to a null pointer.
    */
-  public UnsafeMapData() {
-    keys = new UnsafeArrayData();
-    values = new UnsafeArrayData();
-  }
+  public UnsafeMapData() {}
 
   /**
    * Update this UnsafeMapData to point to different backing data.
@@ -87,28 +92,39 @@ public void pointTo(Object baseObject, long baseOffset, int sizeInBytes) {
     final int valueArraySize = sizeInBytes - (int)keyArraySize - 8;
     assert valueArraySize >= 0 : "valueArraySize (" + valueArraySize + ") should >= 0";
 
-    keys.pointTo(baseObject, baseOffset + 8, (int)keyArraySize);
-    values.pointTo(baseObject, baseOffset + 8 + keyArraySize, valueArraySize);
-
-    assert keys.numElements() == values.numElements();
-
     this.baseObject = baseObject;
     this.baseOffset = baseOffset;
     this.sizeInBytes = sizeInBytes;
+    this.keyArraySize = keyArraySize;
+    // Defer building the key/value array views until they are actually needed (see field comment).
+    this.keys = null;
+    this.values = null;
   }
 
   @Override
   public int numElements() {
-    return keys.numElements();
+    // The key array is laid out at `baseOffset + 8`, and an UnsafeArrayData's element count is its
+    // first 8 bytes, so this reads the same value as `keyArray().numElements()` without
+    // materializing the key/value array views.
+    return (int) Platform.getLong(baseObject, baseOffset + 8);
   }
 
   @Override
   public UnsafeArrayData keyArray() {
+    if (keys == null) {
+      keys = new UnsafeArrayData();
+      keys.pointTo(baseObject, baseOffset + 8, (int) keyArraySize);
+    }
     return keys;
   }
 
   @Override
   public UnsafeArrayData valueArray() {
+    if (values == null) {
+      final int valueArraySize = sizeInBytes - (int) keyArraySize - 8;
+      values = new UnsafeArrayData();
+      values.pointTo(baseObject, baseOffset + 8 + keyArraySize, valueArraySize);
+    }
     return values;
   }
 

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/UnsafeMapSuite.scala

@@ -61,4 +61,40 @@ class UnsafeMapSuite extends SparkFunSuite {
     assert(mapDataSer.valueArray().getLong(0) == 19286)
     assert(mapDataSer.getBaseObject.asInstanceOf[Array[Byte]].length == 1024)
   }
+
+  test("numElements reads the count without materializing the key/value arrays") {
+    assert(unsafeMapData.numElements() == 1)
+    assert(unsafeMapData.numElements() == unsafeMapData.keyArray().numElements())
+  }
+
+  test("keyArray/valueArray return the map's entries") {
+    assert(unsafeMapData.keyArray().getLong(0) == 19285)
+    assert(unsafeMapData.valueArray().getLong(0) == 19286)
+  }
+
+  test("keyArray/valueArray are cached across repeated calls") {
+    assert(unsafeMapData.keyArray() eq unsafeMapData.keyArray())
+    assert(unsafeMapData.valueArray() eq unsafeMapData.valueArray())
+  }
+
+  test("copy preserves numElements and entries") {
+    val copied = unsafeMapData.copy()
+    assert(copied.numElements() == 1)
+    assert(copied.keyArray().getLong(0) == 19285)
+    assert(copied.valueArray().getLong(0) == 19286)
+  }
+
+  test("empty map has numElements 0") {
+    val baseObject = new Array[Byte](64)
+    val offset = 16
+    // Layout: [key array numBytes][empty key array][empty value array] (empty array = 8B header).
+    Platform.putLong(baseObject, offset, 8L)
+    Platform.putLong(baseObject, offset + 8, 0L)
+    Platform.putLong(baseObject, offset + 16, 0L)
+    val emptyMap = new UnsafeMapData
+    emptyMap.pointTo(baseObject, offset, 24)
+    assert(emptyMap.numElements() == 0)
+    assert(emptyMap.keyArray().numElements() == 0)
+    assert(emptyMap.valueArray().numElements() == 0)
+  }
 }