fix leakage in rowtocolumn (#683)

Signed-off-by: Yuan Zhou <[email protected]>
oap-project · Jan 9, 2022 · 878e660 · 878e660
1 parent 8e9f4fe
commit 878e660
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/native-sql-engine/core/src/main/scala/com/intel/oap/execution/ArrowRowToColumnarExec.scala b/native-sql-engine/core/src/main/scala/com/intel/oap/execution/ArrowRowToColumnarExec.scala
@@ -38,6 +38,7 @@ import org.apache.spark.sql.execution.vectorized.WritableColumnVector
 import org.apache.spark.sql.types._
 import org.apache.spark.sql.util.ArrowUtils
 import org.apache.spark.sql.vectorized.{ColumnVector, ColumnarBatch}
+import org.apache.spark.TaskContext
 import org.apache.spark.unsafe.Platform
 
 
@@ -104,7 +105,9 @@ class ArrowRowToColumnarExec(child: SparkPlan) extends RowToColumnarExec(child =
           override def hasNext: Boolean = {
             rowIterator.hasNext
           }
-
+          TaskContext.get().addTaskCompletionListener[Unit] { _ =>
+            arrowBuf.close()
+          }
           override def next(): ColumnarBatch = {
             var isUnsafeRow = true
             var firstRow = InternalRow.apply()