Add support for dtypes as returnType

Author: icexelloss

python/pyspark/sql/functions.py

@@ -28,7 +28,7 @@
 from pyspark import since, SparkContext
 from pyspark.rdd import _prepare_for_python_RDD, ignore_unicode_prefix
 from pyspark.serializers import PickleSerializer, AutoBatchedSerializer
-from pyspark.sql.types import StringType, DataType, _parse_datatype_string
+from pyspark.sql.types import StringType, DataType, _parse_datatype_string, from_pandas_dtypes
 from pyspark.sql.column import Column, _to_java_column, _to_seq
 from pyspark.sql.dataframe import DataFrame
 
@@ -2207,6 +2207,10 @@ def pandas_udf(f=None, returnType=StringType()):
     |         8|      JOHN DOE|          22|
     +----------+--------------+------------+
     """
+    import pandas as pd
+    if isinstance(returnType, pd.Series):
+        returnType = from_pandas_dtypes(returnType)
+
     return _create_udf(f, returnType=returnType, vectorized=True)
 
 

python/pyspark/sql/tests.py

@@ -3395,9 +3395,16 @@ def assertFramesEqual(self, expected, result):
                ("\n\nResult:\n%s\n%s" % (result, result.dtypes)))
         self.assertTrue(expected.equals(result), msg=msg)
 
-    def test_groupby_apply(self):
+    @property
+    def data(self):
         from pyspark.sql.functions import pandas_udf, array, explode, col, lit
-        df = self.spark.range(10).toDF('id').withColumn("vs", array([lit(i) for i in range(20, 30)])).withColumn("v", explode(col('vs'))).drop('vs')
+        return self.spark.range(10).toDF('id') \
+            .withColumn("vs", array([lit(i) for i in range(20, 30)])) \
+            .withColumn("v", explode(col('vs'))).drop('vs')
+
+    def test_groupby_apply_simple(self):
+        from pyspark.sql.functions import pandas_udf
+        df = self.data
 
         def foo(df):
             ret = df
@@ -3417,6 +3424,26 @@ def foo(df):
         expected = df.toPandas().groupby('id').apply(foo).reset_index(drop=True)
         self.assertFramesEqual(expected, result)
 
+    def test_groupby_apply_dtypes(self):
+        from pyspark.sql.functions import pandas_udf
+        df = self.data
+
+        def foo(df):
+            ret = df
+            ret = ret.assign(v3=df.v * 5.0 + 1)
+            return ret
+
+        sample_df = df.filter(df.id == 1).toPandas()
+
+        foo_udf = pandas_udf(
+            foo,
+            foo(sample_df).dtypes
+        )
+
+        result = df.groupby('id').apply(foo_udf).sort('id').toPandas()
+        expected = df.toPandas().groupby('id').apply(foo).reset_index(drop=True)
+        self.assertFramesEqual(expected, result)
+
 
 if __name__ == "__main__":
     from pyspark.sql.tests import *

python/pyspark/sql/types.py

@@ -1597,7 +1597,7 @@ def convert(self, obj, gateway_client):
 register_input_converter(DateConverter())
 
 
-def toArrowType(dt):
+def to_arrow_type(dt):
     """ Convert Spark data type to pyarrow type
     """
     import pyarrow as pa
@@ -1623,6 +1623,31 @@ def toArrowType(dt):
         raise TypeError("Unsupported type in conversion to Arrow: " + str(dt))
     return arrow_type
 
+def from_pandas_type(dt):
+    """ Convert pandas data type to Spark data type
+    """
+    import pandas as pd
+    import numpy as np
+    if dt == np.int32:
+        return IntegerType()
+    elif dt == np.int64:
+        return LongType()
+    elif dt == np.float32:
+        return FloatType()
+    elif dt == np.float64:
+        return DoubleType()
+    elif dt == np.object:
+        return StringType()
+    elif dt == np.dtype('datetime64[ns]') or type(dt) == pd.api.types.DatetimeTZDtype:
+        return TimestampType()
+    else:
+        raise ValueError("Unsupported numpy type in conversion to Spark: {}".format(dt))
+
+def from_pandas_dtypes(dtypes):
+    """ Convert pandas DataFrame dtypes to Spark schema
+    """
+    return StructType([StructField(dtypes.axes[0][i], from_pandas_type(dtypes[i]))
+                       for i in range(len(dtypes))])
 
 def _test():
     import doctest

python/pyspark/worker.py

@@ -32,7 +32,7 @@
 from pyspark.serializers import write_with_length, write_int, read_long, \
     write_long, read_int, SpecialLengths, PythonEvalType, UTF8Deserializer, PickleSerializer, \
     BatchedSerializer, ArrowStreamPandasSerializer
-from pyspark.sql.types import toArrowType
+from pyspark.sql.types import to_arrow_type
 from pyspark import shuffle
 from pyspark.sql.types import StructType, IntegerType, LongType, FloatType, DoubleType
 
@@ -76,7 +76,7 @@ def wrap_udf(f, return_type):
 
 def wrap_pandas_udf(f, return_type):
     if isinstance(return_type, StructType):
-        arrow_return_types = list(toArrowType(field.dataType) for field in return_type)
+        arrow_return_types = list(to_arrow_type(field.dataType) for field in return_type)
 
         def fn(*a):
             import pandas as pd
@@ -89,7 +89,7 @@ def fn(*a):
         return fn
 
     else:
-        arrow_return_type = toArrowType(return_type)
+        arrow_return_type = to_arrow_type(return_type)
 
         def verify_result_length(*a):
             result = f(*a)