User define function to return series

Viewed 65

I am trying to convert a python function to UDF to use on a spark dataframe. The function concatenates the column name with the columns values and then rowwise concatenation of the resulting strings. e.g.

import pyspark.sql.functions as F

dd = pd.DataFrame({'AV':['dfgt', 'dsfafgs'], 'AC':['dhghd', 'erytsh']})

def get_vector(dataframe):
    d2 = dataframe.astype(str).radd(dataframe.columns + ':')
    res = pd.Series(d2.fillna('').values.tolist()).str.join('/')
    return res

get_vector(dd)

#0        AC:dhghd/AV:dfgt
#1    AC:erytsh/AV:dsfafgs

To convert to UDF I tried,

dd_spark = spark.createDataFrame(dd)

@F.pandas_udf(dd_spark.schema, F.PandasUDFType.SCALAR)
def get_vector(dataframe):
    d2 = dataframe.astype(str).radd(dataframe.columns + ':')
    res = pd.Series(d2.fillna('').values.tolist()).str.join('/')
    return res

Which gives the following error:

Traceback (most recent call last):
  File "/usr/lib64/python2.7/threading.py", line 812, in __bootstrap_inner
    self.run()
  File "/usr/lib64/python2.7/threading.py", line 765, in run
    self.__target(*self.__args, **self.__kwargs)
  File "<stdin>", line 37, in get_data
  File "/opt/cloudera/parcels/SPARK2-2.3.0.cloudera2-1.cdh5.13.3.p0.316101/lib/spark2/python/pyspark/sql/udf.py", line 64, in _create_udf
    return udf_obj._wrapped()
  File "/opt/cloudera/parcels/SPARK2-2.3.0.cloudera2-1.cdh5.13.3.p0.316101/lib/spark2/python/pyspark/sql/udf.py", line 186, in _wrapped
    wrapper.returnType = self.returnType
  File "/opt/cloudera/parcels/SPARK2-2.3.0.cloudera2-1.cdh5.13.3.p0.316101/lib/spark2/python/pyspark/sql/udf.py", line 119, in returnType
    "not supported" % str(self._returnType_placeholder))
NotImplementedError: Invalid returnType with scalar Pandas UDFs: StructType(List(StructField(AC,StringType,true),StructField(AV,StringType,true),StructField(UI,StringType,true),StructField(C,StringType,true),StructField(I,StringType,true),StructField(A,StringType,true))) is not supported
1 Answers

The first parameter of pandas_udf should be the return type, not the type of the input data.

Replacing dd_spark.schema with T.StringType() solves the problem:

from pyspark.sql import functions as F
from pyspark.sql import types as T

@F.pandas_udf(T.StringType(), F.PandasUDFType.SCALAR)
def get_vector(dataframe):
    d2 = dataframe.astype(str).radd(dataframe.columns + ':')
    res = pd.Series(d2.fillna('').values.tolist()).str.join('/')
    return res

dd_spark.withColumn("vector", get_vector(
    F.struct([dd_spark[x] for x in dd_spark.columns]))).show(truncate=False)

prints

+-------+------+--------------------+                                           
|AV     |AC    |vector              |
+-------+------+--------------------+
|dfgt   |dhghd |AV:dfgt/AC:dhghd    |
|dsfafgs|erytsh|AV:dsfafgs/AC:erytsh|
+-------+------+--------------------+
Related