• 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|
    +---------------+
    
    1. 很多时候逻辑比较复杂,匿名函数不能完成工作,可以自己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
    

    1. 复杂数据类型,ArrayTypeMapTypeStructType

      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, 真的]|
    +-------------------------+
    
    1. 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)|
    +----------------------------------------------------------+
    
    1. 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中可以嵌套structmaparray,其他同理。

  • 相关阅读:
    《UIP在NIOS上的移植》
    切勿使用:指向局部变量的指针作为函数的返回指针!
    Oeacle创建表空间
    Oracle SQL 语言分类
    线程整理
    输入输出
    异常处理
    哈希算法
    java链表
    课上重点整理
  • 原文地址:https://www.cnblogs.com/leebxo/p/13492631.html
Copyright © 2020-2023  润新知