PySpark中如何将数组元素内部的空格替换为下划线

问题原因 #

直接对数组类型的objects列调用regexp_replace时，Spark会隐式将数组转换为逗号拼接的字符串，此时替换空格会把数组元素分隔位置的空格也一并替换，无法实现仅替换元素内部空格的需求。

实现方法 #

Spark 2.4及以上版本内置数组高阶函数transform，可以遍历数组内的单个元素执行自定义处理，不会破坏原有数组结构，有两种常用写法：

通用兼容写法（支持Spark 2.4+所有版本） #

通过expr调用SQL原生的transform函数，兼容性最好：

from pyspark.sql import functions as F

df = df.withColumn(
    "concat_obj",
    F.expr("transform(objects, x -> regexp_replace(x, ' ', '_'))")
)

原生API写法（仅支持Spark 3.0+版本） #

高版本PySpark可以直接调用内置的transformAPI，无需写SQL表达式：

from pyspark.sql import functions as F

df = df.withColumn(
    "concat_obj",
    F.transform(
        "objects",
        lambda x: F.regexp_replace(x, " ", "_")
    )
)

执行结果 #

两种写法的输出完全符合预期：

• id=1时，concat_obj字段值为["sun", "solar_system", "mars", "milky_way"]
• id=2时，concat_obj字段值为["moon", "cosmic_rays", "orion_nebula"]

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