[GLUTEN-7031] Move task lifecycle management / memory consumer facili…

…ties to gluten-core (#7088)
apache · Sep 3, 2024 · 7ca509f · 7ca509f
1 parent e108f7b
commit 7ca509f
Show file tree

Hide file tree

Showing 63 changed files with 85 additions and 57 deletions.
diff --git a/backends-clickhouse/src/main/java/org/apache/gluten/memory/CHThreadGroup.java b/backends-clickhouse/src/main/java/org/apache/gluten/memory/CHThreadGroup.java
@@ -17,8 +17,8 @@
 package org.apache.gluten.memory;
 
 import org.apache.spark.TaskContext;
-import org.apache.spark.util.TaskResource;
-import org.apache.spark.util.TaskResources;
+import org.apache.spark.task.TaskResource;
+import org.apache.spark.task.TaskResources;
 
 public class CHThreadGroup implements TaskResource {
 

diff --git a/...rc/test/scala/org/apache/gluten/execution/metrics/GlutenClickHouseTPCDSMetricsSuite.scala b/...rc/test/scala/org/apache/gluten/execution/metrics/GlutenClickHouseTPCDSMetricsSuite.scala
@@ -23,7 +23,7 @@ import org.apache.gluten.vectorized.GeneralInIterator
 import org.apache.spark.SparkConf
 import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.execution.InputIteratorTransformer
-import org.apache.spark.util.TaskResources
+import org.apache.spark.task.TaskResources
 
 import scala.collection.JavaConverters._
 

diff --git a/...src/test/scala/org/apache/gluten/execution/metrics/GlutenClickHouseTPCHMetricsSuite.scala b/...src/test/scala/org/apache/gluten/execution/metrics/GlutenClickHouseTPCHMetricsSuite.scala
@@ -23,7 +23,7 @@ import org.apache.gluten.vectorized.GeneralInIterator
 import org.apache.spark.SparkConf
 import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.execution.InputIteratorTransformer
-import org.apache.spark.util.TaskResources
+import org.apache.spark.task.TaskResources
 
 import scala.collection.JavaConverters._
 

diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxTransformerApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxTransformerApi.scala
@@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression}
 import org.apache.spark.sql.connector.read.InputPartition
 import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, PartitionDirectory}
 import org.apache.spark.sql.types._
-import org.apache.spark.util.TaskResources
+import org.apache.spark.task.TaskResources
 import org.apache.spark.util.collection.BitSet
 
 import com.google.protobuf.{Any, Message}

diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxValidatorApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxValidatorApi.scala
@@ -25,7 +25,7 @@ import org.apache.spark.sql.catalyst.expressions.Expression
 import org.apache.spark.sql.catalyst.plans.physical.Partitioning
 import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.types._
