【博学谷学习记录】超强总结，用心分享 | Spark SQL函数定义

#博学谷IT技术支持#

1. 窗口函数

DSL写法

# 1- 创建SparkSession对象
spark = SparkSession.builder.appName('df_write').master('local[*]').getOrCreate()

# 2-读取外部文件的数据
df = spark.read.csv(
    path='file:///export/data/workspace/ky06_pyspark/_03_SparkSql/data/pv.csv',
    header=True,
    inferSchema=True
)
df.createTempView('t1')

# 3- 执行相关的操作
# 统计每个cookie中, pv数量排名前二内容是哪些
df.select(
    '*',
    F.row_number().over(win.partitionBy('cookieid').orderBy(F.desc('pv'))).alias('rank1')
).where('rank1 <= 2').show()

2. 自定义函数

2.1 函数分类

• UDF函数
  • 一进一出
  • 例如split()、substr()
• UDAF函数
  • 多进一出
  • 例如sum()、avg()、count()
• UDTF函数
  • 一进多出，进去一行数据，产生多行或多列数据
  • 例如explode()

2.2 自定义原生函数

1. 创建一个python函数
2. 将python函数注册到Spark SQL中
   1. 方式一：用于SQL/DSL中，spark.udf.register(UDF函数名称,python函数名称,返回值类型)，UDF函数名称用于在SQL语法中使用。
   2. 方式二：用于DSL方式，F.udf(python函数名称,返回值类型)，或者在python函数上添加@F.udf(returnType=返回值类型)。
3. 用SQL/DSL方式使用
4. 示例

# 1- 创建SparkSession对象
spark = SparkSession.builder.appName('df_write').master('local[*]').getOrCreate()

# 2- 初始化一些数据
df = spark.createDataFrame(
    data=[(1,'张三 北京'),(2,'李四 上海'),(3,'王五 广州'),(4,'赵六 深圳'),(5,'田七 杭州')],
    schema='id int,name_address string'
)
df.createTempView('t1')
# 3- 执行相关的操作:
# 需求: 自定义一个函数, 将姓名和地址拆分开
schema = StructType().add('name',StringType()).add('address',StringType())

@F.udf(returnType=schema)
def split_data(data):
    res = data.split(' ')
    return {'name':res[0],'address':res[1]}
# 使用字典返回, key值必须和schema中定义字段名称保持一致

# 注册函数
spark.udf.register('split_data',split_data)
df.select(
    '*',
    split_data('name_address')['name'].alias('name'),
    split_data('name_address')['address'].alias('address')
).show()

3. 基于Pandas完成UDF函数

底层基于Arrow完成数据传输，主要是通过pandas_udf()完成对Pandas函数的包装，并注册成一个Spark SQL函数，支持语法糖模式。

自定义python函数要求：SeriesToSeries，即传入的参数类型必须是Series，输出返回的类型也必须是Series类型。

3.1 自定义UDF函数

举例：自定义两列数据，完成两列的求和计算

# 1- 创建SparkSession对象
spark = SparkSession.builder.appName('pandas_udf').master('local[*]').getOrCreate()

# 添加Arrow的相关的配置 : 在pandas的 UDF中, 底层默认会采用Arrow来进行数据传输的
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", True)

# 2- 初始化量两列数据集: a 和 b
df = spark.createDataFrame(
    data=[(1,2),(2,5),(3,2),(4,7),(8,3),(5,2)],
    schema='a int,b int'
)
df.createTempView('t1')

df.printSchema()
df.show()

# 3- 执行相关的操作:  自定义函数, 完成对 a 和 b的求和

# 3.1 自定义一个python函数, 完成对两列数据计算操作
@F.pandas_udf(returnType=IntegerType())
def sum_a_b(a:pd.Series,b:pd.Series) -> pd.Series:
    return a + b

# 3.2 将py函数进行注册:
# 语法糖模式: 仅支持 DSL
# 支持在SQL中使用
spark.udf.register('sum_a_b',sum_a_b)

# 3.3 测试使用
df.select('a','b',sum_a_b('a','b').alias('sum_ab')).show()

3.2 自定义UDAF函数

要求：SeriesTo标量，传入的类型必须是Series，返回类型必须是python基本数据类型。

举例：对某一列数据求平均值

# 1- 创建SparkSession对象
spark = SparkSession.builder.appName('pandas_udf').master('local[*]').getOrCreate()

# 2- 初始化相关的数据集
df = spark.createDataFrame(
    data=[(1,39.5),(2,63.5),(3,72.0),(4,85.0),(5,23.0),(6,25.0)],
    schema='id int,score float'
)
df.createTempView('t1')

# 3- 执行相关操作:
# 需求: 对某一列数据求平均值
# 3.1 自定义一个python的函数
@F.pandas_udf(returnType='float')
def score_avg(score:pd.Series) -> float:
    return score.mean()

# 3.2 注册函数
spark.udf.register('score_avg',score_avg)

# 3.3  使用函数
spark.sql("""
    select
        round(score_avg(score),2) as score_avg
    from t1
""").show()

df.select(score_avg('score').alias('score_avg')).show()

# 支持在窗口函数中使用
df.select(
    '*',
    score_avg('score').over(win.orderBy('score')).alias('score_avg')
).show()