[SPARK-24334] Fix race condition in ArrowPythonRunner causes unclean shutdown of Arrow memory allocator

python/pyspark/sql/tests.py

@@ -4931,6 +4931,30 @@ def foo3(key, pdf):
         expected4 = udf3.func((), pdf)
         self.assertPandasEqual(expected4, result4)
 
+    # Regression test for SPARK-24334
+    def test_memory_leak(self):
+        from pyspark.sql.functions import pandas_udf, col, PandasUDFType, array, lit, explode
+
+        # Have all data in a single executor thread so it can trigger the race condition easier
+        with self.sql_conf({'spark.sql.shuffle.partitions': 1}):
+            df = self.spark.range(0, 1000)
+            df = df.withColumn('id', array([lit(i) for i in range(0, 300)])) \
+                   .withColumn('id', explode(col('id'))) \
+                   .withColumn('v',  array([lit(i) for i in range(0, 1000)]))
+
+            @pandas_udf(df.schema, PandasUDFType.GROUPED_MAP)
+            def foo(pdf):
+                # Throw exception in the UDF
+                xxx
+                return pdf
+
+            result = df.groupby('id').apply(foo)
+
+            with QuietTest(self.sc):
+                with self.assertRaises(py4j.protocol.Py4JJavaError) as context:
+                    result.count()
+                self.assertTrue('Memory leaked' not in str(context.exception))
+
 
 @unittest.skipIf(
     not _have_pandas or not _have_pyarrow,

sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala

@@ -70,19 +70,13 @@ class ArrowPythonRunner(
         val arrowSchema = ArrowUtils.toArrowSchema(schema, timeZoneId)
         val allocator = ArrowUtils.rootAllocator.newChildAllocator(
           s"stdout writer for $pythonExec", 0, Long.MaxValue)
-
         val root = VectorSchemaRoot.create(arrowSchema, allocator)
-        val arrowWriter = ArrowWriter.create(root)
-
-        context.addTaskCompletionListener { _ =>
-          root.close()
-          allocator.close()
-        }
-
-        val writer = new ArrowStreamWriter(root, null, dataOut)
-        writer.start()
 
         Utils.tryWithSafeFinally {
+          val arrowWriter = ArrowWriter.create(root)
+          val writer = new ArrowStreamWriter(root, null, dataOut)
+          writer.start()
+
           while (inputIterator.hasNext) {
             val nextBatch = inputIterator.next()
 
@@ -94,8 +88,18 @@ class ArrowPythonRunner(
             writer.writeBatch()
             arrowWriter.reset()
           }
-        } {
           writer.end()
+        } {
+          // If we close root and allocator in TaskCompletionListener, there could be a race
+          // condition where the writer thread keeps writing to the VectorSchemaRoot while
+          // it's being closed by the TaskCompletion listener.
+          // Closing root and allocator here is cleaner because root and allocator is owned
+          // by the writer thread and is only visible to the writer thread.
+          //
+          // If the writer thread is interrupted by TaskCompletionListener, it should either
+          // (1) in the try block, in which case it will get an InterruptedException when
+          // performing io, and goes into the finally block or (2) in the finally block,
+          // in which case it will ignore the interruption and close the resources.
           root.close()
           allocator.close()
         }