chore: Register task completion listener to ensure CometExecIterator is always closed (#3959)

Author: wForget

spark/src/main/scala/org/apache/comet/CometExecIterator.scala

@@ -135,6 +135,12 @@ class CometExecIterator(
   private var currentBatch: ColumnarBatch = null
   private var closed: Boolean = false
 
+  // Register a task completion listener to ensure native resources are released
+  // when the task is done.
+  TaskContext.get().addTaskCompletionListener[Unit] { _ =>
+    this.close()
+  }
+
   private def getNextBatch: Option[ColumnarBatch] = {
     assert(partitionIndex >= 0 && partitionIndex < numParts)
 

spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala

@@ -137,7 +137,6 @@ private[spark] class CometExecRDD(
 
     Option(context).foreach { ctx =>
       ctx.addTaskCompletionListener[Unit] { _ =>
-        it.close()
         subqueries.foreach(sub => CometScalarSubquery.removeSubquery(it.id, sub))
       }
     }

spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala

@@ -174,7 +174,6 @@ case class CometTakeOrderedAndProjectExec(
 
         Option(TaskContext.get()).foreach { context =>
           context.addTaskCompletionListener[Unit] { _ =>
-            it.close()
             cleanSubqueries(it.id, this)
           }
         }