[NSE-746]Fix memory allocation in row to columnar (oap-project#809)

* fix row to columnar Signed-off-by: Yuan Zhou <yuan.zhou@intel.com> * fix overflow Signed-off-by: Yuan Zhou <yuan.zhou@intel.com>
zhouyuan · Apr 7, 2022 · 156c2ed · 156c2ed
1 parent 89a2727
commit 156c2ed
Showing 1 changed file with 3 additions and 5 deletions.
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
@@ -120,15 +120,13 @@ case class ArrowRowToColumnarExec(child: SparkPlan) extends UnaryExecNode {
           private val converters = new RowToColumnConverter(localSchema)
           private var last_cb: ColumnarBatch = null
           private var elapse: Long = 0
-          // Allocate large buffer to store the numRows rows
-          val bufferSize = 134217728  // 128M can estimator the buffer size based on the data type
           val allocator = SparkMemoryUtils.contextAllocator()
           var arrowBuf: ArrowBuf = null
           override def hasNext: Boolean = {
             rowIterator.hasNext
           }
           TaskContext.get().addTaskCompletionListener[Unit] { _ =>
-            if (arrowBuf != null) {
+            if (arrowBuf != null && arrowBuf.isOpen()) {
               arrowBuf.close()
             }
           }
@@ -144,7 +142,7 @@ case class ArrowRowToColumnarExec(child: SparkPlan) extends UnaryExecNode {
               isUnsafeRow = false
             }
 
-            if (arrowBuf != null && isUnsafeRow) {
+            if (isUnsafeRow) {
               val rowLength = new ListBuffer[Long]()
               var rowCount = 0
               var offset = 0
@@ -154,7 +152,7 @@ case class ArrowRowToColumnarExec(child: SparkPlan) extends UnaryExecNode {
               val unsafeRow = firstRow.asInstanceOf[UnsafeRow]
               val sizeInBytes = unsafeRow.getSizeInBytes
               // allocate buffer based on 1st row
-              val estimatedBufSize = sizeInBytes * numRows * 1.2
+              val estimatedBufSize = sizeInBytes.toDouble * numRows * 1.2
               arrowBuf = allocator.buffer(estimatedBufSize.toLong)
               Platform.copyMemory(unsafeRow.getBaseObject, unsafeRow.getBaseOffset,
                 null, arrowBuf.memoryAddress() + offset, sizeInBytes)