Spark DataFrame新增列遇TypeError：如何依据指定列表判断赋值1或0？

你的代码报错TypeError: Column is not iterable，本质是把Spark分布式的Column对象当成了Python本地可迭代集合——Spark的Column是对分布式数据列的抽象，不能用for val in rdd_df["a"]这种本地遍历的方式操作，必须用Spark原生的API来实现列级别的条件逻辑。

下面是两种正确的实现方案：

方案一：使用Spark原生when/otherwise + isin（推荐） #

这是最符合Spark分布式计算逻辑的方式，性能最优，不需要自定义函数：

from pyspark import SparkContext
from pyspark.sql import SQLContext
from pyspark.sql.functions import when, col

sc = SparkContext.getOrCreate()
sqlContext = SQLContext(sc)

# 构造测试数据
rdd = sc.parallelize([{"a":'y'}, {"a":'y', "b":2}, {"a":'n', "c":3}])
rdd_df = sqlContext.read.json(rdd)
yes_list = ['y']

# 新增条件列
result_df = rdd_df.withColumn(
    "new_col",
    when(col("a").isin(yes_list), 1).otherwise(0)
)

result_df.show()

代码说明： #

• col("a").isin(yes_list)：Spark内置函数，检查列a的每个值是否在yes_list中，返回布尔类型的Column
• when(条件, 结果1).otherwise(结果2)：Spark的条件分支函数，相当于SQL里的CASE WHEN，对整个列进行批量处理

方案二：使用UDF（用户自定义函数，仅适用于复杂逻辑） #

如果你的条件逻辑比“是否在列表中”更复杂，可以用UDF，但性能不如原生API（因为UDF会把数据拉到本地节点处理，涉及序列化/反序列化开销）：

from pyspark import SparkContext
from pyspark.sql import SQLContext
from pyspark.sql.functions import udf
from pyspark.sql.types import IntegerType

sc = SparkContext.getOrCreate()
sqlContext = SQLContext(sc)

rdd = sc.parallelize([{"a":'y'}, {"a":'y', "b":2}, {"a":'n', "c":3}])
rdd_df = sqlContext.read.json(rdd)
yes_list = ['y']

# 定义自定义函数
def check_in_list(val):
    return 1 if val in yes_list else 0

# 注册UDF，指定返回类型为Integer
check_udf = udf(check_in_list, IntegerType())

# 新增列
result_df = rdd_df.withColumn("new_col", check_udf(col("a")))

result_df.show()

注意事项： #

• 尽量优先使用Spark原生函数，避免不必要的UDF
• 如果UDF中用到外部变量（比如yes_list），要确保变量能被序列化到各个节点

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