-import org.apache.spark.util.TaskResources
+import org.apache.spark.task.TaskResources
 
 class VeloxValidatorApi extends ValidatorApi {
 

diff --git a/...velox/src/main/scala/org/apache/gluten/datasource/v2/ArrowCSVPartitionReaderFactory.scala b/...velox/src/main/scala/org/apache/gluten/datasource/v2/ArrowCSVPartitionReaderFactory.scala
@@ -33,7 +33,8 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.sources.Filter
 import org.apache.spark.sql.types.{StructField, StructType}
 import org.apache.spark.sql.vectorized.ColumnarBatch
-import org.apache.spark.util.{SerializableConfiguration, TaskResources}
+import org.apache.spark.task.TaskResources
+import org.apache.spark.util.SerializableConfiguration
 
 import org.apache.arrow.c.ArrowSchema
 import org.apache.arrow.vector.types.pojo.Schema

diff --git a/backends-velox/src/main/scala/org/apache/gluten/datasource/v2/ArrowCSVTable.scala b/backends-velox/src/main/scala/org/apache/gluten/datasource/v2/ArrowCSVTable.scala
@@ -29,7 +29,7 @@ import org.apache.spark.sql.execution.datasources.FileFormat
 import org.apache.spark.sql.execution.datasources.v2.FileTable
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
-import org.apache.spark.util.TaskResources
+import org.apache.spark.task.TaskResources
 
 import org.apache.hadoop.fs.FileStatus
 

diff --git a/backends-velox/src/main/scala/org/apache/gluten/execution/RowToVeloxColumnarExec.scala b/backends-velox/src/main/scala/org/apache/gluten/execution/RowToVeloxColumnarExec.scala
@@ -34,8 +34,8 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.utils.SparkArrowUtil
 import org.apache.spark.sql.vectorized.ColumnarBatch
+import org.apache.spark.task.TaskResources
 import org.apache.spark.unsafe.Platform
-import org.apache.spark.util.TaskResources
 
 import org.apache.arrow.c.ArrowSchema
 import org.apache.arrow.memory.ArrowBuf

diff --git a/...ends-velox/src/main/scala/org/apache/gluten/expression/VeloxBloomFilterMightContain.scala b/...ends-velox/src/main/scala/org/apache/gluten/expression/VeloxBloomFilterMightContain.scala
@@ -25,7 +25,7 @@ import org.apache.spark.sql.catalyst.expressions.{BinaryExpression, Expression}
 import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.catalyst.expressions.codegen.Block.BlockHelper
 import org.apache.spark.sql.types.DataType
-import org.apache.spark.util.TaskResources
+import org.apache.spark.task.TaskResources
 
 /**
  * Velox's bloom-filter implementation uses different algorithms internally comparing to vanilla

diff --git a/...lox/src/main/scala/org/apache/gluten/expression/aggregate/VeloxBloomFilterAggregate.scala b/...lox/src/main/scala/org/apache/gluten/expression/aggregate/VeloxBloomFilterAggregate.scala
@@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.TypedImperativeAggreg
 import org.apache.spark.sql.catalyst.trees.TernaryLike
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.DataType
-import org.apache.spark.util.TaskResources
+import org.apache.spark.task.TaskResources
 import org.apache.spark.util.sketch.BloomFilter
 
 /**

diff --git a/backends-velox/src/main/scala/org/apache/spark/sql/execution/BroadcastUtils.scala b/backends-velox/src/main/scala/org/apache/spark/sql/execution/BroadcastUtils.scala
@@ -29,7 +29,7 @@ import org.apache.spark.sql.catalyst.plans.physical.{BroadcastMode, BroadcastPar
 import org.apache.spark.sql.execution.joins.{HashedRelation, HashedRelationBroadcastMode, LongHashedRelation}
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.vectorized.ColumnarBatch
-import org.apache.spark.util.TaskResources
+import org.apache.spark.task.TaskResources
 
 import scala.collection.mutable.ArrayBuffer;
 

diff --git a/backends-velox/src/test/java/org/apache/gluten/columnarbatch/ColumnarBatchTest.java b/backends-velox/src/test/java/org/apache/gluten/columnarbatch/ColumnarBatchTest.java
@@ -24,7 +24,7 @@
 import org.apache.spark.sql.types.DataTypes;
 import org.apache.spark.sql.types.StructType;
 import org.apache.spark.sql.vectorized.ColumnarBatch;
-import org.apache.spark.util.TaskResources$;
+import org.apache.spark.task.TaskResources$;
 import org.junit.Assert;
 import org.junit.Test;
 

diff --git a/backends-velox/src/test/java/org/apache/gluten/utils/VeloxBloomFilterTest.java b/backends-velox/src/test/java/org/apache/gluten/utils/VeloxBloomFilterTest.java
@@ -18,7 +18,7 @@
 
 import org.apache.gluten.test.VeloxBackendTestBase;
 
-import org.apache.spark.util.TaskResources$;
+import org.apache.spark.task.TaskResources$;
 import org.apache.spark.util.sketch.BloomFilter;
 import org.apache.spark.util.sketch.IncompatibleMergeException;
 import org.junit.Assert;

diff --git a/backends-velox/src/test/java/org/apache/gluten/vectorized/ArrowColumnVectorTest.java b/backends-velox/src/test/java/org/apache/gluten/vectorized/ArrowColumnVectorTest.java
@@ -19,7 +19,7 @@
 import org.apache.spark.sql.execution.vectorized.MutableColumnarRow;
 import org.apache.spark.sql.types.Decimal;
 import org.apache.spark.sql.types.StructType;
-import org.apache.spark.util.TaskResources$;
+import org.apache.spark.task.TaskResources$;
 import org.junit.Assert;
 import org.junit.Test;
 

diff --git a/cpp/core/CMakeLists.txt b/cpp/core/CMakeLists.txt
@@ -38,14 +38,14 @@ endif()
 
 set(CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/CMake" ${CMAKE_MODULE_PATH})
 
+set(GLUTEN_PROTO_SRC_DIR
+    ${GLUTEN_HOME}/gluten-core/src/main/resources/org/apache/gluten/proto)
+message(STATUS "Set Gluten Proto Directory in ${GLUTEN_PROTO_SRC_DIR}")
+
 set(SUBSTRAIT_PROTO_SRC_DIR
     ${GLUTEN_HOME}/gluten-substrait/src/main/resources/substrait/proto)
 message(STATUS "Set Substrait Proto Directory in ${SUBSTRAIT_PROTO_SRC_DIR}")
 
-set(GLUTEN_PROTO_SRC_DIR
-    ${GLUTEN_HOME}/gluten-substrait/src/main/resources/org/apache/gluten/proto)
-message(STATUS "Set Gluten Proto Directory in ${GLUTEN_PROTO_SRC_DIR}")
-
 find_program(CCACHE_FOUND ccache)
 if(CCACHE_FOUND)
   set_property(GLOBAL PROPERTY RULE_LAUNCH_COMPILE ccache)

diff --git a/.../velox/src/main/scala/org/apache/spark/shuffle/VeloxCelebornColumnarBatchSerializer.scala b/.../velox/src/main/scala/org/apache/spark/shuffle/VeloxCelebornColumnarBatchSerializer.scala
@@ -32,7 +32,7 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.utils.SparkSchemaUtil
 import org.apache.spark.sql.vectorized.ColumnarBatch
-import org.apache.spark.util.{TaskResource, TaskResources}
+import org.apache.spark.task.{TaskResource, TaskResources}
 
 import org.apache.arrow.c.ArrowSchema
 import org.apache.arrow.memory.BufferAllocator

diff --git a/gluten-core/pom.xml b/gluten-core/pom.xml
@@ -128,6 +128,28 @@
     <outputDirectory>target/scala-${scala.binary.version}/classes</outputDirectory>
     <testOutputDirectory>target/scala-${scala.binary.version}/test-classes</testOutputDirectory>
     <plugins>
+      <!-- compile proto buffer files using copied protoc binary -->
+      <plugin>
+        <groupId>org.xolstice.maven.plugins</groupId>
+        <artifactId>protobuf-maven-plugin</artifactId>
+        <executions>
+          <execution>
+            <id>compile-gluten-proto</id>
+            <phase>generate-sources</phase>
+            <goals>
+              <goal>compile</goal>
+              <goal>test-compile</goal>
+            </goals>
+            <configuration>
+              <protocArtifact>
+                com.google.protobuf:protoc:${protobuf.version}:exe:${os.detected.classifier}
+              </protocArtifact>
+              <protoSourceRoot>src/main/resources/org/apache/gluten/proto</protoSourceRoot>
+              <clearOutputDirectory>false</clearOutputDirectory>
+            </configuration>
+          </execution>
+        </executions>
+      </plugin>
       <plugin>
         <groupId>org.apache.maven.plugins</groupId>
         <artifactId>maven-resources-plugin</artifactId>

diff --git a/...he/gluten/memory/MemoryUsageRecorder.java → ...he/gluten/memory/MemoryUsageRecorder.java b/...he/gluten/memory/MemoryUsageRecorder.java → ...he/gluten/memory/MemoryUsageRecorder.java
diff --git a/...luten/memory/MemoryUsageStatsBuilder.java → ...luten/memory/MemoryUsageStatsBuilder.java b/...luten/memory/MemoryUsageStatsBuilder.java → ...luten/memory/MemoryUsageStatsBuilder.java
diff --git a/...ten/memory/SimpleMemoryUsageRecorder.java → ...ten/memory/SimpleMemoryUsageRecorder.java b/...ten/memory/SimpleMemoryUsageRecorder.java → ...ten/memory/SimpleMemoryUsageRecorder.java
diff --git a/...get/DynamicOffHeapSizingMemoryTarget.java → ...get/DynamicOffHeapSizingMemoryTarget.java b/...get/DynamicOffHeapSizingMemoryTarget.java → ...get/DynamicOffHeapSizingMemoryTarget.java
diff --git a/...n/memory/memtarget/KnownNameAndStats.java → ...n/memory/memtarget/KnownNameAndStats.java b/...n/memory/memtarget/KnownNameAndStats.java → ...n/memory/memtarget/KnownNameAndStats.java
diff --git a/...memory/memtarget/LoggingMemoryTarget.java → ...memory/memtarget/LoggingMemoryTarget.java b/...memory/memtarget/LoggingMemoryTarget.java → ...memory/memtarget/LoggingMemoryTarget.java
diff --git a/...gluten/memory/memtarget/MemoryTarget.java → ...gluten/memory/memtarget/MemoryTarget.java b/...gluten/memory/memtarget/MemoryTarget.java → ...gluten/memory/memtarget/MemoryTarget.java
diff --git a/...en/memory/memtarget/MemoryTargetUtil.java → ...en/memory/memtarget/MemoryTargetUtil.java b/...en/memory/memtarget/MemoryTargetUtil.java → ...en/memory/memtarget/MemoryTargetUtil.java
diff --git a/...memory/memtarget/MemoryTargetVisitor.java → ...memory/memtarget/MemoryTargetVisitor.java b/...memory/memtarget/MemoryTargetVisitor.java → ...memory/memtarget/MemoryTargetVisitor.java
diff --git a/...luten/memory/memtarget/MemoryTargets.java → ...luten/memory/memtarget/MemoryTargets.java b/...luten/memory/memtarget/MemoryTargets.java → ...luten/memory/memtarget/MemoryTargets.java
diff --git a/...en/memory/memtarget/NoopMemoryTarget.java → ...en/memory/memtarget/NoopMemoryTarget.java b/...en/memory/memtarget/NoopMemoryTarget.java → ...en/memory/memtarget/NoopMemoryTarget.java
diff --git a/.../gluten/memory/memtarget/OverAcquire.java → .../gluten/memory/memtarget/OverAcquire.java b/.../gluten/memory/memtarget/OverAcquire.java → .../gluten/memory/memtarget/OverAcquire.java
diff --git a/...ache/gluten/memory/memtarget/Spiller.java → ...ache/gluten/memory/memtarget/Spiller.java b/...ache/gluten/memory/memtarget/Spiller.java → ...ache/gluten/memory/memtarget/Spiller.java
diff --git a/...che/gluten/memory/memtarget/Spillers.java → ...che/gluten/memory/memtarget/Spillers.java b/...che/gluten/memory/memtarget/Spillers.java → ...che/gluten/memory/memtarget/Spillers.java
diff --git a/...ory/memtarget/ThrowOnOomMemoryTarget.java → ...ory/memtarget/ThrowOnOomMemoryTarget.java b/...ory/memtarget/ThrowOnOomMemoryTarget.java → ...ory/memtarget/ThrowOnOomMemoryTarget.java
@@ -20,7 +20,7 @@
 
 import org.apache.spark.memory.SparkMemoryUtil;
 import org.apache.spark.sql.internal.SQLConf;
-import org.apache.spark.util.TaskResources;
+import org.apache.spark.task.TaskResources;
 import org.apache.spark.util.Utils;
 
 public class ThrowOnOomMemoryTarget implements MemoryTarget {

diff --git a/...en/memory/memtarget/TreeMemoryTarget.java → ...en/memory/memtarget/TreeMemoryTarget.java b/...en/memory/memtarget/TreeMemoryTarget.java → ...en/memory/memtarget/TreeMemoryTarget.java
diff --git a/...n/memory/memtarget/TreeMemoryTargets.java → ...n/memory/memtarget/TreeMemoryTargets.java b/...n/memory/memtarget/TreeMemoryTargets.java → ...n/memory/memtarget/TreeMemoryTargets.java
diff --git a/...emtarget/spark/RegularMemoryConsumer.java → ...emtarget/spark/RegularMemoryConsumer.java b/...emtarget/spark/RegularMemoryConsumer.java → ...emtarget/spark/RegularMemoryConsumer.java
@@ -30,7 +30,7 @@
 import org.apache.spark.memory.MemoryConsumer;
 import org.apache.spark.memory.MemoryMode;
 import org.apache.spark.memory.TaskMemoryManager;
-import org.apache.spark.util.TaskResources;
+import org.apache.spark.task.TaskResources;
 
 import java.util.Map;
 import java.util.stream.Collectors;

diff --git a/...y/memtarget/spark/TreeMemoryConsumer.java → ...y/memtarget/spark/TreeMemoryConsumer.java b/...y/memtarget/spark/TreeMemoryConsumer.java → ...y/memtarget/spark/TreeMemoryConsumer.java
diff --git a/.../memtarget/spark/TreeMemoryConsumers.java → .../memtarget/spark/TreeMemoryConsumers.java b/.../memtarget/spark/TreeMemoryConsumers.java → .../memtarget/spark/TreeMemoryConsumers.java
diff --git a/...rg/apache/gluten/utils/TaskListener.scala → ...org/apache/gluten/task/TaskListener.scala b/...rg/apache/gluten/utils/TaskListener.scala → ...org/apache/gluten/task/TaskListener.scala
@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.gluten.utils
+package org.apache.gluten.task
 
 import org.apache.spark.TaskFailedReason
 

diff --git a/...rces/org/apache/gluten/proto/config.proto → ...rces/org/apache/gluten/proto/config.proto b/...rces/org/apache/gluten/proto/config.proto → ...rces/org/apache/gluten/proto/config.proto
diff --git a/...rces/org/apache/gluten/proto/memory.proto → ...rces/org/apache/gluten/proto/memory.proto b/...rces/org/apache/gluten/proto/memory.proto → ...rces/org/apache/gluten/proto/memory.proto
diff --git a/...apache/spark/memory/SparkMemoryUtil.scala → ...apache/spark/memory/SparkMemoryUtil.scala b/...apache/spark/memory/SparkMemoryUtil.scala → ...apache/spark/memory/SparkMemoryUtil.scala
@@ -16,7 +16,7 @@
  */
 package org.apache.spark.memory
 
-import org.apache.gluten.memory.memtarget.{DynamicOffHeapSizingMemoryTarget, KnownNameAndStats, LoggingMemoryTarget, MemoryTarget, MemoryTargetVisitor, NoopMemoryTarget, OverAcquire, ThrowOnOomMemoryTarget, TreeMemoryTargets}
+import org.apache.gluten.memory.memtarget._
 import org.apache.gluten.memory.memtarget.spark.{RegularMemoryConsumer, TreeMemoryConsumer}
 import org.apache.gluten.proto.MemoryUsageStats
 

diff --git a/.../org/apache/spark/util/TaskResource.scala → .../org/apache/spark/task/TaskResource.scala b/.../org/apache/spark/util/TaskResource.scala → .../org/apache/spark/task/TaskResource.scala
@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.spark.util
+package org.apache.spark.task
 
 /**
  * Manages the lifecycle for a specific type of memory resource managed by Spark. See also

diff --git a/...org/apache/spark/util/TaskResources.scala → ...org/apache/spark/task/TaskResources.scala b/...org/apache/spark/util/TaskResources.scala → ...org/apache/spark/task/TaskResources.scala
@@ -14,15 +14,17 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.spark.util
+package org.apache.spark.task
+
+import org.apache.gluten.task.TaskListener
 
 import org.apache.spark.{TaskContext, TaskFailedReason, TaskKilledException, UnknownReason}
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.util.{TaskCompletionListener, TaskFailureListener}
 
 import _root_.org.apache.gluten.memory.SimpleMemoryUsageRecorder
 import _root_.org.apache.gluten.sql.shims.SparkShimLoader
-import _root_.org.apache.gluten.utils.TaskListener
 
 import java.util
 import java.util.{Collections, Properties, UUID}
@@ -288,7 +290,7 @@ class TaskResourceRegistry extends Logging {
   }
 
   /** Release all managed resources according to priority and reversed order */
-  private[util] def releaseAll(): Unit = lock {
+  private[task] def releaseAll(): Unit = lock {
     val table = new util.ArrayList(priorityToResourcesMapping.entrySet())
     Collections.sort(
       table,
@@ -310,7 +312,7 @@ class TaskResourceRegistry extends Logging {
   }
 
   /** Release single resource by ID */
-  private[util] def releaseResource(id: String): Unit = lock {
+  private[task] def releaseResource(id: String): Unit = lock {
     if (!resources.containsKey(id)) {
       throw new IllegalArgumentException(
         String.format("TaskResource with ID %s is not registered", id))
@@ -328,7 +330,7 @@ class TaskResourceRegistry extends Logging {
     resources.remove(id)
   }
 
-  private[util] def addResourceIfNotRegistered[T <: TaskResource](id: String, factory: () => T): T =
+  private[task] def addResourceIfNotRegistered[T <: TaskResource](id: String, factory: () => T): T =
     lock {
       if (resources.containsKey(id)) {
         return resources.get(id).asInstanceOf[T]
@@ -338,7 +340,7 @@ class TaskResourceRegistry extends Logging {
       resource
     }
 
-  private[util] def addResource[T <: TaskResource](id: String, resource: T): T = lock {
+  private[task] def addResource[T <: TaskResource](id: String, resource: T): T = lock {
     if (resources.containsKey(id)) {
       throw new IllegalArgumentException(
         String.format("TaskResource with ID %s is already registered", id))
@@ -347,19 +349,19 @@ class TaskResourceRegistry extends Logging {
     resource
   }
 
-  private[util] def isResourceRegistered(id: String): Boolean = lock {
+  private[task] def isResourceRegistered(id: String): Boolean = lock {
     resources.containsKey(id)
   }
 
-  private[util] def getResource[T <: TaskResource](id: String): T = lock {
+  private[task] def getResource[T <: TaskResource](id: String): T = lock {
     if (!resources.containsKey(id)) {
       throw new IllegalArgumentException(
         String.format("TaskResource with ID %s is not registered", id))
     }
     resources.get(id).asInstanceOf[T]
   }
 
-  private[util] def getSharedUsage(): SimpleMemoryUsageRecorder = lock {
+  private[task] def getSharedUsage(): SimpleMemoryUsageRecorder = lock {
     sharedUsage
   }
 }
diff --git a/gluten-data/src/main/java/org/apache/gluten/columnarbatch/IndicatorVector.java b/gluten-data/src/main/java/org/apache/gluten/columnarbatch/IndicatorVector.java
@@ -16,7 +16,7 @@
  */
 package org.apache.gluten.columnarbatch;
 
-import org.apache.spark.util.TaskResources;
+import org.apache.spark.task.TaskResources;
 
 import java.util.concurrent.atomic.AtomicLong;
 

diff --git a/gluten-data/src/main/java/org/apache/gluten/columnarbatch/IndicatorVectorPool.java b/gluten-data/src/main/java/org/apache/gluten/columnarbatch/IndicatorVectorPool.java
@@ -16,7 +16,7 @@
  */
 package org.apache.gluten.columnarbatch;
 
-import org.apache.spark.util.TaskResource;
+import org.apache.spark.task.TaskResource;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 

diff --git a/gluten-data/src/main/java/org/apache/gluten/memory/arrow/alloc/ArrowBufferAllocators.java b/gluten-data/src/main/java/org/apache/gluten/memory/arrow/alloc/ArrowBufferAllocators.java
@@ -23,8 +23,8 @@
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.spark.memory.TaskMemoryManager;
-import org.apache.spark.util.TaskResource;
-import org.apache.spark.util.TaskResources;
+import org.apache.spark.task.TaskResource;
+import org.apache.spark.task.TaskResources;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 

diff --git a/gluten-data/src/main/java/org/apache/gluten/memory/arrow/pool/ArrowNativeMemoryPool.java b/gluten-data/src/main/java/org/apache/gluten/memory/arrow/pool/ArrowNativeMemoryPool.java
@@ -17,8 +17,8 @@
 package org.apache.gluten.memory.arrow.pool;
 
 import org.apache.arrow.dataset.jni.NativeMemoryPool;
-import org.apache.spark.util.TaskResource;
-import org.apache.spark.util.TaskResources;
+import org.apache.spark.task.TaskResource;
+import org.apache.spark.task.TaskResources;
 import org.apache.spark.util.Utils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

diff --git a/gluten-data/src/main/java/org/apache/gluten/memory/listener/ReservationListeners.java b/gluten-data/src/main/java/org/apache/gluten/memory/listener/ReservationListeners.java
@@ -22,7 +22,7 @@
 import org.apache.gluten.memory.memtarget.*;
 
 import org.apache.spark.memory.TaskMemoryManager;
-import org.apache.spark.util.TaskResources;
+import org.apache.spark.task.TaskResources;
 
 import java.util.Collections;
 import java.util.Map;

diff --git a/gluten-data/src/main/scala/org/apache/gluten/metrics/HashAggregateMetricsUpdater.scala b/gluten-data/src/main/scala/org/apache/gluten/metrics/HashAggregateMetricsUpdater.scala
@@ -20,7 +20,7 @@ import org.apache.gluten.substrait.AggregationParams
 
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.utils.SparkMetricsUtil
-import org.apache.spark.util.TaskResources
+import org.apache.spark.task.TaskResources
 
 trait HashAggregateMetricsUpdater extends MetricsUpdater {
   def updateAggregationMetrics(

diff --git a/gluten-data/src/main/scala/org/apache/gluten/metrics/JoinMetricsUpdater.scala b/gluten-data/src/main/scala/org/apache/gluten/metrics/JoinMetricsUpdater.scala
@@ -21,7 +21,7 @@ import org.apache.gluten.substrait.JoinParams
 
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.utils.SparkMetricsUtil
-import org.apache.spark.util.TaskResources
+import org.apache.spark.task.TaskResources
 
 import java.util
 

diff --git a/gluten-data/src/main/scala/org/apache/gluten/metrics/SortMetricsUpdater.scala b/gluten-data/src/main/scala/org/apache/gluten/metrics/SortMetricsUpdater.scala
@@ -18,7 +18,7 @@ package org.apache.gluten.metrics
 
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.utils.SparkMetricsUtil
-import org.apache.spark.util.TaskResources
+import org.apache.spark.task.TaskResources
 
 class SortMetricsUpdater(val metrics: Map[String, SQLMetric]) extends MetricsUpdater {
 

diff --git a/gluten-data/src/main/scala/org/apache/gluten/runtime/Runtime.scala b/gluten-data/src/main/scala/org/apache/gluten/runtime/Runtime.scala
@@ -27,7 +27,7 @@ import org.apache.gluten.proto.MemoryUsageStats
 
 import org.apache.spark.memory.SparkMemoryUtil
 import org.apache.spark.sql.internal.{GlutenConfigUtil, SQLConf}
-import org.apache.spark.util.TaskResource
+import org.apache.spark.task.TaskResource
 
 import org.slf4j.LoggerFactory