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.
    """</span>
    <span class="token comment"># This is to check whether the input function is from a user-defined function or</span>
    <span class="token comment"># Python function.</span>
    <span class="token keyword">if</span> <span class="token builtin">hasattr</span><span class="token punctuation">(</span>f<span class="token punctuation">,</span> <span class="token string">'asNondeterministic'</span><span class="token punctuation">)</span><span class="token punctuation">:</span>
        <span class="token keyword">if</span> returnType <span class="token keyword">is</span> <span class="token keyword">not</span> <span class="token boolean">None</span><span class="token punctuation">:</span>
            <span class="token keyword">raise</span> TypeError<span class="token punctuation">(</span>
                <span class="token string">"Invalid returnType: data type can not be specified when f is"</span>
                <span class="token string">"a user-defined function, but got %s."</span> <span class="token operator">%</span> returnType<span class="token punctuation">)</span>
        <span class="token keyword">if</span> f<span class="token punctuation">.</span>evalType <span class="token keyword">not</span> <span class="token keyword">in</span> <span class="token punctuation">[</span>PythonEvalType<span class="token punctuation">.</span>SQL_BATCHED_UDF<span class="token punctuation">,</span>
                              PythonEvalType<span class="token punctuation">.</span>SQL_SCALAR_PANDAS_UDF<span class="token punctuation">,</span>
                              PythonEvalType<span class="token punctuation">.</span>SQL_GROUPED_AGG_PANDAS_UDF<span class="token punctuation">]</span><span class="token punctuation">:</span>
            <span class="token keyword">raise</span> ValueError<span class="token punctuation">(</span>
                <span class="token string">"Invalid f: f must be SQL_BATCHED_UDF, SQL_SCALAR_PANDAS_UDF or "</span>
                <span class="token string">"SQL_GROUPED_AGG_PANDAS_UDF"</span><span class="token punctuation">)</span>
        register_udf <span class="token operator">=</span> UserDefinedFunction<span class="token punctuation">(</span>f<span class="token punctuation">.</span>func<span class="token punctuation">,</span> returnType<span class="token operator">=</span>f<span class="token punctuation">.</span>returnType<span class="token punctuation">,</span> name<span class="token operator">=</span>name<span class="token punctuation">,</span>
                                           evalType<span class="token operator">=</span>f<span class="token punctuation">.</span>evalType<span class="token punctuation">,</span>
                                           deterministic<span class="token operator">=</span>f<span class="token punctuation">.</span>deterministic<span class="token punctuation">)</span>
        return_udf <span class="token operator">=</span> f
    <span class="token keyword">else</span><span class="token punctuation">:</span>
        <span class="token keyword">if</span> returnType <span class="token keyword">is</span> <span class="token boolean">None</span><span class="token punctuation">:</span> <span class="token comment">#這里指定了返回類型默認為StringType()</span>
            returnType <span class="token operator">=</span> StringType<span class="token punctuation">(</span><span class="token punctuation">)</span>
        register_udf <span class="token operator">=</span> UserDefinedFunction<span class="token punctuation">(</span>f<span class="token punctuation">,</span> returnType<span class="token operator">=</span>returnType<span class="token punctuation">,</span> name<span class="token operator">=</span>name<span class="token punctuation">,</span>
                                           evalType<span class="token operator">=</span>PythonEvalType<span class="token punctuation">.</span>SQL_BATCHED_UDF<span class="token punctuation">)</span>
        return_udf <span class="token operator">=</span> register_udf<span class="token punctuation">.</span>_wrapped<span class="token punctuation">(</span><span class="token punctuation">)</span>
    self<span class="token punctuation">.</span>sparkSession<span class="token punctuation">.</span>_jsparkSession<span class="token punctuation">.</span>udf<span class="token punctuation">(</span><span class="token punctuation">)</span><span class="token punctuation">.</span>registerPython<span class="token punctuation">(</span>name<span class="token punctuation">,</span> register_udf<span class="token punctuation">.</span>_judf<span class="token punctuation">)</span>
    <span class="token keyword">return</span> 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
<span class="token keyword">for</span> word <span class="token keyword">in</span> word_list<span class="token punctuation">:</span>
    word_dict<span class="token punctuation">[</span>word<span class="token punctuation">]</span> <span class="token operator">+=</span> <span class="token number">1</span>

<span class="token keyword">return</span> 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，其他同理。