如何获取Spark Dataset列值并动态用于SQL查询？

你的核心问题是误解了Spark中Column的本质：Column是延迟计算的逻辑表达式，不是实际数据值，而且Spark是分布式计算框架，Driver端代码无法直接获取Executor端每条记录的具体值，所以你原来的写法完全不符合Spark的运行机制。

为什么原来的代码行不通？ #

1. Column对象只描述数据操作逻辑，不能直接转成字符串拼接SQL，你看到的Level.minus(1)或companies只是表达式的toString结果，不是实际数值。
2. someMethod在Driver端执行，而Dataset的记录在Executor端处理，你没法在Driver端获取每条记录的LEVEL和COMPANIES值。
3. 直接调用collectAsList()会把整个Dataset的数据拉到Driver端，不仅性能极差，还可能触发内存溢出，而且无法对应到当前处理的那条记录。

正确实现方案 #

根据你的业务逻辑（当LEVEL>1时，关联DS1和DS2，获取符合条件的level数组更新COMPANIES），推荐两种实现方式：

方式一：用Spark内置API实现（推荐，性能最优）

利用Spark的DataFrame关联、分组聚合API实现，避免自定义UDF，Spark会自动优化执行计划：

import org.apache.spark.sql.functions._
import org.apache.spark.sql.Column

// 1. 处理需要更新的记录：炸开COMPANIES数组，方便关联
val explodedDS1 = DS1.filter(col("LEVEL") > 1)
  .withColumn("company_id", explode(col("COMPANIES")))

// 2. 按照你的SQL逻辑关联DS1（别名cs）和DS2（别名cp）
// 关联条件：cs.level = 当前记录LEVEL-1，且cs.company_private_id在当前COMPANIES数组中
val joined = explodedDS1.join(DS1.as("cs"), 
  col("cs.level") === col("LEVEL") - 1 && col("cs.company_private_id") === col("company_id"),
  "inner"
)

// 3. 按原记录的所有字段分组，收集符合条件的cs.level到数组
val updatedPart = joined.groupBy(DS1.columns.map(col): _*)
  .agg(collect_set(col("cs.level")).as("new_COMPANIES"))

// 4. 合并原DS1和更新后的部分，完成最终更新
val finalDS = DS1.join(updatedPart, DS1.columns.map(col): _*, "left_outer")
  .withColumn("COMPANIES", 
    when(col("new_COMPANIES").isNotNull, col("new_COMPANIES"))
    .otherwise(col("COMPANIES"))
  )
  .drop("new_COMPANIES", "company_id") // 清理临时列

方式二：自定义UDF（适合复杂业务逻辑）

如果业务逻辑过于复杂，无法用内置API实现，可以用UDF，但绝对不能在UDF中调用sparkSession.sql（UDF在Executor端运行，无法访问Driver端的SparkSession）。需要先把依赖数据广播到Executor：

import org.apache.spark.sql.functions._
import org.apache.spark.sql.SparkSession

// 1. 提前把DS1中需要的预处理数据加载到广播变量（Executor端可访问）
val ds1LevelCompanyMap = DS1.select("level", "company_private_id")
  .rdd.groupBy(_.getInt(0))
  .mapValues(_.map(_.getInt(1)).toSet)
  .collectAsMap()

val broadcastMap = sparkSession.sparkContext.broadcast(ds1LevelCompanyMap)

// 2. 定义UDF：输入COMPANIES数组和LEVEL，输出更新后的数组
val updateCompaniesUdf = udf((companies: Array[Int], level: Int) => {
  if (level <= 1) {
    companies // LEVEL<=1时保留原数组
  } else {
    val targetLevel = level - 1
    // 从广播变量中获取目标level对应的company集合
    val targetCompanies = broadcastMap.value.getOrElse(targetLevel, Set.empty[Int])
    // 筛选出COMPANIES数组中符合条件的元素，对应返回targetLevel的数组
    companies.filter(targetCompanies.contains).map(_ => targetLevel)
  }
})

// 3. 应用UDF更新列
val finalDS = DS1.withColumn("COMPANIES", updateCompaniesUdf(col("COMPANIES"), col("LEVEL")))

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