fix arrow eval python

apache · Aug 14, 2024 · 3b1030b · 3b1030b
1 parent 2012a33
commit 3b1030b
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/backends-velox/src/main/scala/org/apache/spark/api/python/ColumnarArrowEvalPythonExec.scala b/backends-velox/src/main/scala/org/apache/spark/api/python/ColumnarArrowEvalPythonExec.scala
@@ -334,17 +334,16 @@ case class ColumnarArrowEvalPythonExec(
         val inputBatchIter = contextAwareIterator.map {
           inputCb =>
             start_time = System.nanoTime()
-            ColumnarBatches.ensureLoaded(ArrowBufferAllocators.contextInstance, inputCb)
-            ColumnarBatches.retain(inputCb)
+            val loadedCb = ColumnarBatches.ensureLoaded(ArrowBufferAllocators.contextInstance, inputCb)
             // 0. cache input for later merge
-            inputCbCache += inputCb
-            numInputRows += inputCb.numRows
+            inputCbCache += loadedCb
+            numInputRows += loadedCb.numRows
             // We only need to pass the referred cols data to python worker for evaluation.
             var colsForEval = new ArrayBuffer[ColumnVector]()
             for (i <- originalOffsets) {
-              colsForEval += inputCb.column(i)
+              colsForEval += loadedCb.column(i)
             }
-            new ColumnarBatch(colsForEval.toArray, inputCb.numRows())
+            new ColumnarBatch(colsForEval.toArray, loadedCb.numRows())
         }
 
         val outputColumnarBatchIterator =