Spark SQL 进阶：自定义 UDF/UDAF 函数开发与复杂查询优化技巧

Spark SQL 进阶：自定义 UDF/UDAF 开发与查询优化

一、自定义函数开发

1. UDF（用户定义函数）

用于单行数据转换，输入一行输出一个值：

from pyspark.sql import functions as F
from pyspark.sql.types import FloatType

# 注册UDF：温度转换（华氏度→摄氏度）
@F.udf(returnType=FloatType())
def fahrenheit_to_celsius(f):
    return (f - 32) * 5/9

# 使用示例
df = spark.createDataFrame([(68,), (95,)], ["temp_f"])
df.withColumn("temp_c", fahrenheit_to_celsius("temp_f")).show()

输出：

+------+------+
|temp_f|temp_c|
+------+------+
|    68|  20.0|
|    95|  35.0|
+------+------+

2. UDAF（用户定义聚合函数）

处理多行数据聚合（Spark 3.0+推荐方式）：

from pyspark.sql import Aggregator
from pyspark.sql.types import StructType, StructField, DoubleType

# 计算几何平均数
class GeometricMean(Aggregator):
    def zero(self):
        return (1.0, 0)  # (乘积, 计数)
    
    def reduce(self, buffer, value):
        return (buffer[0] * value, buffer[1] + 1)
    
    def merge(self, buf1, buf2):
        return (buf1[0] * buf2[0], buf1[1] + buf2[1])
    
    def finish(self, buffer):
        return buffer[0] ** (1/buffer[1]) if buffer[1] > 0 else 0.0
    
    def outputSchema(self):
        return DoubleType()

# 注册使用
geom_mean_udaf = GeometricMean()
df = spark.createDataFrame([(2,), (8,), (32,)], ["value"])
df.agg(geom_mean_udaf("value").alias("geom_mean")).show()

输出：

+---------+
|geom_mean|
+---------+
|      8.0|  # ∛(2×8×32) = 8
+---------+

二、复杂查询优化技巧

1. 避免UDF性能陷阱

优先使用内置函数：

# 不推荐（UDF黑盒）
@F.udf
def square(x):
    return x * x

# 推荐（内置函数优化）
df.select(F.col("value") ** 2)

2. 谓词下推优化

# 自动优化：先过滤再JOIN
df1.join(df2, "id").filter(df1["date"] > "2023-01-01")

3. 动态分区裁剪

启用配置：

spark.conf.set("spark.sql.optimizer.dynamicPartitionPruning", "true")

4. 广播连接优化

处理大小表JOIN：

spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "10485760")  # 10MB
df_large.join(F.broadcast(df_small), "key")

5. 自适应查询执行（AQE）

Spark 3.0+ 自动优化：

spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")

三、高级实践案例

场景：计算用户购买频率的百分位数

# 1. 注册UDAF（中位数计算）
class MedianUDAF(Aggregator):
    ...  # 实现略（需排序分桶）

# 2. 优化查询
(spark.table("orders")
 .groupBy("user_id")
 .agg(F.collect_list("amount").alias("amounts"))
 .withColumn("median", median_udaf("amounts"))  # 自定义UDAF
 .filter("median > 1000")  # 谓词下推
 .write.partitionBy("region")  # 动态分区裁剪
 .parquet("output/"))

四、性能对比基准

最佳实践：

1. 优先使用内置函数，避免UDF
2. 对>100MB表启用广播连接
3. 始终开启AQE（Spark 3.0+）
4. 定期ANALYZE TABLE更新统计信息