diff --git a/gluten-core/src/main/java/org/apache/gluten/memory/memtarget/RetryOnOomMemoryTarget.java b/gluten-core/src/main/java/org/apache/gluten/memory/memtarget/RetryOnOomMemoryTarget.java
index 780fc91ff30c..baa5b7324dcd 100644
--- a/gluten-core/src/main/java/org/apache/gluten/memory/memtarget/RetryOnOomMemoryTarget.java
+++ b/gluten-core/src/main/java/org/apache/gluten/memory/memtarget/RetryOnOomMemoryTarget.java
@@ -17,6 +17,7 @@
 package org.apache.gluten.memory.memtarget;
 
 import org.apache.gluten.memory.MemoryUsageStatsBuilder;
+import org.apache.gluten.proto.MemoryUsageStats;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -98,4 +99,8 @@ public TreeMemoryTarget parent() {
   public Spiller getNodeSpiller() {
     return target.getNodeSpiller();
   }
+
+  public TreeMemoryTarget target() {
+    return target;
+  }
 }