PySpark 自定义函数 UDF

转自：https://www.jianshu.com/p/06c6f9e50974

1. 最简单的注册UDF ---- 直接将lambda表达式注册成UDF
   下面是一个简单的清洗函数

from pyspark.sql.types import StringType
spark.udf.register('sex_distinct', lambda x: 'M' if x == u'男' else 'F', StringType())
spark.sql("""
select sex_distinct('男')
""").show()

结果

+---------------+
|sex_distinct(男)|
+---------------+
|              M|
+---------------+

2. 很多时候逻辑比较复杂，匿名函数不能完成工作，可以自己def一个函数，将def的函数名填入上面lambda函数所在位置就行

from pyspark.sql.types import StringType
def sex_distinct(sex_chinese):
    if sex_chinese == u'男':
        return u'M'
    else:
        return u'F'
spark.udf.register('sex_distinct_rename', sex_distinct, StringType())
spark.sql("""

select sex_distinct_rename('女')

""").show()

源码分析

    def register(self, name, f, returnType=None):
        """注册python的函数或自定义的函数为udf
:param name: sql语句中的函数名

:param f: 函数，可以python的，也可以是自定义的

:param returnType:

["DataType", "NullType", "StringType", "BinaryType", "BooleanType", "DateType",

"TimestampType", "DecimalType", "DoubleType", "FloatType", "ByteType", "IntegerType",

"LongType", "ShortType", "ArrayType", "MapType", "StructField", "StructType"]

可以看出规律了吧，和sql中的一一对应

:return: a user-defined function.
To register a nondeterministic Python function, users need to first build

a nondeterministic user-defined function for the Python function and then register it

as a SQL function.
returnType can be optionally specified when f is a Python function but not

when f is a user-defined function. Please see below.
1. 当f是python内部的函数（所谓python内部的函数就是python自带的函数）
returnType 默认是 string type 并且可以按需指定. 返回类型必须匹配指定类型.

这种情况约等于

register(name, f, returnType=StringType()).
>>> strlen = spark.udf.register("stringLengthString", lambda x: len(x))

>>> spark.sql("SELECT stringLengthString('test')").collect()

[Row(stringLengthString(test)=u'4')]
>>> spark.sql("SELECT 'foo' AS text").select(strlen("text")).collect()

[Row(stringLengthString(text)=u'3')]
>>> from pyspark.sql.types import IntegerType

>>> _ = spark.udf.register("stringLengthInt", lambda x: len(x), IntegerType())

>>> spark.sql("SELECT stringLengthInt('test')").collect()

[Row(stringLengthInt(test)=4)]
2. 当f是用户自定义的函数
Spark uses the return type of the given user-defined function as the return type of

the registered user-defined function. returnType should not be specified.

In this case, this API works as if register(name, f).
>>> from pyspark.sql.types import IntegerType

>>> from pyspark.sql.functions import udf

>>> slen = udf(lambda s: len(s), IntegerType())

>>> _ = spark.udf.register("slen", slen)

>>> spark.sql("SELECT slen('test')").collect()

[Row(slen(test)=4)]
>>> import random

>>> from pyspark.sql.functions import udf

>>> from pyspark.sql.types import IntegerType

>>> random_udf = udf(lambda: random.randint(0, 100), IntegerType()).asNondeterministic()

>>> new_random_udf = spark.udf.register("random_udf", random_udf)

>>> spark.sql("SELECT random_udf()").collect()  # doctest: +SKIP

[Row(random_udf()=82)]
>>> from pyspark.sql.functions import pandas_udf, PandasUDFType

>>> @pandas_udf("integer", PandasUDFType.SCALAR)  # doctest: +SKIP

... def add_one(x):

...     return x + 1

...

>>> _ = spark.udf.register("add_one", add_one)  # doctest: +SKIP

>>> spark.sql("SELECT add_one(id) FROM range(3)").collect()  # doctest: +SKIP

[Row(add_one(id)=1), Row(add_one(id)=2), Row(add_one(id)=3)]
>>> @pandas_udf("integer", PandasUDFType.GROUPED_AGG)  # doctest: +SKIP

... def sum_udf(v):

...     return v.sum()

...

>>> _ = spark.udf.register("sum_udf", sum_udf)  # doctest: +SKIP

>>> q = "SELECT sum_udf(v1) FROM VALUES (3, 0), (2, 0), (1, 1) tbl(v1, v2) GROUP BY v2"

>>> spark.sql(q).collect()  # doctest: +SKIP

[Row(sum_udf(v1)=1), Row(sum_udf(v1)=5)]
.. note:: Registration for a user-defined function (case 2.) was added from

Spark 2.3.0.

"""

