fixup

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 {
 

backends-clickhouse/src/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._
 

backends-clickhouse/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._
 

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}

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 {
 

backends-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

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
 

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

backends-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

backends-velox/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
 
 /**

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;
 

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;
 

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;

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;
 

gluten-celeborn/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

gluten-substrait/src/main/java/org/apache/gluten/memory/memtarget/ThrowOnOomMemoryTarget.java → gluten-core/src/main/java/org/apache/gluten/memory/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 {

gluten-substrait/src/main/java/org/apache/gluten/memory/memtarget/spark/RegularMemoryConsumer.java → gluten-core/src/main/java/org/apache/gluten/memory/memtarget/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;

gluten-substrait/src/main/scala/org/apache/gluten/utils/TaskListener.scala → gluten-core/src/main/java/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
 

gluten-substrait/src/main/scala/org/apache/spark/memory/SparkMemoryUtil.scala → gluten-core/src/main/scala/org/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
 

gluten-substrait/src/main/scala/org/apache/spark/util/TaskResource.scala → gluten-core/src/main/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

gluten-substrait/src/main/scala/org/apache/spark/util/TaskResources.scala → gluten-core/src/main/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
   }
 }

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;
 

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;
 

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;
 

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;

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;

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(

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
 

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 {
 

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
 

gluten-data/src/main/scala/org/apache/gluten/runtime/Runtimes.scala

@@ -16,7 +16,7 @@
  */
 package org.apache.gluten.runtime
 
-import org.apache.spark.util.{TaskResource, TaskResources}
+import org.apache.spark.task.{TaskResource, TaskResources}
 
 object Runtimes {
 

gluten-data/src/main/scala/org/apache/gluten/vectorized/ColumnarBatchSerializer.scala

@@ -30,7 +30,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

gluten-data/src/main/scala/org/apache/spark/sql/execution/ColumnarBuildSideRelation.scala

@@ -30,7 +30,7 @@ import org.apache.spark.sql.execution.joins.BuildSideRelation
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.utils.SparkArrowUtil
 import org.apache.spark.sql.vectorized.ColumnarBatch
-import org.apache.spark.util.TaskResources
+import org.apache.spark.task.TaskResources
 
 import org.apache.arrow.c.ArrowSchema
 

gluten-data/src/test/scala/org/apache/gluten/execution/MassiveMemoryAllocationSuite.scala

@@ -22,7 +22,7 @@ import org.apache.gluten.memory.memtarget.{MemoryTarget, Spiller, Spillers}
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.test.SharedSparkSession
-import org.apache.spark.util.TaskResources
+import org.apache.spark.task.TaskResources
 
 import java.util.concurrent.{Callable, Executors, TimeUnit}
 import java.util.concurrent.atomic.AtomicLong