feat: Reset data buf of NativeBatchDecoderIterator on close (#2235)

* Reset data buf of NativeBatchDecoderIterator on close

* address comment

* address comment

Author: wForget

spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala

@@ -48,7 +48,7 @@ case class NativeBatchDecoderIterator(
   private var currentBatch: ColumnarBatch = null
   private var batch = fetchNext()
 
-  import NativeBatchDecoderIterator.threadLocalDataBuf
+  import NativeBatchDecoderIterator._
 
   if (taskContext != null) {
     taskContext.addTaskCompletionListener[Unit](_ => {
@@ -182,14 +182,24 @@ case class NativeBatchDecoderIterator(
           currentBatch = null
         }
         in.close()
+        resetDataBuf()
         isClosed = true
       }
     }
   }
 }
 
 object NativeBatchDecoderIterator {
+
+  private val INITIAL_BUFFER_SIZE = 128 * 1024
+
   private val threadLocalDataBuf: ThreadLocal[ByteBuffer] = ThreadLocal.withInitial(() => {
-    ByteBuffer.allocateDirect(128 * 1024)
+    ByteBuffer.allocateDirect(INITIAL_BUFFER_SIZE)
   })
+
+  private def resetDataBuf(): Unit = {
+    if (threadLocalDataBuf.get().capacity() > INITIAL_BUFFER_SIZE) {
+      threadLocalDataBuf.set(ByteBuffer.allocateDirect(INITIAL_BUFFER_SIZE))
+    }
+  }
 }