pyspark中自定义函数的用法

pyspark中自定义函数比python中多了一部注册，整体流程是“定义-注册-调用”，其中注册和调用两步在sparksql和DSL中又有所区别，具体如下：

from  pyspark.sql import SparkSession

第一步定义一个函数：

    def squared_func(number):
        return number *number

这一步和在python中是一样的。

    #创建一个临时表，用于后边测试函数
    spark.range(1,20).registerTempTable("tmp_test")

第二步和第三步：

（1）在spark sql中

    #  用spark.udf.register注册函数
    from pyspark.sql.types import LongType  #导入LongType数据类型包
    spark.udf.register('func_squared_type',squared_func,returnType=LongType())

    spark.udf.register有三个参数：

   1、注册后的函数名称，字符串

   2、已经定义好的函数，这里用第一步里边的squared_func

   3、返回的数据类型，默认是字符串类型，这里设置为LongType()，注意不要忘记括号。

    #或者直接用lambda定义并注册函数
    spark.udf.register('func_squared_lambda',lambda number:number * number)

   #调用函数

    squared_result1=spark.sql("""
         SELECT id,func_squared_type(id) as id_squared from  tmp_test
     """).collect()
    for result in squared_result1:
         print str(result["id"]) +'\t'+ str(result['id_squared'])

（2）在DSL中

    #用udf注册函数
    from  pyspark.sql.functions import udf
    squared_udf = udf(squared_func,LongType())

    udf有二个参数：

   1、已经定义好的函数，这里用第一步里边的squared_func

   2、返回的数据类型，默认是字符串类型，这里设置为LongType()，注意不要忘记括号。

   #调用函数
    test_df = spark.table("tmp_test")    #读取临时表的数据

    test_df.select("id",squared_udf('id').alias('id_squared')).show()          # .alias（）是取别名