viirya commented on code in PR #60196:
URL: https://github.com/apache/spark/pull/60196#discussion_r4500475750
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -175,6 +175,12 @@ private enum MapState {
*/
private final long pageSizeBytes;
+ /**
+ * The size of the next data page to allocate. It starts at the initial page
size and doubles
+ * with every new page until it reaches {@link #pageSizeBytes}.
+ */
+ private long nextPageSizeBytes;
Review Comment:
`reset()` doesn't restore `nextPageSizeBytes`, so a map created with a small
initial page goes straight back to full-size pages after a reset, while the
javadoc says "Resets this map to initialized state". Only aggregation calls
`reset()` today, and there the two sizes are equal, so nothing is broken now.
Should we keep a final `initialPageSizeBytes` field and restore it here so the
new constructor's contract also holds after a reset?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/joins/HashedRelationSuite.scala:
##########
@@ -141,6 +141,60 @@ abstract class HashedRelationSuite extends
SharedSparkSession {
assert(java.util.Arrays.equals(os2.toByteArray, os.toByteArray))
}
+ test("SPARK-61082: many UnsafeHashedRelations fit in one task when each
holds little data") {
+ // UnsafeHashedRelation takes its page size from SparkEnv (about 64 MiB in
tests), not from this
+ // memory manager. 16 relations rounded up to a full page each would need
about 1 GiB, 4x the
+ // memory available.
+ val smallUmm = newMemoryManager(ByteUnit.MiB.toBytes(64),
ByteUnit.MiB.toBytes(256))
Review Comment:
`UnsafeHashedRelation` takes its page size from `SparkEnv`, so the 64 MiB
passed to `newMemoryManager` has no effect here. Whether this test fails
without the fix then depends on the test JVM's heap: with a 2g heap the derived
page size can be 16 MiB, and 16 relations would roughly fit in 256 MiB anyway.
Could we assert the precondition, e.g.
`assert(SparkEnv.get.memoryManager.pageSizeBytes >= ByteUnit.MiB.toBytes(32))`,
so the test can't silently stop covering the bug?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/joins/HashedRelationSuite.scala:
##########
@@ -141,6 +141,60 @@ abstract class HashedRelationSuite extends
SharedSparkSession {
assert(java.util.Arrays.equals(os2.toByteArray, os.toByteArray))
}
+ test("SPARK-61082: many UnsafeHashedRelations fit in one task when each
holds little data") {
+ // UnsafeHashedRelation takes its page size from SparkEnv (about 64 MiB in
tests), not from this
+ // memory manager. 16 relations rounded up to a full page each would need
about 1 GiB, 4x the
+ // memory available.
+ val smallUmm = newMemoryManager(ByteUnit.MiB.toBytes(64),
ByteUnit.MiB.toBytes(256))
+ val smallMm = new TaskMemoryManager(smallUmm, 0)
+ val buildKey = Seq(BoundReference(0, StringType, false))
+ val toUnsafeRow = UnsafeProjection.create(Array[DataType](StringType,
IntegerType))
+ val toUnsafeKey = UnsafeProjection.create(Array[DataType](StringType))
+ val relations = ArrayBuffer.empty[HashedRelation]
+ try {
+ (0 until 16).foreach { i =>
+ val rows = (0 until 10).map { k =>
+ toUnsafeRow(InternalRow(UTF8String.fromString(s"key$k"), k +
i)).copy()
+ }
+ val relation = HashedRelation(rows.iterator, buildKey,
taskMemoryManager = smallMm)
+ assert(relation.isInstanceOf[UnsafeHashedRelation])
+ relations += relation
+ }
+ relations.zipWithIndex.foreach { case (relation, i) =>
+ assert(relation.estimatedSize < 2 *
UnsafeHashedRelation.INITIAL_PAGE_SIZE_BYTES)
+ (0 until 10).foreach { k =>
+ val key = toUnsafeKey(InternalRow(UTF8String.fromString(s"key$k")))
+ assert(relation.getValue(key).getInt(1) === k + i)
+ }
+ }
+ } catch {
+ // SparkOutOfMemoryError is fatal to ScalaTest and would abort the whole
suite.
+ case e: SparkOutOfMemoryError =>
+ fail("Could not build 16 small UnsafeHashedRelations in one task", e)
+ } finally {
+ relations.foreach(_.close())
+ }
+ assert(smallUmm.executionMemoryUsed === 0)
+ }
+
+ test("SPARK-61082: UnsafeHashedRelation does not round its page up to the
page size") {
+ // The single string column is also the key, so a row is its own key.
+ val buildKey = Seq(BoundReference(0, StringType, false))
+ val toUnsafeRow = UnsafeProjection.create(Array[DataType](StringType))
+ val row = toUnsafeRow(InternalRow(UTF8String.fromString("a"))).copy()
+ val relation = HashedRelation(Iterator(row), buildKey, taskMemoryManager =
mm)
+ val ser = sparkContext.env.serializer.newInstance()
Review Comment:
nit: the SPARK-60998 deserialization test covers both Java and Kryo. Kryo
goes through the same `read` helper here, but checking both serializers would
match that test and protect the Kryo entry point.
##########
core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java:
##########
@@ -817,6 +819,31 @@ public void testPeakMemoryUsed() {
}
}
+ @Test
+ public void dataPagesGrowFromInitialPageSize() {
+ final BytesToBytesMap map = new BytesToBytesMap(taskMemoryManager, 1,
1024, 4096);
+ try {
+ // All records share one key, so the pointer array does not grow and
every increase of the
+ // memory consumption is a new data page.
+ final long[] key = new long[]{0};
+ final List<Long> pageSizes = new ArrayList<>();
+ long previousMemory = map.getTotalMemoryConsumption();
+ while (pageSizes.size() < 5) {
+ assertTrue(map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8)
+ .append(key, Platform.LONG_ARRAY_OFFSET, 8, key,
Platform.LONG_ARRAY_OFFSET, 8));
+ final long memory = map.getTotalMemoryConsumption();
+ if (memory > previousMemory) {
+ pageSizes.add(memory - previousMemory);
+ }
+ previousMemory = memory;
+ }
+ assertEquals(Arrays.asList(1024L, 2048L, 4096L, 4096L, 4096L),
pageSizes);
Review Comment:
The javadoc also says a record larger than the next page gets a page of its
own size. That case is new on this path. Could you add a record bigger than the
1 KiB initial page here and check its page size?
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -259,10 +285,36 @@ public BytesToBytesMap(
throw new IllegalArgumentException("Page size " + pageSizeBytes + "
cannot exceed " +
TaskMemoryManager.MAXIMUM_PAGE_SIZE_BYTES);
}
+ if (initialPageSizeBytes <= 0) {
+ throw new IllegalArgumentException("Initial page size must be greater
than 0");
+ }
+ this.nextPageSizeBytes = Math.min(initialPageSizeBytes, pageSizeBytes);
this.initialCapacity = initialCapacity;
allocate(initialCapacity);
}
+ /**
+ * Creates a map whose first data page has `initialPageSizeBytes` (at most
`pageSizeBytes`), and
+ * whose later pages double in size until they reach `pageSizeBytes`. A map
holding little data
+ * then does not reserve a full page. A record larger than the next page
gets a page of its size.
+ */
+ public BytesToBytesMap(
+ TaskMemoryManager taskMemoryManager,
+ int initialCapacity,
+ long initialPageSizeBytes,
Review Comment:
In the 3-arg overload the third argument is `pageSizeBytes`, but in this one
it is `initialPageSizeBytes`. Since the two `long`s are adjacent, it's easy to
swap them by mistake. Would `(taskMemoryManager, initialCapacity,
pageSizeBytes, initialPageSizeBytes)` be safer?
--
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]
