Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[SPARK-50686][SQL] Hash to sort aggregation fallback - memory usage optimization #49314

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
+24 −5
Open

[SPARK-50686][SQL] Hash to sort aggregation fallback - memory usage optimization #49314

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e6b8e8e
[SPARK-50686][SQL] memory optimization in the hash to sort aggregatio…
akupchinskiy Dec 26, 2024
7f8db0c
Merge branch 'apache:master' into aggregation-memory-usage-optimization
akupchinskiy Dec 27, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 12 additions & 2 deletions core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -964,8 +964,7 @@ public void reset() {
updatePeakMemoryUsed();
numKeys = 0;
numValues = 0;
freeArray(longArray);
longArray = null;
freeInternalArray();
while (dataPages.size() > 0) {
MemoryBlock dataPage = dataPages.removeLast();
freePage(dataPage);
Expand All @@ -976,6 +975,17 @@ public void reset() {
pageCursor = 0;
}

/**
* Free array memory to reduce the memory footprint in case of a fallback
* from a hash-based aggregation to the sort-based one.
*/
public void freeInternalArray() {
if (longArray != null) {
freeArray(longArray);
longArray = null;
}
}

/**
* Grows the size of the hash table and re-hash everything.
*/
Expand Down
4 changes: 3 additions & 1 deletion sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,9 @@ public UnsafeKVExternalSorter(
// so that we can always reuse the pointer array.
if (map.numValues() > pointerArray.size() / 4) {
// Here we ask the map to allocate memory, so that the memory manager won't ask the map
// to spill, if the memory is not enough.
// to spill, if the memory is not enough. Also, we free the redundant internal map array
// to reduce the overall memory footprint.
map.freeInternalArray();
pointerArray = map.allocateArray(map.numValues() * 4L);
}

Expand Down
11 changes: 9 additions & 2 deletions sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeKVExternalSorterSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -242,7 +242,7 @@ class UnsafeKVExternalSorterSuite extends SparkFunSuite with SharedSparkSession
try {
val context = new TaskContextImpl(0, 0, 0, 0, 0, 1, taskMemoryManager, new Properties(), null)
TaskContext.setTaskContext(context)
val expectedSpillSize = map.getTotalMemoryConsumption
val expectedSpillSize = expectedSpillSizeForMapWithDuplicateKeys(map)
val sorter = new UnsafeKVExternalSorter(
schema,
schema,
Expand All @@ -267,7 +267,8 @@ class UnsafeKVExternalSorterSuite extends SparkFunSuite with SharedSparkSession
try {
val context = new TaskContextImpl(0, 0, 0, 0, 0, 1, taskMemoryManager, new Properties(), null)
TaskContext.setTaskContext(context)
val expectedSpillSize = map1.getTotalMemoryConsumption + map2.getTotalMemoryConsumption
val expectedSpillSize = expectedSpillSizeForMapWithDuplicateKeys(map1) +
expectedSpillSizeForMapWithDuplicateKeys(map2)
val sorter1 = new UnsafeKVExternalSorter(
schema,
schema,
Expand Down Expand Up @@ -309,4 +310,10 @@ class UnsafeKVExternalSorterSuite extends SparkFunSuite with SharedSparkSession
}
map
}

private def expectedSpillSizeForMapWithDuplicateKeys(map: BytesToBytesMap): Long = {
val internalArrayMemoryUsed = Option(map.getArray).map(_.memoryBlock().size()).getOrElse(0L)
map.getTotalMemoryConsumption - internalArrayMemoryUsed
}

}
Loading