You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/12 17:50:26 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-5856-da266854dcfb395e4ff4da882c800c4352762fed created (now 5504583774)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-5856-da266854dcfb395e4ff4da882c800c4352762fed
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at 5504583774 test: cover native memory accounting boundaries (#5856)
This branch includes the following new commits:
new 5504583774 test: cover native memory accounting boundaries (#5856)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: test: cover native memory accounting boundaries (#5856)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-5856-da266854dcfb395e4ff4da882c800c4352762fed
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 5504583774748166985e48df2ca2bf0bcf734171
Author: KUAN-HAO HUANG <10...@users.noreply.github.com>
AuthorDate: Sat Sep 12 16:56:35 2026 +0000
test: cover native memory accounting boundaries (#5856)
Co-authored-by: Andy Grove <ag...@apache.org>
---
.../apache/spark/CometTaskMemoryManagerSuite.scala | 99 ++++++++++++++++++----
1 file changed, 82 insertions(+), 17 deletions(-)
diff --git a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala
index 548ecfa488..103e20ed8c 100644
--- a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala
@@ -29,6 +29,82 @@ import org.apache.spark.memory.{MemoryConsumer, TaskMemoryManager, TestMemoryMan
class CometTaskMemoryManagerSuite extends AnyFunSuite {
test("native memory usage is visible to Spark's memory consumer") {
+ withTaskMemoryManager { taskMemoryManager =>
+ val manager = new CometTaskMemoryManager(1L, 0L)
+ val consumer = nativeMemoryConsumer(manager)
+
+ assert(manager.getUsed == 0L)
+ assert(consumer.getUsed == 0L)
+
+ assert(manager.acquireMemory(128L) == 128L)
+ assert(manager.getUsed == 128L)
+ assert(consumer.getUsed == 128L)
+ assert(taskMemoryManager.getMemoryConsumptionForThisTask == 128L)
+
+ manager.releaseMemory(128L)
+ assert(manager.getUsed == 0L)
+ assert(consumer.getUsed == 0L)
+ assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L)
+ }
+ }
+
+ test("partial and zero grants account only for acquired memory") {
+ withTaskMemoryManager { taskMemoryManager =>
+ val manager = new CometTaskMemoryManager(1L, 0L)
+ val consumer = nativeMemoryConsumer(manager)
+
+ def checkUsage(expected: Long): Unit = {
+ assert(manager.getUsed == expected)
+ assert(consumer.getUsed == expected)
+ assert(taskMemoryManager.getMemoryConsumptionForThisTask == expected)
+ }
+
+ assert(manager.acquireMemory(768L) == 768L)
+ checkUsage(768L)
+ assert(manager.acquireMemory(512L) == 256L)
+ checkUsage(1024L)
+ assert(manager.acquireMemory(128L) == 0L)
+ checkUsage(1024L)
+
+ manager.releaseMemory(256L)
+ checkUsage(768L)
+ assert(manager.acquireMemory(128L) == 128L)
+ checkUsage(896L)
+ manager.releaseMemory(896L)
+ checkUsage(0L)
+ }
+ }
+
+ test("managers in the same task retain separate memory accounting") {
+ withTaskMemoryManager { taskMemoryManager =>
+ val first = new CometTaskMemoryManager(1L, 0L)
+ val second = new CometTaskMemoryManager(2L, 0L)
+ val firstConsumer = nativeMemoryConsumer(first)
+ val secondConsumer = nativeMemoryConsumer(second)
+
+ def checkUsage(firstBytes: Long, secondBytes: Long): Unit = {
+ assert(first.getUsed == firstBytes)
+ assert(firstConsumer.getUsed == firstBytes)
+ assert(second.getUsed == secondBytes)
+ assert(secondConsumer.getUsed == secondBytes)
+ assert(taskMemoryManager.getMemoryConsumptionForThisTask == firstBytes + secondBytes)
+ }
+
+ checkUsage(0L, 0L)
+ assert(first.acquireMemory(128L) == 128L)
+ checkUsage(128L, 0L)
+ assert(second.acquireMemory(256L) == 256L)
+ checkUsage(128L, 256L)
+ first.releaseMemory(64L)
+ checkUsage(64L, 256L)
+ second.releaseMemory(256L)
+ checkUsage(64L, 0L)
+ first.releaseMemory(64L)
+ checkUsage(0L, 0L)
+ }
+ }
+
+ private def withTaskMemoryManager(f: TaskMemoryManager => Unit): Unit = {
val memoryManager = new TestMemoryManager(new SparkConf())
memoryManager.limit(1024)
val taskMemoryManager = new TaskMemoryManager(memoryManager, 0L)
@@ -48,24 +124,13 @@ class CometTaskMemoryManagerSuite extends AnyFunSuite {
TaskContext.setTaskContext(taskContext)
try {
- val manager = new CometTaskMemoryManager(1L, 0L)
- val consumer = nativeMemoryConsumer(manager)
-
- assert(manager.getUsed == 0L)
- assert(consumer.getUsed == 0L)
-
- assert(manager.acquireMemory(128L) == 128L)
- assert(manager.getUsed == 128L)
- assert(consumer.getUsed == 128L)
- assert(taskMemoryManager.getMemoryConsumptionForThisTask == 128L)
-
- manager.releaseMemory(128L)
- assert(manager.getUsed == 0L)
- assert(consumer.getUsed == 0L)
- assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L)
+ f(taskMemoryManager)
} finally {
- taskMemoryManager.cleanUpAllAllocatedMemory()
- TaskContext.unset()
+ try {
+ taskMemoryManager.cleanUpAllAllocatedMemory()
+ } finally {
+ TaskContext.unset()
+ }
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org