如何用PySpark将两列数组以分号分隔拼接为字符串？

问题描述 #

现有如下结构的Spark DataFrame：

+--+-----------------------+---------------------+
|id|                   type|                 name|
+--+-----------------------+---------------------+
| 1|[stars]                |[sun, altair, sirius]|
| 2|[solar system, planets]|[mars, earth]        |
| 3|[natural satellites]   |[moon, io, titan]    |
+--+-----------------------+---------------------+

需要新增一列result，将type列的数组、name列的数组分别转为逗号分隔的字符串后，用; 作为分隔符拼接，预期输出如下：

+--+-----------------------+---------------------+-----------------------------------+
|id|                   type|                 name|                             result|
+--+-----------------------+---------------------+-----------------------------------+
| 1|[stars]                |[sun, altair, sirius]|stars; sun, altair, sirius         |
| 2|[solar system, planets]|[mars, venus]        |solar system, planets; mars, venus |
| 3|[natural satellites]   |[moon, io, titan]    |natural satellites; moon, io, titan|
+--+-----------------------+---------------------+-----------------------------------+

解决方案 #

直接使用concat_ws无法满足需求，因为它会把两个数组的所有元素直接拼接，而我们需要先将每个数组内部转为逗号分隔的字符串，再用分号拼接这两个字符串。具体实现步骤如下：

1. 导入PySpark所需函数：from pyspark.sql.functions import concat_ws, col
2. 对type和name列分别调用concat_ws(',', col)，将数组转为逗号分隔的字符串
3. 再次调用concat_ws('; ', ...)，把两个处理后的字符串拼接成目标result列

完整代码示例 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import concat_ws, col

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

# 构造测试数据
data = [
    (1, ["stars"], ["sun", "altair", "sirius"]),
    (2, ["solar system", "planets"], ["mars", "venus"]),
    (3, ["natural satellites"], ["moon", "io", "titan"])
]
df = spark.createDataFrame(data, ["id", "type", "name"])

# 生成result列
result_df = df.withColumn(
    "result",
    concat_ws(
        "; ",
        concat_ws(",", col("type")),
        concat_ws(",", col("name"))
    )
)

# 输出结果
result_df.show(truncate=False)

执行上述代码后，即可得到符合预期格式的DataFrame。

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