[SPARK-34726][SQL] Fix collectToPython timeouts

apache · Mar 12, 2021 · b4b425e · b4b425e
1 parent 7e0fbe0
commit b4b425e
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -3257,12 +3257,11 @@ class Dataset[T] private[sql](
 
   private[sql] def collectToPython(): Array[Any] = {
     EvaluatePython.registerPicklers()
-    withAction("collectToPython", queryExecution) { plan =>
+    val iter = withAction("collectToPython", queryExecution) { plan =>
       val toJava: (Any) => Any = EvaluatePython.toJava(_, schema)
-      val iter: Iterator[Array[Byte]] = new SerDeUtil.AutoBatchedPickler(
-        plan.executeCollect().iterator.map(toJava))
-      PythonRDD.serveIterator(iter, "serve-DataFrame")
+      new SerDeUtil.AutoBatchedPickler(plan.executeCollect().iterator.map(toJava))
     }
+    PythonRDD.serveIterator(iter, "serve-DataFrame")
   }
 
   private[sql] def getRowsToPython(