# This is to check whether the input function is from a user-defined function or

# Python function.

if hasattr(f, 'asNondeterministic'):

if returnType is not None:

raise TypeError(

"Invalid returnType: data type can not be specified when f is"

"a user-defined function, but got %s." % returnType)

if f.evalType not in [PythonEvalType.SQL_BATCHED_UDF,

PythonEvalType.SQL_SCALAR_PANDAS_UDF,

PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF]:

raise ValueError(

"Invalid f: f must be SQL_BATCHED_UDF, SQL_SCALAR_PANDAS_UDF or "

"SQL_GROUPED_AGG_PANDAS_UDF")

register_udf = UserDefinedFunction(f.func, returnType=f.returnType, name=name,

evalType=f.evalType,

deterministic=f.deterministic)

return_udf = f

else:

if returnType is None: #这里指定了返回类型默认为StringType()

returnType = StringType()

register_udf = UserDefinedFunction(f, returnType=returnType, name=name,

evalType=PythonEvalType.SQL_BATCHED_UDF)

return_udf = register_udf._wrapped()

self.sparkSession._jsparkSession.udf().registerPython(name, register_udf._judf)

return return_udf

3. 复杂数据类型，ArrayType、MapType、StructType

   1. ArrayType Demo

from pyspark.sql.types import *
def split_to_array(input_string):

word_list = input_string.split('|')

return word_list
spark.udf.register('split_to_array', split_to_array, ArrayType(StringType()))
spark.sql("""

select split_to_array('我| shi|真的')

""").show()

结果

+-------------------------+
|split_to_array(我| shi|真的)|
+-------------------------+
|            [我,  shi, 真的]|
+-------------------------+

2. MapType Demo

from pyspark.sql.types import *
def word_count(input_string):

word_dict = {}

word_list = input_string.split(' ')

for word in word_list:

word_dict[word] = 0
for word in word_list:

word_dict[word] += 1
return word_dict
spark.udf.register('word_count', word_count, MapType(StringType(), IntegerType()))
spark.sql("""

select word_count('this apple belong to big apple')

""").show(truncate=False)

结果

+----------------------------------------------------------+
|word_count(this apple belong to big apple)                |
+----------------------------------------------------------+
|Map(this -> 1, big -> 1, belong -> 1, to -> 1, apple -> 2)|
+----------------------------------------------------------+

3. StructType Demo

from pyspark.sql.types import *
import hashlib
def string_to_struct(input_string):

my_dict={}

m = hashlib.md5()

m.update(input_string.encode('utf-8'))

my_dict['id'] = m.hexdigest()

my_dict['content'] = input_string

return my_dict
schema = StructType([

StructField("id", StringType(), True),

StructField("content", StringType(), True)

])
spark.udf.register('string_to_struct', string_to_struct, schema)
df = spark.sql("""

select string_to_struct('my name is hello world')

""")
df.show(truncate=False)
df.printSchema()

结果

+---------------------------------------------------------+
|string_to_struct(my name is hello world)                 |
+---------------------------------------------------------+
|[1e030e259e2c7759fb24572ac4d62d3f,my name is hello world]|
+---------------------------------------------------------+
root

|-- string_to_struct(my name is hello world): struct (nullable = true)

|    |-- id: string (nullable = true)

|    |-- content: string (nullable = true)

可以看出规律了吧，python中的类型要和自己定义的复杂类型对应起来。
此外，复杂数据类型支持嵌套，array中可以嵌套struct、map、array，其他同理。