Call clear after each batch.

Author: viirya

python/pyspark/sql/tests.py

@@ -4680,6 +4680,26 @@ def test_supported_types(self):
         self.assertPandasEqual(expected2, result2)
         self.assertPandasEqual(expected3, result3)
 
+    def test_array_type_correct(self):
+        from pyspark.sql.functions import pandas_udf, PandasUDFType, array, col
+
+        df = self.data.withColumn("arr", array(col("id"))).repartition(1, "id")
+
+        output_schema = StructType(
+            [StructField('id', LongType()),
+             StructField('v', IntegerType()),
+             StructField('arr', ArrayType(LongType()))])
+
+        udf = pandas_udf(
+            lambda pdf: pdf,
+            output_schema,
+            PandasUDFType.GROUPED_MAP
+        )
+
+        result = df.groupby('id').apply(udf).sort('id').toPandas()
+        expected = df.toPandas().groupby('id').apply(udf.func).reset_index(drop=True)
+        self.assertPandasEqual(expected, result)
+
     def test_register_grouped_map_udf(self):
         from pyspark.sql.functions import pandas_udf, PandasUDFType
 

sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowWriter.scala

@@ -311,6 +311,7 @@ private[arrow] class ArrayWriter(
   override def reset(): Unit = {
     super.reset()
     elementWriter.reset()
+    valueVector.clear()
   }
 }