基于指定列拆分DataFrame列：解决‘Column is not iterable’错误

问题描述 #

有一个包含两列的Spark DataFrame：id（整数类型）和values（结构体数组类型）。需要按指定列名列表作为前缀拆分name字段——当列表中的列名是name字段的前缀时，拆分出前缀（new_name）和剩余部分（new_value）。当前代码执行时抛出Column is not iterable错误，原代码如下：

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, explode
from pyspark.sql.types import StructType, StructField, StringType, IntegerType, ArrayType

# 初始化SparkSession
spark = SparkSession.builder.appName("SplitNameExample").getOrCreate()

value_schema = ArrayType(
    StructType([
        StructField("name", StringType(), True),
        StructField("location", StringType(), True)
    ])
)

data = [
    (1, [
        {"name": "col1_US_value_name", "location": "usa"},
        {"name": "col2_name_plex", "location": "usa"},
        {"name": "col4_false_val", "location": "usa"},
        {"name": "col3_name_is_fantasy", "location": "usa"}
    ])
]

df = spark.createDataFrame(data, ["id", "values"])

df = df.withColumn("values", explode(col("values")).alias("values"))
df = df.select(col("id"), col("values.name").alias("name"))
df.display()

col_names = ["col1","col2_name","col3_name_is","col4"]

pattern = "|".join(col_names)
print(pattern)
df = df.withColumn("new_name", regexp_extract("name", pattern, 0))
df.display()

df = df.withColumn(
    "new_value", 
    split(df['name'], concat(df['new_name'], lit('_'))).getItem(1) # 尝试split()[1]时也报错'Column is not iterable'
)
df.display()

期望输出结果：

id    name                new_name          new_value
1   col1_US_value_name    col1              US_value_name
1   col2_name_plex        col2_name         plex
1   col4_false_val        col4              false_val
1   col3_name_is_fantasy  col3_name_is      fantasy

错误原因 #

1. Column is not iterable错误：Spark的split函数返回Column对象，不能用Python列表索引语法[]，必须用getItem(n)获取拆分后的第n个元素。
2. 正则匹配逻辑缺陷：原正则未限定从字符串开头匹配，可能误匹配name字段中间的子串；同时未明确优先匹配长前缀的逻辑，可能导致短前缀抢先匹配，破坏拆分规则。

修正后的代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, explode, regexp_extract, split, concat, lit

# 初始化SparkSession
spark = SparkSession.builder.appName("SplitNameExample").getOrCreate()

value_schema = ArrayType(
    StructType([
        StructField("name", StringType(), True),
        StructField("location", StringType(), True)
    ])
)

data = [
    (1, [
        {"name": "col1_US_value_name", "location": "usa"},
        {"name": "col2_name_plex", "location": "usa"},
        {"name": "col4_false_val", "location": "usa"},
        {"name": "col3_name_is_fantasy", "location": "usa"}
    ])
]

df = spark.createDataFrame(data, ["id", "values"])

# 展开数组并提取name字段
df = df.withColumn("values", explode(col("values")))
df = df.select(col("id"), col("values.name").alias("name"))

# 列名列表：长前缀排在前面，避免短前缀优先匹配
col_names = ["col3_name_is", "col2_name", "col1", "col4"]
# 构建正则：限定从开头匹配，确保只匹配前缀
pattern = f"^({'|'.join(col_names)})"

# 提取前缀new_name
df = df.withColumn("new_name", regexp_extract(col("name"), pattern, 1))

# 拆分获取new_value：用concat拼接分隔符，getItem获取拆分后的第二部分
df = df.withColumn(
    "new_value",
    split(col("name"), concat(col("new_name"), lit("_"))).getItem(1)
)

# 展示结果
df.select("id", "name", "new_name", "new_value").show(truncate=False)

执行结果 #

+---+-----------------------+-----------+-------------+
|id |name                   |new_name   |new_value    |
+---+-----------------------+-----------+-------------+
|1  |col1_US_value_name     |col1       |US_value_name|
|1  |col2_name_plex         |col2_name  |plex         |
|1  |col4_false_val         |col4       |false_val    |
|1  |col3_name_is_fantasy   |col3_name_is|fantasy      |
+---+-----------------------+-----------+-------------+

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