Clean up; Add more tests

python/pyspark/sql/tests.py

@@ -5454,6 +5454,15 @@ def test_retain_group_columns(self):
             expected1 = df.groupby(df.id).agg(sum(df.v))
             self.assertPandasEqual(expected1.toPandas(), result1.toPandas())
 
+    def test_array_type(self):
+        from pyspark.sql.functions import pandas_udf, PandasUDFType
+
+        df = self.data
+
+        array_udf = pandas_udf(lambda x: [1.0, 2.0], 'array<double>', PandasUDFType.GROUPED_AGG)
+        result1 = df.groupby('id').agg(array_udf(df['v']).alias('v2'))
+        self.assertEquals(result1.first()['v2'], [1.0, 2.0])
+
     def test_invalid_args(self):
         from pyspark.sql.functions import mean
 
@@ -5556,9 +5565,6 @@ def test_multiple_udfs(self):
                       .withColumn('max_v', max(df['v']).over(w)) \
                       .withColumn('min_w', min(df['w']).over(w))
 
-        result1.explain(True)
-        expected1.explain(True)
-
         self.assertPandasEqual(expected1.toPandas(), result1.toPandas())
 
     def test_replace_existing(self):
@@ -5646,29 +5652,46 @@ def test_mixed_sql_and_udf(self):
         expected3 = expected1
 
         # Test mixing sql window function and udf
-        result4 = df.withColumn('max_v', max_udf(df['v']).over(w)).withColumn('rank', rank().over(ow))
-        expected4 = df.withColumn('max_v', max(df['v']).over(w)).withColumn('rank', rank().over(ow))
+        result4 = df.withColumn('max_v', max_udf(df['v']).over(w)) \
+                    .withColumn('rank', rank().over(ow))
+        expected4 = df.withColumn('max_v', max(df['v']).over(w)) \
+                      .withColumn('rank', rank().over(ow))
 
         self.assertPandasEqual(expected1.toPandas(), result1.toPandas())
         self.assertPandasEqual(expected2.toPandas(), result2.toPandas())
         self.assertPandasEqual(expected3.toPandas(), result3.toPandas())
         self.assertPandasEqual(expected4.toPandas(), result4.toPandas())
 
+    def test_array_type(self):
+        from pyspark.sql.functions import pandas_udf, PandasUDFType
+
+        df = self.data
+        w = self.unbounded_window
+
+        array_udf = pandas_udf(lambda x: [1.0, 2.0], 'array<double>', PandasUDFType.GROUPED_AGG)
+        result1 = df.withColumn('v2', array_udf(df['v']).over(w))
+        self.assertEquals(result1.first()['v2'], [1.0, 2.0])
+
     def test_invalid_args(self):
         from pyspark.sql.functions import mean, pandas_udf, PandasUDFType
 
         df = self.data
         w = self.unbounded_window
-
+        ow = self.ordered_window
         mean_udf = self.pandas_agg_mean_udf
 
         with QuietTest(self.sc):
             with self.assertRaisesRegexp(
                     AnalysisException,
-                    '.*not supported within a window function'):
+                    '.*does not have any WindowFunction'):
                 foo_udf = pandas_udf(lambda x: x, 'v double', PandasUDFType.GROUPED_MAP)
                 df.withColumn('v2', foo_udf(df['v']).over(w))
 
+        with QuietTest(self.sc):
+            with self.assertRaisesRegexp(
+                    AnalysisException,
+                    'Only unbounded window frame is supported with Python UDFs.'):
+                df.withColumn('mean_v', mean_udf(df['v']).over(ow))
 
 
 if __name__ == "__main__":

python/pyspark/worker.py

@@ -133,10 +133,8 @@ def wrap_window_agg_pandas_udf(f, return_type):
 
     def wrapped(*series):
         import pandas as pd
-        import numpy as np
         result = f(*series)
-        # This doesn't work with non primitive types
-        return pd.Series(np.repeat(result, len(series[0])))
+        return pd.Series([result]).repeat(len(series[0]))
 
     return lambda *a: (wrapped(*a), arrow_return_type)
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala

@@ -113,9 +113,14 @@ trait CheckAnalysis extends PredicateHelper {
             failAnalysis("An offset window function can only be evaluated in an ordered " +
               s"row-based window frame with a single offset: $w")
 
+          case w @ WindowExpression(_: PythonUDF,
+          WindowSpecDefinition(_, _, frame: SpecifiedWindowFrame))
+            if !frame.isUnbounded =>
+            failAnalysis(s"Only unbounded window frame is supported with Python UDFs.")
+
           case w @ WindowExpression(e, s) =>
             // Only allow window functions with an aggregate expression or an offset window
-            // function.
+            // function or a Pandas window UDF.
             e match {
               case _: AggregateExpression | _: OffsetWindowFunction | _: AggregateWindowFunction =>
                 w

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/windowExpressions.scala

@@ -309,8 +309,7 @@ object WindowFunctionType {
 
   def functionType(windowExpression: NamedExpression): Option[WindowFunctionType] = {
     windowExpression.collectFirst {
-      case _: WindowFunction => SQL
-      case _: AggregateFunction => SQL
+      case _: WindowFunction | _: AggregateFunction => SQL
       case udf: PythonUDF if PythonUDF.isWindowPandasUDF(udf) => Python
     }
   }

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

@@ -122,21 +122,17 @@ case class WindowInPandasExec(
       }
 
       val inputProj = UnsafeProjection.create(allInputs, child.output)
-      val pythonInput = grouped.map{ case (k, rows) =>
+      val pythonInput = grouped.map { case (k, rows) =>
           rows.map { row =>
             queue.add(row.asInstanceOf[UnsafeRow])
             inputProj(row)
           }
       }
 
-      val pythonEvalType = udfExpressions.head.evalType match {
-        case PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF =>
-          PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF
-      }
-
       val windowFunctionResult = new ArrowPythonRunner(
         pyFuncs, bufferSize, reuseWorker,
-        pythonEvalType, argOffsets, windowInputSchema,
+        PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF,
+        argOffsets, windowInputSchema,
         sessionLocalTimeZone, pandasRespectSessionTimeZone)
         .compute(pythonInput, context.partitionId(), context)