pyspark udf

简单的注册UDF

直接用lambda表达式注册成UDF

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

返回的结果为 M

自己定义函数UDF

如果遇到复杂的开发逻辑，简单的lambda函数不能够满足需求，需要进行自定义函数进行相关UDF的开发

from pyspark.sql.types import *
def get_name(name):
	if name == u'男':
		return u'M'
	else:
		return u'F'
spark.udf.register('get_name',get_name,StringType())
spark.sql("select get_name('男')).show()

	

pyspark udf 源码解析

    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

复杂数据类型：ArrayType、MapType、StructType

ArrayType

from pyspark.sql.types import *
def split_word(name):
	result = name.split('-')
	return name

spark.udf.register("split_name",split_name,ArrayType(StringType()))
spark.sql("select split_word('2021-10-10')").show()

MapType

from pyspark.sql.types import *
def word_count(data):
	data_dict = {}
	data_list = data.split(",")
	for word in data_list:
		data_dict[word]=0
	for word in data_list:
		data_dict[word]+=1
	return data_dict
spark.udf.register('word_count',word_count,MapType(StringType(),IntegerType()))
spark.sql("select word_count('hello,python,hello,world,hello')").show()
	

StructType

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()

此外，复杂数据类型支持嵌套，array中可以嵌套struct、map、array，其他同理。