如何在SparkSQL中结合自定义isEligible函数实现多条件统计？

你有如下SparkSQL DataFrame：

name gender age isActive points 
------------------------------- 
Bob  M      12  true     100 
Hal  M      16  false    80 
Pat  F      21  true     70 
Lin  F      17  false    40 
Zac  M      18  true     20 
Mei  F      19  true     10 
Sal  M      13  false    10 

同时有一个自定义高阶函数：

def isEligible(prog: String) (name: String, gender: String, age: Int, isActive: Boolean, points: Int): Boolean

该函数用于判断用户是否符合某项目资格，现在需要处理上述DataFrame，输出汇总各项目符合人数的DataFrame：

prog1 prog2 prog3 prog4 
----------------------- 
7     3     2     5 

1. 把自定义函数转换成Spark UDF #

Spark无法直接识别Scala的高阶函数，所以我们需要把它包装成Spark支持的UDF。这里我们注册一个接受项目名和用户所有字段的UDF，内部调用你的isEligible逻辑：

import org.apache.spark.sql.functions.{udf, lit, when, sum, col}

// 注册UDF，适配Spark的参数传递方式
val isEligibleUDF = udf((prog: String, name: String, gender: String, age: Int, isActive: Boolean, points: Int) => {
  // 调用你的高阶函数：先传入项目名得到判断函数，再传入用户参数
  isEligible(prog)(name, gender, age, isActive, points)
})

2. 为每个项目生成资格标记列 #

接下来，我们遍历所有需要统计的项目，给原DataFrame添加对应列，用1表示符合资格，0表示不符合：

// 定义需要统计的项目列表
val targetProjects = List("prog1", "prog2", "prog3", "prog4")

// 遍历项目，逐个添加资格标记列
val dfWithEligibility = targetProjects.foldLeft(originalDF) { (currentDF, prog) =>
  currentDF.withColumn(
    prog,
    when(isEligibleUDF(lit(prog), col("name"), col("gender"), col("age"), col("isActive"), col("points")), 1).otherwise(0)
  )
}

3. 聚合统计各项目的符合人数 #

现在只需要对每个项目的标记列求和，就能得到每个项目的总符合人数：

// 生成各项目列的求和表达式
val aggregationExprs = targetProjects.map(prog => sum(col(prog)).alias(prog))

// 执行聚合，得到最终结果
val resultDF = dfWithEligibility.agg(aggregationExprs.head, aggregationExprs.tail: _*)

验证结果 #

执行完上面的代码后，resultDF就是你想要的汇总表，它会输出一行数据，包含每个项目的符合人数，和你给出的示例格式完全一致。

额外提示 #

• 如果你的项目数量很多，这种方式依然高效，因为Spark会把这些操作优化成分布式任务并行执行。
• 要是你用的是Python，逻辑完全一致，只是UDF的定义方式换成@udf装饰器或者spark.udf.register即可。
• 如果你想简化字段传递，也可以把用户的字段打包成struct，但上面的方式可读性更强，维护起来更方便。

内容的提问来源于stack exchange，提问作者user3685285