PySpark之SparkSQL 使用Pandas UDF

一、介绍

• Pandas UDF是用户定义的函数，由Spark执行，使用Arrow传输数据，Pandas执行数据，允许向量化操作。Pandas UDF是使用pandas_udf()作为装饰器或包装函数来定义的，不需要额外的配置。Pandas UDF通常表现为常规的PySpark函数API
• pyspark.sql.functions.pandas_udf(f=None, returnType=None, functionType=None)
  • f: 用户定义的函数；
  • returnType: 用户自定义函数的返回值类型，该值可以是pyspark.sql.types.DataType对象或DDL格式的类型字符串
  • functionType: pyspark.sql.functions.PandasUDFType中的枚举值。 默认值：SCALAR.存在此参数是为了兼容性。 鼓励使用Python类型。

二、类型1：seriesToSeries

• 类型提示可以表示为pandas.Series,…->pandas.Series

• 通过pandas_udf()与具有上述类型提示的函数一起使用，它会创建一个Pandas UDF，其中给定的函数采用一个或多个pandas.Series并输出一个pandas.Series。函数的输出应始终与输入的长度相同。在内部，PySpark将通过拆分为批次并作为数据的子集调用每个批次的函数，然后将结果连接在一起来执行Pandas UDF

import os
from pyspark.sql.types import LongType
from pyspark.sql import SparkSession
import pandas as pd
from pyspark.sql.functions import pandas_udf,col

os.environ['SPARK_HOME'] = '/export/server/spark'
PYSPARK_PYTHON = '/root/anaconda3/envs/pyspark_env/bin/python'
os.environ['PYSPARK_PYTHON'] = PYSPARK_PYTHON
os.environ['PYSPARK_DRIVER_PYTHON'] = PYSPARK_PYTHON


if __name__ == '__main__':
    spark = SparkSession\
        .builder\
        .appName("test")\
        .getOrCreate()
    sc = spark.sparkContext
    # 方式1：普通方式创建pandas_func
    def multiply_func(a:pd.Series,b:pd.Series)->pd.Series:
        return a*b
    multiply = pandas_udf(multiply_func,returnType=LongType())
    x = pd.Series([1,2,3])
    print(multiply_func(x,x))
    df = spark.createDataFrame(pd.DataFrame(x,columns=['x']))
    df.select(multiply(col("x"),col("x"))).show()
    print("="*100)
    # 方式2：装饰器方法
    @pandas_udf(LongType