spark-test commented on code in PR #60048:
URL: https://github.com/apache/spark/pull/60048#discussion_r4469915394
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/joins/HashedRelationSuite.scala:
##########
@@ -272,6 +272,135 @@ abstract class HashedRelationSuite extends
SharedSparkSession {
map2.free()
}
+ test("LongToUnsafeRowMap shrinks an over-allocated sparse array at optimize
time") {
+ val unsafeProj = UnsafeProjection.create(Seq(BoundReference(0, LongType,
false)))
+
+ // Many rows share 3 keys spread far apart, so the initial array is sized
by the row
+ // count (2 * nextPow2(N) longs) while the load factor invariant of
updateIndex only
+ // needs 16 longs for 3 keys. Keys far apart also keep optimize() from
switching to
+ // the dense mode, which is preferred whenever it applies.
+ val keys = Array(-1000000000000L, 0L, 1000000000000L)
+ val N = 99999
+ val map = new LongToUnsafeRowMap(mm, N)
+ var i = 0
+ while (i < N) {
+ val k = keys(i % 3)
+ map.append(k, unsafeProj(InternalRow(k)))
+ i += 1
+ }
+
+ // Capture the page size via the difference, so the assertions do not
depend on the
+ // details of page growth.
+ var n = 1
+ while (n < N) n *= 2
+ val pageBytes = map.getTotalMemoryConsumption - n * 2L * 8
+ map.optimize()
+ // The smallest power-of-two array keeping numKeys * 4 <= array.size() for
3 keys.
+ assert(map.getTotalMemoryConsumption === pageBytes + 16L * 8)
+
+ val row = unsafeProj(InternalRow(0L)).copy()
+ keys.foreach { k =>
+ val r = map.get(k, row)
+ var c = 0
+ while (r.hasNext) {
+ assert(r.next().getLong(0) === k)
+ c += 1
+ }
+ assert(c === N / 3)
+ }
+ assert(map.getValue(12345L, row) eq null)
+ // keys() has production consumers on broadcast relations (e.g. dynamic
partition
+ // pruning), so the shrunk map must still serve it.
+ assert(map.keys().map(_.getLong(0)).toSet === keys.toSet)
+ // A second optimize() must be a no-op: the shrunk array is already
minimal.
+ map.optimize()
+ assert(map.getTotalMemoryConsumption === pageBytes + 16L * 8)
+ map.free()
+ }
+
+ test("LongToUnsafeRowMap keeps a right-sized sparse array unchanged") {
+ val unsafeProj = UnsafeProjection.create(Seq(BoundReference(0, LongType,
false)))
+
+ // With unique keys the growth in updateIndex/growArray already leaves the
array at
+ // the minimal size for the load factor invariant, so optimize() must not
rehash.
+ val N = 100000
+ val keys = (0 until N).map(i => i * 0x9E3779B97F4A7C15L).toArray
+ val map = new LongToUnsafeRowMap(mm, N)
+ keys.foreach { k =>
+ map.append(k, unsafeProj(InternalRow(k)))
+ }
+ val before = map.getTotalMemoryConsumption
+ map.optimize()
+ assert(map.getTotalMemoryConsumption === before)
Review Comment:
**Nit (P3):** This assertion, like the one after the second `optimize()` at
line 317, only compares `getTotalMemoryConsumption` (`array.size() * 8 +
page.size()`), which a same-size rehash leaves unchanged. If the `target <
array.size()` guard in `shrinkArray` were relaxed to `<=` or removed, every
right-sized build would allocate a second full-size array and rehash all keys
at `optimize()`, yet both tests and `afterEach`'s `executionMemoryUsed === 0`
check would still pass, because `rehashInto` frees the old array. Peak memory
is not a reliable substitute: the suite's shared `mm` keeps a running peak
across tests, and even a fresh `TaskMemoryManager` would miss a rehash of the
16-long array in the second `optimize()`, which stays far below the peak the
build already reached. Counting allocations, for example with a test
`TaskMemoryManager` subclass that overrides the public `allocatePage`, makes
both no-rehash claims fail when they break.
**Verification:**
- **Regression:** If optimize() allocates during either check, that check
fails in both memory modes; for example, relaxing the shrink guard to allow a
same-size rehash makes the right-sized-map check fail. With the current code
both checks pass.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/joins/HashedRelationSuite.scala:
##########
@@ -272,6 +272,135 @@ abstract class HashedRelationSuite extends
SharedSparkSession {
map2.free()
}
+ test("LongToUnsafeRowMap shrinks an over-allocated sparse array at optimize
time") {
+ val unsafeProj = UnsafeProjection.create(Seq(BoundReference(0, LongType,
false)))
+
+ // Many rows share 3 keys spread far apart, so the initial array is sized
by the row
+ // count (2 * nextPow2(N) longs) while the load factor invariant of
updateIndex only
+ // needs 16 longs for 3 keys. Keys far apart also keep optimize() from
switching to
+ // the dense mode, which is preferred whenever it applies.
+ val keys = Array(-1000000000000L, 0L, 1000000000000L)
+ val N = 99999
+ val map = new LongToUnsafeRowMap(mm, N)
+ var i = 0
+ while (i < N) {
+ val k = keys(i % 3)
+ map.append(k, unsafeProj(InternalRow(k)))
+ i += 1
+ }
+
+ // Capture the page size via the difference, so the assertions do not
depend on the
+ // details of page growth.
+ var n = 1
+ while (n < N) n *= 2
+ val pageBytes = map.getTotalMemoryConsumption - n * 2L * 8
+ map.optimize()
+ // The smallest power-of-two array keeping numKeys * 4 <= array.size() for
3 keys.
+ assert(map.getTotalMemoryConsumption === pageBytes + 16L * 8)
+
+ val row = unsafeProj(InternalRow(0L)).copy()
+ keys.foreach { k =>
+ val r = map.get(k, row)
+ var c = 0
+ while (r.hasNext) {
+ assert(r.next().getLong(0) === k)
+ c += 1
+ }
+ assert(c === N / 3)
+ }
+ assert(map.getValue(12345L, row) eq null)
+ // keys() has production consumers on broadcast relations (e.g. dynamic
partition
+ // pruning), so the shrunk map must still serve it.
+ assert(map.keys().map(_.getLong(0)).toSet === keys.toSet)
+ // A second optimize() must be a no-op: the shrunk array is already
minimal.
+ map.optimize()
+ assert(map.getTotalMemoryConsumption === pageBytes + 16L * 8)
+ map.free()
+ }
+
+ test("LongToUnsafeRowMap keeps a right-sized sparse array unchanged") {
+ val unsafeProj = UnsafeProjection.create(Seq(BoundReference(0, LongType,
false)))
+
+ // With unique keys the growth in updateIndex/growArray already leaves the
array at
+ // the minimal size for the load factor invariant, so optimize() must not
rehash.
+ val N = 100000
+ val keys = (0 until N).map(i => i * 0x9E3779B97F4A7C15L).toArray
+ val map = new LongToUnsafeRowMap(mm, N)
+ keys.foreach { k =>
+ map.append(k, unsafeProj(InternalRow(k)))
+ }
+ val before = map.getTotalMemoryConsumption
+ map.optimize()
+ assert(map.getTotalMemoryConsumption === before)
+
+ val row = unsafeProj(InternalRow(0L)).copy()
+ keys.foreach { k =>
+ assert(map.getValue(k, row).getLong(0) === k)
+ }
+ map.free()
+ }
+
+ test("LongToUnsafeRowMap serialization carries the shrunk sparse array") {
+ val unsafeProj = UnsafeProjection.create(Seq(BoundReference(0, LongType,
false)))
+
+ val keys = Array(-1000000000000L, 0L, 1000000000000L)
+ val N = 999
+ val map = new LongToUnsafeRowMap(mm, N)
+ var i = 0
+ while (i < N) {
+ val k = keys(i % 3)
+ map.append(k, unsafeProj(InternalRow(k)))
+ i += 1
+ }
+ map.optimize()
+
+ // write() ships the sparse array verbatim: with the shrunk 16-long index
the payload
+ // is the index plus N rows (24 bytes each) and the stream framing, well
under the
+ // un-shrunk index alone (2 * nextPow2(N) * 8 = 16384 bytes).
+ val os = new ByteArrayOutputStream()
+ val out = new ObjectOutputStream(os)
+ map.writeExternal(out)
+ out.flush()
+ assert(os.toByteArray.length < 16L * 8 + N * 24 + 1024)
+ map.free()
+
+ val in = new ObjectInputStream(new ByteArrayInputStream(os.toByteArray))
+ val map2 = new LongToUnsafeRowMap(mm, 1)
+ map2.readExternal(in)
+ val row = unsafeProj(InternalRow(0L)).copy()
+ keys.foreach { k =>
+ val r = map2.get(k, row)
+ var c = 0
+ while (r.hasNext) {
+ assert(r.next().getLong(0) === k)
+ c += 1
+ }
+ assert(c === N / 3)
+ }
+ map2.free()
+ }
+
+ test("LongToUnsafeRowMap collapses an empty map via the existing dense-mode
path") {
+ val unsafeProj = UnsafeProjection.create(Seq(BoundReference(0, LongType,
false)))
+ // No appends: minKey/maxKey stay at their sentinel values, so maxKey -
minKey wraps
+ // around to exactly 1, which is always < 1024. optimize() therefore
always takes the
+ // pre-existing dense-mode branch for an empty map, never shrinkArray()
(numKeys == 0
+ // is unreachable there: any append() that is not a duplicate also updates
+ // minKey/maxKey, so a non-empty map never leaves them at their sentinel
values).
+ // This test is a regression guard on that pre-existing behavior, not new
coverage of
+ // the shrink added by this PR.
Review Comment:
**Nit (P3):** Once this is merged, 'existing', 'pre-existing' and 'this PR',
here and in the test name on line 383, no longer refer to any version of the
code. The comment can state the lasting property directly: with no appends,
`maxKey - minKey` wraps around to exactly 1, so `optimize()` always takes the
dense branch for an empty map and never reaches `shrinkArray()`. A name such as
"LongToUnsafeRowMap optimize() turns an empty map into dense mode" describes
the same test without the history.
**Verification:**
- **Inspection:** The empty-map test's name and comment contain no PR- or
history-relative wording and explain why optimize() takes the dense branch.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/joins/HashedRelationSuite.scala:
##########
@@ -272,6 +272,135 @@ abstract class HashedRelationSuite extends
SharedSparkSession {
map2.free()
}
+ test("LongToUnsafeRowMap shrinks an over-allocated sparse array at optimize
time") {
+ val unsafeProj = UnsafeProjection.create(Seq(BoundReference(0, LongType,
false)))
+
+ // Many rows share 3 keys spread far apart, so the initial array is sized
by the row
+ // count (2 * nextPow2(N) longs) while the load factor invariant of
updateIndex only
+ // needs 16 longs for 3 keys. Keys far apart also keep optimize() from
switching to
+ // the dense mode, which is preferred whenever it applies.
+ val keys = Array(-1000000000000L, 0L, 1000000000000L)
+ val N = 99999
+ val map = new LongToUnsafeRowMap(mm, N)
+ var i = 0
+ while (i < N) {
+ val k = keys(i % 3)
+ map.append(k, unsafeProj(InternalRow(k)))
+ i += 1
+ }
+
+ // Capture the page size via the difference, so the assertions do not
depend on the
+ // details of page growth.
+ var n = 1
+ while (n < N) n *= 2
+ val pageBytes = map.getTotalMemoryConsumption - n * 2L * 8
+ map.optimize()
+ // The smallest power-of-two array keeping numKeys * 4 <= array.size() for
3 keys.
+ assert(map.getTotalMemoryConsumption === pageBytes + 16L * 8)
+
+ val row = unsafeProj(InternalRow(0L)).copy()
+ keys.foreach { k =>
+ val r = map.get(k, row)
+ var c = 0
+ while (r.hasNext) {
+ assert(r.next().getLong(0) === k)
+ c += 1
+ }
+ assert(c === N / 3)
+ }
+ assert(map.getValue(12345L, row) eq null)
+ // keys() has production consumers on broadcast relations (e.g. dynamic
partition
+ // pruning), so the shrunk map must still serve it.
+ assert(map.keys().map(_.getLong(0)).toSet === keys.toSet)
+ // A second optimize() must be a no-op: the shrunk array is already
minimal.
+ map.optimize()
+ assert(map.getTotalMemoryConsumption === pageBytes + 16L * 8)
+ map.free()
+ }
+
+ test("LongToUnsafeRowMap keeps a right-sized sparse array unchanged") {
+ val unsafeProj = UnsafeProjection.create(Seq(BoundReference(0, LongType,
false)))
+
+ // With unique keys the growth in updateIndex/growArray already leaves the
array at
+ // the minimal size for the load factor invariant, so optimize() must not
rehash.
+ val N = 100000
+ val keys = (0 until N).map(i => i * 0x9E3779B97F4A7C15L).toArray
+ val map = new LongToUnsafeRowMap(mm, N)
+ keys.foreach { k =>
+ map.append(k, unsafeProj(InternalRow(k)))
+ }
+ val before = map.getTotalMemoryConsumption
+ map.optimize()
+ assert(map.getTotalMemoryConsumption === before)
+
+ val row = unsafeProj(InternalRow(0L)).copy()
+ keys.foreach { k =>
+ assert(map.getValue(k, row).getLong(0) === k)
+ }
+ map.free()
+ }
+
+ test("LongToUnsafeRowMap serialization carries the shrunk sparse array") {
+ val unsafeProj = UnsafeProjection.create(Seq(BoundReference(0, LongType,
false)))
+
+ val keys = Array(-1000000000000L, 0L, 1000000000000L)
+ val N = 999
+ val map = new LongToUnsafeRowMap(mm, N)
+ var i = 0
+ while (i < N) {
+ val k = keys(i % 3)
+ map.append(k, unsafeProj(InternalRow(k)))
+ i += 1
+ }
+ map.optimize()
+
+ // write() ships the sparse array verbatim: with the shrunk 16-long index
the payload
+ // is the index plus N rows (24 bytes each) and the stream framing, well
under the
Review Comment:
**Nit (P3):** For N = 999 the payload is not under 16,384 bytes: `write()`
emits 1 + 4 * 8 + 8 + 16 * 8 + 8 + 999 * 24 = 24,153 data bytes, and the rows
alone take 23,976. The assertion below still works, but for a different reason:
its bound (16 * 8 + 999 * 24 + 1024 = 25,128) leaves 1,024 bytes for stream
framing over the shrunk payload, while an un-shrunk index would add 16,256
bytes. Saying that in the comment keeps a later reader from tightening the
bound toward 16,384 and breaking the test.
**Verification:**
- **Inspection:** The comment's sizes agree with the serialized layout for
999 rows: 24,153 data bytes, a 25,128-byte bound, and 16,256 extra bytes for an
un-shrunk index.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelation.scala:
##########
@@ -840,27 +840,52 @@ private[execution] final class LongToUnsafeRowMap(
}
}
- private def growArray(): Unit = {
- var old_array = array
- val n = Math.toIntExact(array.size())
+ /**
+ * Rehashes every occupied (key, address) pair of the current array into a
freshly
+ * allocated array of the given size, then frees the old array. Shared by
growArray()
+ * and shrinkArray(), which only differ in how they compute newSize.
+ */
+ private def rehashInto(newSize: Long): Unit = {
+ val oldArray = array
numKeys = 0
- array = allocateArray(n * 2)
+ array = allocateArray(newSize)
array.zeroOut()
- mask = n * 2 - 2
+ mask = Math.toIntExact(newSize - 2)
var i = 0
- while (i < old_array.size()) {
- if (old_array.get(i + 1) > 0) {
- val key = old_array.get(i)
- updateIndex(key, findKeyPosition(key), old_array.get(i + 1))
+ while (i < oldArray.size()) {
+ if (oldArray.get(i + 1) > 0) {
+ val key = oldArray.get(i)
+ updateIndex(key, findKeyPosition(key), oldArray.get(i + 1))
}
i += 2
}
- freeArray(old_array)
- old_array = null // release the reference to old array
+ freeArray(oldArray)
+ }
+
+ private def growArray(): Unit = {
+ rehashInto(array.size() * 2)
+ }
+
+ /**
+ * Shrink the sparse array to the smallest power-of-two size that keeps the
load factor
+ * invariant (numKeys * 4 <= array.size(), see updateIndex). The initial
size comes from
+ * the estimated row count, so a build with many rows sharing few keys ends
with a
+ * mostly-empty array, and a serialized broadcast keeps paying for the
headroom on every
+ * executor. Never grow: a map at the array-size ceiling may hold more keys
than a
+ * minimal table could take (see updateIndex).
+ */
+ private def shrinkArray(): Unit = {
+ var target = 2L
+ while (target < numKeys * 4) target *= 2
+ if (target < array.size()) {
Review Comment:
**Blocking (P1):** The description says shuffled hash joins are unaffected
because their growth ladder already lands on the minimal size, but that only
holds once a task's map has grown. `ShuffledHashJoinExec.buildHashedRelation`
passes no size estimate, so each task's `LongToUnsafeRowMap` starts from
`HashedRelation.apply`'s default of 64 rows, a 128-long array, and
`updateIndex` does not grow it before the 33rd distinct key. A build partition
with 2 to 16 distinct keys spread over a range of at least 1024 skips the dense
branch, and this condition is then true: the map is rehashed into 8 to 64
longs, where the base branch left it alone.
For those tasks (a `SHUFFLE_HASH` hint,
`spark.sql.join.preferSortMergeJoin=false`, or AQE's shuffled-hash selection,
on a single integral or packed long key), each task now pays an extra allocate,
zero, rehash and free, and every streamed lookup probes at a load factor in
(0.25, 0.5] instead of at most 0.25. An exact emulation of
`firstSlot`/`nextSlot` over random key sets gives, for 8 keys, 1.06/1.15 ->
1.33/2.08 slot checks per hit/miss. Both arrays already fit in L1, so nothing
offsets the extra checks, and the memory saved is at most 960 bytes per task,
none of it serialized.
**Recommended change:** Make shrinkArray skip the rehash when the current
array is no larger than the 128-long array that a default-capacity map
allocates, using a named constant tied to the 64-row default estimate, and
document that floor next to the never-grow rationale. Add hashed-relation test
coverage for a default-capacity map with a few wide-range keys.
**Why this works:** Default-capacity maps either keep their 128-long array
or grow to exactly the minimal size, so with the floor the shrink never fires
for them. Maps built from larger capacity estimates, which is how broadcast
builds over-allocate, keep the current shrink.
**Scope:** Add a size floor to the sparse-index shrink policy so
default-capacity maps keep their base-branch behavior, and cover it with a
hashed-relation unit test.
**Compatibility:** Every broadcast build whose index exceeds the floor keeps
the PR's shrink, reduced footprint and smaller serialized payload.
**Risks:** Broadcast builds of at most 64 rows also keep their at most 1 KiB
of headroom. That saving is immaterial, and removing the floor rolls the change
back. The floor encodes the 64-row default estimate; if that default changes,
the floor must change with it, so derive both from one constant.
**Constraints:** Keep the existing never-grow condition (target <
array.size()) and the order in which optimize() tries dense mode before
shrinking.
**Exclusions:**
- How far larger broadcast indexes should shrink, including keeping
headroom, strengthening the slot hash, or requiring a minimum reduction, stays
with the open questions on probe cost and on shrinks that only halve the array.
**Success:** optimize() leaves a default-capacity map with 2 to 16 distinct
keys spanning at least 1024 at 128 longs and allocates no second array.
Over-allocated maps larger than the floor still shrink to the smallest power of
two that keeps numKeys * 4 <= size. Unique-key, grown, dense-eligible, empty
and size-ceiling maps keep their current outcomes.
**Verification:**
- **Regression:** A map built with the default capacity estimate and a few
distinct keys spread over a range of at least 1024 keeps its memory footprint
across optimize() and allocates no second array, in both on-heap and off-heap
modes.
- **Behavior:** An over-allocated map built from a large capacity estimate
with a few wide-range keys still shrinks to the minimal array and still returns
every key's rows, the key set and misses correctly.
- **Compatibility:** Maps that grew during the build, dense-eligible and
empty maps, and serialization round trips keep their existing results.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelation.scala:
##########
@@ -840,27 +840,52 @@ private[execution] final class LongToUnsafeRowMap(
}
}
- private def growArray(): Unit = {
- var old_array = array
- val n = Math.toIntExact(array.size())
+ /**
+ * Rehashes every occupied (key, address) pair of the current array into a
freshly
+ * allocated array of the given size, then frees the old array. Shared by
growArray()
+ * and shrinkArray(), which only differ in how they compute newSize.
+ */
+ private def rehashInto(newSize: Long): Unit = {
+ val oldArray = array
numKeys = 0
- array = allocateArray(n * 2)
+ array = allocateArray(newSize)
array.zeroOut()
- mask = n * 2 - 2
+ mask = Math.toIntExact(newSize - 2)
var i = 0
- while (i < old_array.size()) {
- if (old_array.get(i + 1) > 0) {
- val key = old_array.get(i)
- updateIndex(key, findKeyPosition(key), old_array.get(i + 1))
+ while (i < oldArray.size()) {
+ if (oldArray.get(i + 1) > 0) {
+ val key = oldArray.get(i)
+ updateIndex(key, findKeyPosition(key), oldArray.get(i + 1))
}
i += 2
}
- freeArray(old_array)
- old_array = null // release the reference to old array
+ freeArray(oldArray)
+ }
+
+ private def growArray(): Unit = {
+ rehashInto(array.size() * 2)
+ }
+
+ /**
+ * Shrink the sparse array to the smallest power-of-two size that keeps the
load factor
+ * invariant (numKeys * 4 <= array.size(), see updateIndex). The initial
size comes from
+ * the estimated row count, so a build with many rows sharing few keys ends
with a
+ * mostly-empty array, and a serialized broadcast keeps paying for the
headroom on every
+ * executor. Never grow: a map at the array-size ceiling may hold more keys
than a
+ * minimal table could take (see updateIndex).
Review Comment:
**Nit (P3):** Read against the first sentence's definition (the smallest
power of two that keeps `numKeys * 4 <= size`, which fits every key by
construction), "more keys than a minimal table could take" does not explain the
guard. The actual reason is that at the `1 << 30` ceiling `updateIndex` stops
growing and keeps accepting keys until `numKeys > array.size() / 2 * 0.75`, so
`numKeys * 4` can exceed `array.size()` and the computed target would be larger
than the current array. Something like "Never grow: at the array-size ceiling
updateIndex accepts more keys than the load factor invariant allows, so the
target can exceed array.size()" says that directly.
**Verification:**
- **Inspection:** shrinkArray's doc says that the target can exceed the
current array at the size ceiling and that the guard prevents growing,
consistent with updateIndex.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
