SparkSQL 之 UDF、UDAF 函数代码实现

摘要本文从 UDF/UDAF/UDTF 三大函数类型、两种注册方式、弱类型 vs 强类型 UDAF、Aggregator 生命周期、性能陷阱五个维度配合 2 张架构图 完整代码彻底掌握 SparkSQL 自定义函数实现。关键词UDF, UDAF, UDTF, Aggregator, functions.udf, spark.udf.register一、三大函数类型UDF 一对一: 1 行 → 1 行 (name → UPPER(name)) UDAF 多对一: N 行 → 1 行 (多行 → SUM/AVG) UDTF 一对多: 1 行 → N 行 (一行 → explode 多行)二、函数分类 注册UDF 两种注册方式importorg.apache.spark.sql.functions._// SQL 注册spark.udf.register(myUpper,(s:String)s.toUpperCase)spark.sql(SELECT myUpper(name) FROM users)// DSL 注册valmyUpperUdfudf((s:String)s.toUpperCase)df.withColumn(upper_name,myUpperUdf(col(name)))三、UDAF 深度对比 Aggregator 生命周期强类型 Aggregator推荐importorg.apache.spark.sql.expressions.AggregatorcaseclassAverage(varsum:Double,varcount:Long)objectAverageAggregatorextendsAggregator[Double,Average,Double]{defzero:AverageAverage(0.0,0L)// 初始缓冲defreduce(b:Average,a:Double):Average{b.suma;b.count1;b}// 分区内defmerge(b1:Average,b2:Average):Average{b1.sumb2.sum;b1.countb2.count;b1}// 跨分区deffinish(reduction:Average):Doublereduction.sum/reduction.count// 输出defbufferEncoder:Encoder[Average]Encoders.productdefoutputEncoder:Encoder[Double]Encoders.scalaDouble}valavgUDAFAverageAggregator.toColumn.name(avg_score)ds.select(avgUDAF).show()弱类型 UserDefinedAggregateFunctionclassMyAvgUDAFextendsUserDefinedAggregateFunction{definputSchemaStructType(StructField(input,DoubleType)::Nil)defbufferSchemaStructType(StructField(sum,DoubleType)::StructField(count,LongType)::Nil)defdataTypeDoubleTypedefdeterministictruedefinitialize(buffer:MutableAggregationBuffer){buffer(0)0.0;buffer(1)0L}defupdate(buffer:MutableAggregationBuffer,input:Row){/* 累加 */}defmerge(b1:MutableAggregationBuffer,b2:Row){/* 合并 */}defevaluate(buffer:Row)buffer.getDouble(0)/buffer.getLong(1)}spark.udf.register(myAvg,newMyAvgUDAF)四、性能陷阱与最佳实践⚠️ UDF 是黑盒 → Catalyst 无法优化 · 无法谓词下推 · 无法 WholeStageCodegen · 逐行序列化调用 ✅ 优化建议: 1. 优先用 Spark SQL 内置函数 2. 复杂逻辑用 Scala 表达式组合 3. 必须用 UDF → Pandas UDF (Arrow 向量化, 快 100x) 4. UDAF 优先用强类型 Aggregator五、总结分类UDF 一对一 / UDAF 多对一 / UDTF 一对多注册SQL 用 registerDSL 用 functions.udfUDAF强类型 Aggregator 优于弱类型 UDAF作者大数据技术实践者博客blog.starzy.cnGitHubstarzy1990.github.io专注 AI Agent · LangGraph · RAG · 大数据架构 · 数据工程实践