如何在PySpark DataFrame中将分隔列动态拆分至新列

需求说明 #

将PySpark DataFrame中以~为分隔符的单列值，根据给定的列名列表动态拆分为多个新列，要求使用Spark DataFrame原生函数实现。

输入示例 #

原始DataFrame：

指定列名列表：

col_names = ["id", "names", "sal", "country"]

解决方案代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import split

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

# 创建示例DataFrame
data = [("1~Ram~1000~US",), ("2~john~2000~UK",), ("3~Marry~7000~IND",)]
df = spark.createDataFrame(data, ["Raw_column_name"])

# 拆分列并动态指定新列名
split_col = split(df["Raw_column_name"], "~")
df_split = df.select(*[split_col[i].alias(col_names[i]) for i in range(len(col_names))])

# 可选：转换数据类型（比如将sal转为整数）
# df_split = df_split.withColumn("sal", df_split["sal"].cast("int"))

# 查看结果
df_split.show()

输出结果 #

+---+-----+----+-------+
| id|names| sal|country|
+---+-----+----+-------+
|  1|  Ram|1000|     US|
|  2| john|2000|     UK|
|  3|Marry|7000|    IND|
+---+-----+----+-------+

关键说明 #

1. split函数：split(col, sep)将指定列按分隔符拆分为数组类型列，这里用~作为分隔符。
2. 动态列生成：通过列表推导式遍历col_names，将拆分后的数组元素逐个取出并指定别名，实现动态匹配列名的需求。
3. 注意事项：需确保col_names的长度与拆分后数组的元素数量一致，避免索引越界。如果需要转换列的数据类型，可以在拆分后使用cast函数处理。

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