如何基于排除码向Spark DataFrame条件性新增计算列

针对你的需求，这里提供两种可行的实现方案，核心解决跨DataFrame的条件判断问题：

方法一：基于广播变量的isin判断（适合排除码数据量较小的场景） #

先从排除码DataFrame中提取所有排除码，再在原始DataFrame中通过条件判断新增列。这种方式无需关联两个DataFrame，操作更轻量。

Scala 实现 #

// 提取排除码集合
val excludeCodes = exclude_df.select("exclusioncode").as[String].collect().toSet

// 导入所需函数
import org.apache.spark.sql.functions.{when, lit, col}

// 新增条件列
val result_df = raw_df
  .withColumn("offeramount1", 
    when(!col("code").isin(excludeCodes.toSeq: _*), col("amount") * 0.75).otherwise(lit(null))
  )
  .withColumn("offeramount2", 
    when(!col("code").isin(excludeCodes.toSeq: _*), col("amount") * 0.65).otherwise(lit(null))
  )

Python 实现 #

# 提取排除码列表
exclude_codes = [row[0] for row in exclude_df.select("exclusioncode").collect()]

# 导入所需函数
from pyspark.sql.functions import when, col

# 新增条件列
result_df = raw_df\
  .withColumn("offeramount1", 
    when(~col("code").isin(exclude_codes), col("amount") * 0.75).otherwise(None)
  )\
  .withColumn("offeramount2", 
    when(~col("code").isin(exclude_codes), col("amount") * 0.65).otherwise(None)
  )

方法二：基于左连接的条件判断（适合排除码数据量较大的场景） #

如果排除码数据量很大，直接collect()会占用Driver节点内存，此时可以通过左连接标记排除状态，再新增列。

Scala 实现 #

// 导入所需函数
import org.apache.spark.sql.functions.{when, lit, col}

// 给排除码DF添加标记列
val marked_exclude_df = exclude_df.withColumn("is_excluded", lit(true))

// 左连接原始DF和标记后的排除码DF
val joined_df = raw_df.join(
  marked_exclude_df, 
  raw_df("code") === marked_exclude_df("exclusioncode"), 
  "left_outer"
)

// 新增条件列并清理冗余列
val result_df = joined_df
  .withColumn("offeramount1", 
    when(col("is_excluded").isNull, col("amount") * 0.75).otherwise(lit(null))
  )
  .withColumn("offeramount2", 
    when(col("is_excluded").isNull, col("amount") * 0.65).otherwise(lit(null))
  )
  .drop("exclusioncode", "is_excluded")

Python 实现 #

# 导入所需函数
from pyspark.sql.functions import when, lit, col

# 给排除码DF添加标记列
marked_exclude_df = exclude_df.withColumn("is_excluded", lit(True))

# 左连接原始DF和标记后的排除码DF
joined_df = raw_df.join(
  marked_exclude_df, 
  raw_df["code"] == marked_exclude_df["exclusioncode"], 
  "left_outer"
)

# 新增条件列并清理冗余列
result_df = joined_df\
  .withColumn("offeramount1", 
    when(col("is_excluded").isNull, col("amount") * 0.75).otherwise(None)
  )\
  .withColumn("offeramount2", 
    when(col("is_excluded").isNull, col("amount") * 0.65).otherwise(None)
  )\
  .drop("exclusioncode", "is_excluded")

说明 #

你之前直接用withColumn无法跨DataFrame比较的原因是：Spark的列表达式只能引用当前DataFrame的列，无法直接访问另一个DataFrame的数据集。上述两种方案分别通过提取排除码到本地集合或关联DataFrame的方式，解决了跨数据集的条件判断问题。

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