From a1db477e864d115de3aa08464b8c2d1db527ec3a Mon Sep 17 00:00:00 2001 From: David Mollitor Date: Fri, 18 Sep 2026 20:47:29 +0000 Subject: [PATCH] [SPARK-59653][SQL] Lazily materialize UnsafeMapData key/value arrays for escape analysis UnsafeMapData's constructor eagerly allocated two nested UnsafeArrayData (keys, values). That nested object graph defeats JIT escape analysis, so a transient count-only map access -- size(map), cardinality, and the size(x) > 0 filter InferFiltersFromGenerate inserts below every non-outer explode/inline -- allocates a wrapper plus two sub-wrappers per row just to read an element count that is already stored inline in the layout. UnsafeArrayData is a flat object that the JIT already scalar-replaces (so size(array_col) allocates nothing); only maps carry this per-row garbage. Build the key/value array views lazily in keyArray()/valueArray() and read numElements() directly from the key-array header, making a freshly pointed-to UnsafeMapData a flat, scalar-replaceable object. There is no storage-format change and no call-site changes; serialization, copy(), and MapData.foreach are unaffected. On a representative size(map_col) workload (whole-stage codegen on, assertions disabled), per-row getMap wrapper allocation dropped from ~4022 to ~94 JFR allocation samples (~98%), results unchanged. Generated-by: Isaac Co-authored-by: Isaac --- .../catalyst/expressions/UnsafeMapData.java | 40 +++++++++++++------ .../sql/catalyst/util/UnsafeMapSuite.scala | 36 +++++++++++++++++ 2 files changed, 64 insertions(+), 12 deletions(-) diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeMapData.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeMapData.java index 58973e7a3acf1..779c8123649f2 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeMapData.java +++ b/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; } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/UnsafeMapSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/UnsafeMapSuite.scala index 443534fd0a06a..53584ae90a828 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/UnsafeMapSuite.scala +++ b/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) + } }